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 delete_environment(self, environment_name: str) -> None: """Deletes an existing Cloud Composer environment. Args: environment_name: Name 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('Deleting "%s" Composer environment from "%s" project.', fully_qualified_name, self.project_id) try: request = self.client.projects().locations().environments().delete( name=fully_qualified_name) 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_NOT_FOUND_CODE: logging.info('The Composer environment %s does not exists.', fully_qualified_name) return logging.exception( 'Error occurred while deleting Composer environment.') raise Error('Error occurred while deleting Composer environment.')
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/{p}/serviceAccounts/{s}@{p}.iam.gserviceaccount.com'.format( p=project_id, s=service_account_name) 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: Optional[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.')