def set_service_account_role(project_id, service_account_name,
role_name) -> None:
"""Adds role to a given service account.
The roles grant service accounts appropriate permissions to use specific
resource. The role_name can be either primitive, predefined or custom roles.
Please see https://cloud.google.com/iam/docs/understanding-roles for list of
allowed primitive and predefined roles.
Args:
project_id: GCP project id.
service_account_name: The service account name.
role_name: The role to be added to the service account. The role_name
doesn't need "roles/" prefix to be added. Allowed values -
https://cloud.google.com/iam/docs/understanding-roles. e.g - editor
"""
logging.info('Adding "%s" role to "%s" service account in "%s" project',
role_name, service_account_name, project_id)
# Read existing binding.
service_account_email = _get_service_account_email(project_id,
service_account_name)
member = 'serviceAccount:{}'.format(service_account_email)
binding = {'role': f'roles/{role_name}', 'members': [member]}
request = _get_resource_manager_client().projects().getIamPolicy(
resource=project_id)
policy = utils.execute_request(request)
# Modify binding.
policy['bindings'].append(binding)
# Write binding.
set_iam_policy_request_body = {'policy': policy}
request = _get_resource_manager_client().projects().setIamPolicy(
resource=project_id, body=set_iam_policy_request_body)
utils.execute_request(request)
def get_environment(self, environment_name: str) -> Dict[str, Any]:
"""Retrieves details of a Composer environment.
Args:
environment_name: Name of the existing Composer environment. The fully
qualified environment name will be constructed as follows -
'projects/{project_id}/locations/{location}/environments/
{environment_name}'.
Returns:
environment: Details of Composer environment.
Raises:
Error: If the request was not processed successfully.
"""
fully_qualified_name = self._get_fully_qualified_env_name(
environment_name)
logging.info('Retrieving Composer environment details for "%s"',
fully_qualified_name)
try:
request = self.client.projects().locations().environments().get(
name=fully_qualified_name)
composer_environment_details = utils.execute_request(request)
return composer_environment_details
except errors.HttpError:
logging.exception(
'Error while retrieving Composer environment details.')
raise Error('Error while retrieving Composer environment details.')
def override_airflow_configs(
self, environment_name: str,
airflow_config_overrides: Dict[str, str]) -> None:
"""Overrides Airflow configurations on the existing Composer environment.
Args:
environment_name: Name of the existing Composer environment. The fully
qualified environment name will be constructed as follows -
'projects/{project_id}/locations/{location}/environments/
{environment_name}'.
airflow_config_overrides: Airflow configurations to be overridden in the
Composer environment.
Raises:
Error: If the request was not processed successfully.
"""
fully_qualified_name = self._get_fully_qualified_env_name(
environment_name)
logging.info(
'Overriding "%s" Airflow configurations in "%s" Composer '
'environment.', airflow_config_overrides, fully_qualified_name)
try:
request_body = {
'name': fully_qualified_name,
'config': {
'softwareConfig': {
'airflowConfigOverrides': airflow_config_overrides
}
}
}
request = (self.client.projects().locations().environments().patch(
name=fully_qualified_name,
body=request_body,
updateMask='config.softwareConfig.airflowConfigOverrides'))
operation = utils.execute_request(request)
operation_client = self.client.projects().locations().operations()
utils.wait_for_operation(operation_client, operation)
logging.info(
'Airflow configurations "%s" has been overridden in "%s" Composer '
'environment.', airflow_config_overrides, fully_qualified_name)
except errors.HttpError:
logging.exception(
'Error occurred while overriding Airflow configurations.')
raise Error(
'Error occurred while overriding Airflow configurations.')
def install_python_packages(self, environment_name: str,
packages: Dict[str, str]) -> None:
"""Install Python packages on the existing Composer environment.
Args:
environment_name: Name of the existing Composer environment. The fully
qualified environment name will be constructed as follows -
'projects/{project_id}/locations/{location}/environments/
{environment_name}'.
packages: Dictionary of Python packages to be installed in the Composer
environment. Each entry in the dictionary has dependency name as the key
and version as the value. e.g -
{'tensorflow' : "<=1.0.1", 'apache-beam': '==2.12.0', 'flask': '>1.0.3'}
Raises:
Error: If the list of packages is empty.
"""
if not packages:
raise Error('Package list cannot be empty.')
fully_qualified_name = self._get_fully_qualified_env_name(
environment_name)
logging.info('Installing "%s" packages in "%s" Composer environment.',
packages, fully_qualified_name)
try:
request_body = {
'name': fully_qualified_name,
'config': {
'softwareConfig': {
'pypiPackages': packages
}
}
}
request = (self.client.projects().locations().environments().patch(
name=fully_qualified_name,
body=request_body,
updateMask='config.softwareConfig.pypiPackages'))
operation = utils.execute_request(request)
operation_client = self.client.projects().locations().operations()
utils.wait_for_operation(operation_client, operation)
logging.info(
'Installed "%s" packages in "%s" Composer environment.',
packages, fully_qualified_name)
except errors.HttpError:
logging.exception('Error occurred while installing packages.')
raise Error('Error occurred while installing python packages.')
def set_environment_variables(
self, environment_name: str,
environment_variables: Dict[str, str]) -> None:
"""Sets environment variables on the existing Composer environment.
Args:
environment_name: Name of the existing Composer environment. The fully
qualified environment name will be constructed as follows -
'projects/{project_id}/locations/{location}/environments/
{environment_name}'.
environment_variables: Environment variables to be added to the Composer
environment.
Raises:
Error: If the request was not processed successfully.
"""
fully_qualified_name = self._get_fully_qualified_env_name(
environment_name)
logging.info(
'Setting "%s" environment variables in "%s" Composer '
'environment.', environment_variables, fully_qualified_name)
try:
request_body = {
'name': fully_qualified_name,
'config': {
'softwareConfig': {
'envVariables': environment_variables
}
}
}
request = (self.client.projects().locations().environments().patch(
name=fully_qualified_name,
body=request_body,
updateMask='config.softwareConfig.envVariables'))
operation = utils.execute_request(request)
operation_client = self.client.projects().locations().operations()
utils.wait_for_operation(operation_client, operation)
logging.info(
'Updated "%s" environment variables in "%s" Composer '
'environment.', environment_variables, fully_qualified_name)
except errors.HttpError:
logging.exception(
'Error occurred while setting environment variables.')
raise Error('Error occurred while setting environment variables.')
def enable_apis(self, apis: List[str]) -> None:
"""Enables multiple Cloud APIs for a GCP project.
Args:
apis: The list of APIs to be enabled.
Raises:
Error: If the request was not processed successfully.
"""
parent = f'projects/{self.project_id}'
request_body = {'serviceIds': apis}
try:
request = self.client.services().batchEnable(parent=parent,
body=request_body)
operation = utils.execute_request(request)
utils.wait_for_operation(self.client.operations(), operation)
except errors.HttpError:
logging.exception('Error occurred while enabling Cloud APIs.')
raise Error('Error occurred while enabling Cloud APIs.')
def _create_service_account_key(project_id: str, service_account_name: str,
file_object: Any) -> None:
"""Creates key for service account and writes private key to the file object.
Args:
project_id: GCP project id.
service_account_name: The service account name.
file_object: The file object to which the private key will be written.
"""
name = 'projects/{}/serviceAccounts/{}@{}.iam.gserviceaccount.com'.format(
project_id, service_account_name, project_id)
logging.info(
'Creating service account key for "%s" service account in "%s" project',
service_account_name, project_id)
request = _get_service_account_client().keys().create(name=name, body={})
service_account_key = utils.execute_request(request)
private_key = base64.b64decode(
service_account_key['privateKeyData']).decode()
private_key_json = json.loads(private_key)
json.dump(private_key_json, file_object)
def create_service_account(project_id: str, service_account_name: str,
role_name: str, file_name: str) -> Dict[str, Any]:
"""Create a new service account.
Args:
project_id: GCP project id.
service_account_name: The service account name.
role_name: The role to be assigned to the service account.
file_name: The file where service account key will be stored.
Returns:
service_account: The newly created service account.
Raises:
ValueError: If the service_account_name is empty.
ValueError: If the file_name is empty.
"""
if not service_account_name:
raise ValueError('Service account name cannot be empty.')
if not file_name:
raise ValueError('The file name cannot be empty.')
service_account_details = get_service_account(project_id,
service_account_name)
if service_account_details:
return service_account_details
logging.info('Creating "%s" service account in "%s" project',
service_account_name, project_id)
request = _get_service_account_client().create(
name='projects/' + project_id,
body={
'accountId': service_account_name,
'serviceAccount': {
'displayName': service_account_name.upper()
},
})
service_account_details = utils.execute_request(request)
set_service_account_role(project_id, service_account_name, role_name)
create_service_account_key(project_id, service_account_name, file_name)
return service_account_details
def create_environment(self,
environment_name: str,
zone: str = 'b',
disk_size_gb: int = _DISC_SIZE,
machine_type: str = _MACHINE_TYPE,
image_version: str = None,
python_version: str = _PYTHON_VERSION) -> None:
"""Creates new Cloud Composer environment.
Args:
environment_name: Name of Composer environment.
zone: Optional. Zone where the Composer environment will be created. It
defaults to 'b' since zone 'b' is present in all the regions.
Allowed values - https://cloud.google.com/compute/docs/regions-zones/.
disk_size_gb: Optional. The disk size in GB used for node VMs. It defaults
to 20GB since it is the minimum size.
machine_type: Optional. The parameter will specify what type of VM to
create.It defaults to 'n1-standard-1'. Allowed values -
https://cloud.google.com/compute/docs/machine-types.
image_version: The version of Composer and Airflow running in the
environment. If this is not provided, a default version is used as per
https://cloud.google.com/composer/docs/concepts/versioning/composer-versions.
python_version: The version of Python used to run the Apache Airflow. It
defaults to '3'.
Raises:
Error: If the provided disk size is less than 20GB.
"""
if disk_size_gb < 20:
raise Error(
('The minimum disk size needs to be 20GB to create Composer '
'environment'))
fully_qualified_name = self._get_fully_qualified_env_name(
environment_name)
parent = f'projects/{self.project_id}/locations/{self.location}'
composer_zone = f'{self.location}-{zone}'
location = f'projects/{self.project_id}/zones/{composer_zone}'
machine_type = (f'projects/{self.project_id}/zones/{composer_zone}/'
f'machineTypes/{machine_type}')
software_config = {'pythonVersion': python_version}
if image_version:
software_config['imageVersion'] = image_version
request_body = {
'name': fully_qualified_name,
'config': {
'nodeConfig': {
'location': location,
'machineType': machine_type,
'diskSizeGb': disk_size_gb
},
'softwareConfig': software_config
}
}
logging.info('Creating "%s" Composer environment for "%s" project.',
fully_qualified_name, self.project_id)
try:
request = self.client.projects().locations().environments().create(
parent=parent, body=request_body)
operation = utils.execute_request(request)
operation_client = self.client.projects().locations().operations()
utils.wait_for_operation(operation_client, operation)
except errors.HttpError as error:
if error.__dict__['resp'].status == _HTTP_CONFLICT_CODE:
logging.info('The Composer environment %s already exists.',
fully_qualified_name)
return
logging.exception(
'Error occurred while creating Composer environment.')
raise Error('Error occurred while creating Composer environment.')