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.')
