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 test_wait_for_completion_of_operation_handles_errors(self): mock_get_operation = (self.mock_client.get.return_value.execute) mock_get_operation.side_effect = [ self.operation_not_completed, self.operation_not_completed, self.operation_failed ] with self.assertRaises(utils.Error): utils.wait_for_operation(self.mock_client, self.operation_not_completed) self.assertEqual(3, mock_get_operation.call_count)
def test_wait_for_completion_of_operation(self): mock_get_operation = (self.mock_client.get.return_value.execute) mock_get_operation.side_effect = [ self.operation_not_completed, self.operation_not_completed, self.operation_completed ] utils.wait_for_operation(self.mock_client, self.operation_not_completed) self.assertEqual(3, mock_get_operation.call_count) self.assertEqual(2, self.sleep_mock.call_count)
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_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.')