# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import getpass import io import json import logging import os import pathlib import platform import tarfile import time import warnings from functools import cached_property from typing import Dict, List, Optional, Tuple import docker from docker import errors as docker_errors from docker.types import Mount from google.api_core import exceptions as api_exception from google.auth import exceptions as auth_exception from google.cloud import artifactregistry_v1 from google.cloud.orchestration.airflow import service_v1 from composer_local_dev import console, constants, errors, files, utils LOG = logging.getLogger(__name__) DOCKER_FILES = pathlib.Path(__file__).parent / "docker_files" def timeout_occurred(start_time): """Returns whether time since start is greater than OPERATION_TIMEOUT.""" return time.time() - start_time >= constants.OPERATION_TIMEOUT_SECONDS def get_image_mounts( env_path: pathlib.Path, dags_path: str, gcloud_config_path: str, kube_config_path: Optional[str], requirements: pathlib.Path, database_mounts: Dict[pathlib.Path, str] ) -> List[docker.types.Mount]: """ Return list of docker volumes to be mounted inside container. Following paths are mounted: - requirements for python packages to be installed - dags, plugins and data for paths which contains dags, plugins and data - gcloud_config_path which contains user credentials to gcloud - kube_config_path which contains user cluster credentials for K8S [Optional] - environment airflow sqlite db file location - database_mounts which contains the path for database mounts """ mount_paths = { requirements: "composer_requirements.txt", dags_path: "gcs/dags/", env_path / "plugins": "gcs/plugins/", env_path / "data": "gcs/data/", gcloud_config_path: ".config/gcloud", **database_mounts, } # Add kube_config_path only if it's provided if kube_config_path: mount_paths[kube_config_path] = ".kube/" return [ docker.types.Mount( source=str(source), target=f"{constants.AIRFLOW_HOME}/{target}", type="bind", ) for source, target in mount_paths.items() ] def get_default_environment_variables( dag_dir_list_interval: int, project_id: str, default_db_variables: Dict[str, str] ) -> Dict: """Return environment variables that will be set inside container.""" return { "AIRFLOW__API__AUTH_BACKEND": "airflow.api.auth.backend.default", "AIRFLOW__WEBSERVER__EXPOSE_CONFIG": "true", "AIRFLOW__CORE__LOAD_EXAMPLES": "false", "AIRFLOW__SCHEDULER__DAG_DIR_LIST_INTERVAL": dag_dir_list_interval, "AIRFLOW__CORE__DAGS_FOLDER": "/home/airflow/gcs/dags", "AIRFLOW__CORE__PLUGINS_FOLDER": "/home/airflow/gcs/plugins", "AIRFLOW__CORE__DATA_FOLDER": "/home/airflow/gcs/data", "AIRFLOW__WEBSERVER__RELOAD_ON_PLUGIN_CHANGE": "True", "COMPOSER_PYTHON_VERSION": "3", # By default, the container runs as the user `airflow` with UID 999. Set # this env variable to "True" to make it run as the current host user. "COMPOSER_CONTAINER_RUN_AS_HOST_USER": "False", "COMPOSER_HOST_USER_NAME": f"{getpass.getuser()}", "COMPOSER_HOST_USER_ID": f"{os.getuid() if platform.system() != 'Windows' else ''}", "AIRFLOW_HOME": "/home/airflow/airflow", "AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT": f"google-cloud-platform://?" f"extra__google_cloud_platform__project={project_id}&" f"extra__google_cloud_platform__scope=" f"https://www.googleapis.com/auth/cloud-platform", **default_db_variables, } def parse_env_variable( line: str, env_file_path: pathlib.Path ) -> Tuple[str, str]: """Parse line in format of key=value and return (key, value) tuple.""" try: key, value = line.split("=", maxsplit=1) except ValueError: raise errors.FailedToParseVariablesError(env_file_path, line) return key.strip(), value.strip() def load_environment_variables(env_dir_path: pathlib.Path) -> Dict: """ Load environment variable to be sourced in the local Composer environment. Raises an error if the variables.env file does not exist in the ``env_dir_path``. Args: env_dir_path (pathlib.Path): Path to the local composer environment. Returns: Dict: Environment variables. """ env_file_path = env_dir_path / "variables.env" LOG.info("Loading environment variables from %s", env_file_path) if not env_file_path.is_file(): raise errors.ComposerCliError( f"Environment variables file '{env_file_path}' not found." ) env_vars = dict() with open(env_file_path) as fp: for line in fp: line = line.strip() if not line or line.startswith("#"): continue key, value = parse_env_variable(line, env_file_path) if key in constants.NOT_MODIFIABLE_ENVIRONMENT_VARIABLES: LOG.warning( "'%s' environment variable cannot be set " "and will be ignored.", key, ) elif key in constants.STRICT_ENVIRONMENT_VARIABLES: possible_values = constants.STRICT_ENVIRONMENT_VARIABLES[key] if value not in possible_values: LOG.warning( "'%s' environment variable can be set " "to the one of the following values: '%s'", key, ",".join(possible_values), ) else: env_vars[key] = value else: env_vars[key] = value return env_vars def filter_not_modifiable_env_vars(env_vars: Dict) -> Dict: """ Filter out environment variables that cannot be modified by user. """ filtered_env_vars = dict() for key, val in env_vars.items(): if key in constants.NOT_MODIFIABLE_ENVIRONMENT_VARIABLES: LOG.warning( "'%s' environment variable cannot be set and will be ignored.", key, ) elif key in constants.STRICT_ENVIRONMENT_VARIABLES: possible_values = constants.STRICT_ENVIRONMENT_VARIABLES[key] if val not in possible_values: LOG.warning( "'%s' environment variable can be set " "to the one of the following values: '%s'", key, ",".join(possible_values), ) else: env_vars[key] = val else: filtered_env_vars[key] = val return filtered_env_vars def get_software_config_from_environment( project: str, location: str, environment: str ): """Get software configuration from the Composer environment. Args: project (str): Composer GCP Project ID. location (str): Location of the Composer environment. environment (str): Composer environment name. Returns: SoftwareConfig: Software configuration of the Composer environment. """ LOG.info("Getting Cloud Composer environment configuration.") client = service_v1.EnvironmentsClient() name = f"projects/{project}/locations/{location}/environments/{environment}" request = service_v1.GetEnvironmentRequest(name=name) LOG.debug(f"GetEnvironmentRequest: %s", request) try: response = client.get_environment(request=request) except api_exception.GoogleAPIError as err: raise errors.ComposerCliError( constants.COMPOSER_SOFTWARE_CONFIG_API_ERROR.format(err=str(err)) ) LOG.debug(f"GetEnvironmentResponse: %s", response) return response.config.software_config def parse_airflow_override_to_env_var(airflow_override: str) -> str: """ Parse airflow override variable name in format of section-key to AIRFLOW__SECTION__KEY. """ section, key = airflow_override.split("-", maxsplit=1) return f"AIRFLOW__{section.upper()}__{key.upper()}" def get_airflow_overrides(software_config): """ Returns dictionary with environment variable names and their values mapped from Airflow overrides in Composer Software Config. """ return { parse_airflow_override_to_env_var(k): v for k, v in software_config.airflow_config_overrides.items() } def get_env_variables(software_config): """ Returns dictionary with environment variable names (with unset values) mapped from Airflow environment variables in Composer Software Config. """ return {k: "" for k, _ in software_config.env_variables.items()} def assert_image_exists(image_version: str): """Asserts that image version exists. Raises if the image does not exist. Warns if the API error occurs, or we cannot access to Artifact Registry. Args: image_version: Image version in format of 'composer-x.y.z-airflow-a.b.c' """ airflow_v, composer_v = utils.get_airflow_composer_versions(image_version) image_tag = utils.get_image_version_tag(airflow_v, composer_v) LOG.info("Asserting that %s composer image version exists", image_tag) image_url = constants.ARTIFACT_REGISTRY_IMAGE_URL.format( airflow_v=airflow_v, composer_v=composer_v ) client = artifactregistry_v1.ArtifactRegistryClient() request = artifactregistry_v1.GetTagRequest(name=image_url) LOG.debug(f"GetTagRequest for %s: %s", image_tag, str(request)) try: client.get_tag(request=request) except api_exception.NotFound: raise errors.ImageNotFoundError(image_version=image_tag) from None except api_exception.PermissionDenied: warnings.warn( constants.IMAGE_TAG_PERMISSION_DENIED_WARN.format( image_tag=image_tag ) ) except ( auth_exception.GoogleAuthError, api_exception.GoogleAPIError, ) as err: raise errors.InvalidAuthError(err) def get_docker_image_tag_from_image_version(image_version: str) -> str: """ Parse image version to Airflow and Composer versions and return image tag with those versions if it exists. Args: image_version: Image version in format of 'composer-x.y.z-airflow-a.b.c' Returns: Composer image tag in Artifact Registry """ airflow_v, composer_v = utils.get_airflow_composer_versions(image_version) return constants.DOCKER_REGISTRY_IMAGE_TAG.format( airflow_v=airflow_v, composer_v=composer_v ) def is_mount_permission_error(error: docker_errors.APIError) -> bool: """Checks if error is possibly a Docker mount permission error.""" return ( error.is_client_error() and error.response.status_code == constants.BAD_REQUEST_ERROR_CODE and "invalid mount config" in error.explanation ) def copy_to_container(container, src: pathlib.Path) -> None: """Copy entrypoint file to Docker container.""" logging.debug("Copying entrypoint file to Docker container.") stream = io.BytesIO() with tarfile.open(fileobj=stream, mode="w|") as tar, open(src, "rb") as f: info = tar.gettarinfo(fileobj=f) info.name = src.name tar.addfile(info, f) container.put_archive(constants.AIRFLOW_HOME, stream.getvalue()) class EnvironmentStatus: def __init__(self, name: str, version: str, status: str): self.name = name self.version = version self.status = status.capitalize() def get_image_version(env): """ Return environment image version. If the environment is running, get image version from the container tag. Otherwise, get image version from the configuration. """ try: container = env.get_container(env.container_name) except errors.EnvironmentNotRunningError: logging.debug( constants.IMAGE_VERSION_CONTAINER_MISSING.format(env_name=env.name) ) return env.image_version if not container.image.tags: LOG.warning( constants.IMAGE_VERSION_TAG_MISSING.format(env_name=env.name) ) return env.image_version tag = container.image.tags[0] image_tag = tag.split(":")[-1] airflow_v, composer_v = utils.get_airflow_composer_versions(image_tag) airflow_v = utils.format_airflow_version_dotted(airflow_v) return utils.get_image_version_tag(airflow_v, composer_v) def get_environments_status( envs: List[pathlib.Path], ) -> List[EnvironmentStatus]: """Get list of environment statuses.""" environments_status = [] for env_path in envs: try: env = Environment.load_from_config(env_path, None) env_status = env.status() image_version = get_image_version(env) except errors.InvalidConfigurationError: env_status = "Could not parse the config" image_version = "x" environment_status = EnvironmentStatus( env_path.name, image_version, env_status ) environments_status.append(environment_status) return environments_status class EnvironmentConfig: def __init__(self, env_dir_path: pathlib.Path, port: Optional[int]): self.env_dir_path = env_dir_path self.config = self.load_configuration_from_file() self.project_id = self.get_str_param("composer_project_id") self.image_version = self.get_str_param("composer_image_version") self.location = self.get_str_param("composer_location") self.dags_path = self.get_str_param("dags_path") self.dag_dir_list_interval = self.parse_int_param( "dag_dir_list_interval", allowed_range=(0,) ) self.port = ( port if port is not None else self.parse_int_param("port", allowed_range=(0, 65536)) ) self.database_engine = self.get_str_param("database_engine") def load_configuration_from_file(self) -> Dict: """ Load environment configuration from json file. Returns: Dict: Environment configuration dictionary. """ config_path = self.env_dir_path / "config.json" LOG.info("Loading configuration file from %s", config_path) if not config_path.is_file(): raise errors.ComposerCliError( f"Configuration file '{config_path}' not found." ) with open(config_path) as fp: try: config = json.load(fp) except json.JSONDecodeError as err: raise errors.FailedToParseConfigError(config_path, err) return config def get_str_param(self, name: str): """ Get parameter from the config. Raises an error if the parameter does not exist in the config. """ try: return self.config[name] except KeyError: raise errors.MissingRequiredParameterError(name) from None def parse_int_param( self, name: str, allowed_range: Optional[Tuple[int, int]] = None, ): """ Get parameter from the config and convert it to integer. Raises an error if the parameter value is not a valid integer. Optional ``allowed_range`` argument can be used to validate if parameter value is in the given range. Args: name: Name of the parameter in the config allowed_range: Tuple containing allowed range of values Returns: Parameter value converted to integer """ try: value = self.get_str_param(name) value = int(value) except ValueError as err: raise errors.FailedToParseConfigParamIntError(name, value) from None if allowed_range is None: return value if value < allowed_range[0] or ( len(allowed_range) > 1 and value > allowed_range[1] ): raise errors.FailedToParseConfigParamIntRangeError( name, value, allowed_range ) return value class Environment: def __init__( self, env_dir_path: pathlib.Path, project_id: str, image_version: str, location: str, dags_path: Optional[str], dag_dir_list_interval: int = 10, database_engine: str = constants.DatabaseEngine.sqlite3, port: Optional[int] = None, pypi_packages: Optional[Dict] = None, environment_vars: Optional[Dict] = None, ): self.name = env_dir_path.name self.container_name = f"{constants.CONTAINER_NAME}-{self.name}" self.db_container_name = f"{constants.DB_CONTAINER_NAME}-{self.name}" self.docker_network_name = f"{constants.DOCKER_NETWORK_NAME}-{self.name}" self.env_dir_path = env_dir_path self.airflow_db = self.env_dir_path / "airflow.db" self.entrypoint_file = DOCKER_FILES / "entrypoint.sh" self.run_file = DOCKER_FILES / "run_as_user.sh" self.requirements_file = self.env_dir_path / "requirements.txt" self.project_id = project_id self.image_version = image_version self.image_tag = get_docker_image_tag_from_image_version(image_version) self.db_image_tag = 'postgres:14-alpine' self.airflow_db_folder = self.env_dir_path / "postgresql_data" self.location = location self.dags_path = files.resolve_dags_path(dags_path, env_dir_path) self.dag_dir_list_interval = dag_dir_list_interval self.database_engine = database_engine self.is_database_sqlite3 = self.database_engine == constants.DatabaseEngine.sqlite3 self.port: int = port if port is not None else 8080 self.pypi_packages = ( pypi_packages if pypi_packages is not None else dict() ) self.environment_vars = ( environment_vars if environment_vars is not None else dict() ) self.docker_client = self.get_client() def get_client(self): try: return docker.from_env() except docker.errors.DockerException as err: logging.debug("Docker not found.", exc_info=True) raise errors.DockerNotAvailableError(err) from None def get_container( self, container_name: str, assert_running: bool = False, ignore_not_found: bool = False ): """ Returns created docker container and raises when it's not created. Args: container_name: name of the container assert_running: assert that container is running ignore_not_found: change the behaviour of raising error in case of not found the container """ try: container = self.docker_client.containers.get(container_name) if ( assert_running and container.status != constants.ContainerStatus.RUNNING ): raise errors.EnvironmentNotRunningError() from None return container except docker_errors.NotFound: logging.debug("Container not found.", exc_info=True) if not ignore_not_found: raise errors.EnvironmentNotFoundError() from None @classmethod def load_from_config(cls, env_dir_path: pathlib.Path, port: Optional[int]): """Create local environment using 'config.json' configuration file.""" config = EnvironmentConfig(env_dir_path, port) environment_vars = load_environment_variables(env_dir_path) return cls( env_dir_path=env_dir_path, project_id=config.project_id, image_version=config.image_version, location=config.location, dags_path=config.dags_path, dag_dir_list_interval=config.dag_dir_list_interval, port=config.port, database_engine=config.database_engine, environment_vars=environment_vars, ) @classmethod def from_source_environment( cls, source_environment: str, project: str, location: str, env_dir_path: pathlib.Path, web_server_port: Optional[int], dags_path: Optional[str], database_engine: str, ): """ Create Environment using configuration retrieved from Composer environment. """ software_config = get_software_config_from_environment( project, location, source_environment ) pypi_packages = {k: v for k, v in software_config.pypi_packages.items()} env_variables = get_env_variables(software_config) airflow_overrides = get_airflow_overrides(software_config) env_variables.update(airflow_overrides) env_variables = filter_not_modifiable_env_vars(env_variables) return cls( env_dir_path=env_dir_path, project_id=project, image_version=software_config.image_version, location=location, dags_path=dags_path, dag_dir_list_interval=10, port=web_server_port, pypi_packages=pypi_packages, environment_vars=env_variables, database_engine=database_engine, ) def pypi_packages_to_requirements(self): """Create requirements file using environment PyPi packagest list.""" reqs = sorted( f"{key}{value}" for key, value in self.pypi_packages.items() ) reqs_lines = "\n".join(reqs) with open(self.env_dir_path / "requirements.txt", "w") as fp: fp.write(reqs_lines) def environment_vars_to_env_file(self): """ Write fetched environment variables keys to `variables.env` file. """ env_vars = sorted( f"# {key}=" for key, _ in self.environment_vars.items() ) env_vars_lines = "\n".join(env_vars) with open(self.env_dir_path / "variables.env", "w") as fp: fp.write(env_vars_lines) def assert_requirements_exist(self): """Asserts that PyPi requirements file exist in environment directory.""" req_file = self.env_dir_path / "requirements.txt" if not req_file.is_file(): raise errors.ComposerCliError(f"Missing '{req_file}' file.") def write_environment_config_to_config_file(self): """Saves environment configuration to config.json file.""" config = { "composer_image_version": self.image_version, "composer_location": self.location, "composer_project_id": self.project_id, "dags_path": self.dags_path, "dag_dir_list_interval": int(self.dag_dir_list_interval), "port": int(self.port), "database_engine": self.database_engine, } with open(self.env_dir_path / "config.json", "w") as fp: json.dump(config, fp, indent=4) @cached_property def database_extras(self) -> Dict[str, Dict]: env_path = self.env_dir_path extras = { constants.DatabaseEngine.sqlite3: { "mounts": { "folders": {}, "files": { env_path / "airflow.db": "airflow/airflow.db", } }, "env_vars": {}, "ports": {}, }, constants.DatabaseEngine.postgresql: { "mounts": { "folders": { env_path / "postgresql_data": "/var/lib/postgresql/data", }, "files": { env_path / ".keep": "airflow/.keep", } }, "env_vars": { "PGDATA": "/var/lib/postgresql/data/pgdata", "POSTGRES_USER": "postgres", "POSTGRES_PASSWORD": "airflow", "POSTGRES_DB": "airflow", "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN": f"postgresql+psycopg2://postgres:airflow@{self.db_container_name}:5432/airflow", }, "ports": { f"5432/tcp": "25432", }, }, } if self.database_engine in extras: return extras[self.database_engine] return extras[constants.DatabaseEngine.sqlite3] def create_container(self, **kwargs): try: return self.docker_client.containers.create(**kwargs) except docker_errors.APIError as err: logging.debug( "Received docker API error when creating container.", exc_info=True, ) if err.status_code == constants.CONFLICT_ERROR_CODE: raise errors.EnvironmentAlreadyRunningError( self.name ) from None raise def create_docker_container(self): """Creates docker container. Raises when docker container with the same name already exists. """ LOG.debug("Creating container") db_extras = self.database_extras grouped_db_mounts = db_extras["mounts"] db_mounts = { **grouped_db_mounts['files'], **grouped_db_mounts['folders'], } mounts = get_image_mounts( self.env_dir_path, self.dags_path, utils.resolve_gcloud_config_path(), utils.resolve_kube_config_path(), self.requirements_file, db_mounts, ) db_vars = db_extras["env_vars"] default_vars = get_default_environment_variables( self.dag_dir_list_interval, self.project_id, db_vars ) env_vars = {**default_vars, **self.environment_vars} if ( platform.system() == "Windows" and env_vars["COMPOSER_CONTAINER_RUN_AS_HOST_USER"] == "True" ): raise Exception( "COMPOSER_CONTAINER_RUN_AS_HOST_USER must be set to `False` on Windows" ) ports = { f"8080/tcp": self.port, } entrypoint = f"sh {constants.ENTRYPOINT_PATH}" memory_limit = constants.DOCKER_CONTAINER_MEMORY_LIMIT try: container = self.create_container( image=self.image_tag, name=self.container_name, entrypoint=entrypoint, environment=env_vars, mounts=mounts, ports=ports, mem_limit=memory_limit, detach=True, ) except docker_errors.ImageNotFound: LOG.debug( "Failed to create container with ImageNotFound error. " "Pulling the image..." ) self.pull_image() container = self.create_container( image=self.image_tag, name=self.container_name, entrypoint=entrypoint, environment=env_vars, mounts=mounts, ports=ports, mem_limit=memory_limit, detach=True, ) except docker_errors.APIError as err: error = f"Failed to create container with an error: {err}" if is_mount_permission_error(err): error += constants.DOCKER_PERMISSION_ERROR_HINT.format( docs_faq_url=constants.COMPOSER_FAQ_MOUNTING_LINK ) raise errors.EnvironmentStartError(error) copy_to_container(container, self.entrypoint_file) copy_to_container(container, self.run_file) return container def create_db_docker_container(self): """Creates docker container for database. Raises when docker container with the same name already exists. """ db_extras = self.database_extras grouped_db_mounts = db_extras["mounts"] db_mounts = { **grouped_db_mounts['files'], **grouped_db_mounts['folders'], } mounts = get_image_mounts( self.env_dir_path, self.dags_path, utils.resolve_gcloud_config_path(), utils.resolve_kube_config_path(), self.requirements_file, db_mounts, ) db_vars = db_extras["env_vars"] db_ports = db_extras["ports"] memory_limit = constants.DOCKER_CONTAINER_MEMORY_LIMIT self.docker_client.images.pull(self.db_image_tag) LOG.info("DB_VARS") LOG.info(db_vars) try: container = self.create_container( image=self.db_image_tag, name=self.db_container_name, environment=db_vars, mounts=mounts, ports=db_ports, mem_limit=memory_limit, detach=True, ) return container except docker_errors.APIError as err: error = f"Failed to create container for database with an error: {err}" if is_mount_permission_error(err): error += constants.DOCKER_PERMISSION_ERROR_HINT.format( docs_faq_url=constants.COMPOSER_FAQ_MOUNTING_LINK ) raise errors.EnvironmentStartError(error) def get_docker_network(self): try: return self.docker_client.networks.get(self.docker_network_name) except docker.errors.NotFound as _: return self.docker_client.networks.create(self.docker_network_name) except docker.errors.APIError as err: error = f"Failed to create/get network an error: {err}" raise errors.EnvironmentStartError(error) def pull_image(self): """Pull Composer docker image.""" try: # TODO: (b/237054183): Print detailed status (progress bar of image pulling) with console.get_console().status(constants.PULL_IMAGE_MSG): self.docker_client.images.pull(self.image_tag) except (docker_errors.ImageNotFound, docker_errors.APIError): logging.debug("Failed to pull composer image.", exc_info=True) raise errors.ImageNotFoundError(self.image_version) from None def pull_db_image(self): try: # TODO: (b/237054183): Print detailed status (progress bar of image pulling) with console.get_console().status(constants.DB_PULL_IMAGE_MSG): self.docker_client.images.pull(self.db_image_tag) except (docker_errors.ImageNotFound, docker_errors.APIError): logging.debug(f"Failed to pull database image ({self.db_image_tag}).", exc_info=True) raise errors.ImageNotFoundError(self.db_image_tag) from None def create_database_files(self): db_extras = self.database_extras db_mounts = db_extras["mounts"] for host_path in db_mounts['files'].keys(): files.create_empty_file(host_path, skip_if_exist=False) for host_path in db_mounts['folders'].keys(): files.create_empty_folder(host_path) def create(self): """Creates Composer local environment. Directory with environment name will be created under `composer` path and environment configuration will be saved to config.json and requirements.txt files. """ assert_image_exists(self.image_version) files.create_environment_directories(self.env_dir_path, self.dags_path) self.create_database_files() self.write_environment_config_to_config_file() self.pypi_packages_to_requirements() self.environment_vars_to_env_file() console.get_console().print( constants.CREATE_MESSAGE.format( env_dir=self.env_dir_path, env_name=self.name, config_path=self.env_dir_path / "config.json", requirements_path=self.env_dir_path / "requirements.txt", env_variables_path=self.env_dir_path / "variables.env", dags_path=self.dags_path, ) ) def assert_container_is_active(self, container_name): """ Asserts docker container is in running or created state (is active). """ status = self.get_container(container_name).status if status not in ( constants.ContainerStatus.RUNNING, constants.ContainerStatus.CREATED, ): raise errors.EnvironmentStartError() def wait_for_db_start(self): start_time = time.time() with console.get_console().status("[bold green]Starting database..."): self.assert_container_is_active(self.db_container_name) for line in self.get_container(self.db_container_name).logs(stream=True, timestamps=True): line = line.decode('utf-8').strip() console.get_console().print(line) if "database system is ready to accept connections" in line: start_duration = time.time() - start_time LOG.info("Database is started in %.2f seconds", start_duration) return if timeout_occurred(start_time): raise errors.EnvironmentStartTimeoutError() self.assert_container_is_active(self.db_container_name) raise errors.EnvironmentStartError() def wait_for_start(self): """ Poll environment logs to see if it is ready. When Airflow scheduler starts, it prints 'searching for files' in the logs. We are using it as marker of the environment readiness. """ start_time = time.time() with console.get_console().status("[bold green]Starting environment..."): self.assert_container_is_active(self.container_name) for line in self.get_container(self.container_name).logs(stream=True, timestamps=True): line = line.decode("utf-8").strip() console.get_console().print(line) # TODO: (b/234684803) Improve detecting container readiness if "Searching for files" in line: start_duration = time.time() - start_time LOG.info( "Environment started in %.2f seconds", start_duration ) return if timeout_occurred(start_time): raise errors.EnvironmentStartTimeoutError() self.assert_container_is_active(self.container_name) raise errors.EnvironmentStartError() def get_or_create_container(self, container_name: str): """ Get existing container or create new container if it does not exist. """ try: return self.get_container(container_name) except errors.EnvironmentNotRunningError: if container_name == self.container_name: # if the given container name is the main container return self.create_docker_container() else: # if the given container name is db container return self.create_db_docker_container() def start_container(self, container_name: str = None, assert_not_running=True): """ Start the given container """ container = self.get_or_create_container(container_name) if ( assert_not_running and container.status == constants.ContainerStatus.RUNNING ): raise errors.EnvironmentAlreadyRunningError(self.name) from None try: container.start() return container except docker.errors.APIError as err: logging.debug( "Starting environment failed with Docker API error.", exc_info=True, ) # TODO: (b/234552960) Test on different OS/language setting if ( err.status_code == constants.SERVER_ERROR_CODE and "port is already allocated" in str(err) ): container.remove() raise errors.ComposerCliError( constants.PORT_IN_USE_ERROR.format(port=self.port) ) error = f"Environment ({container_name}) failed to start with an error: {err}" raise errors.EnvironmentStartError(error) from None def start(self, assert_not_running=True): """Starts local composer environment. Before starting we are asserting that are required files in the environment directory. The docker container is created and started. This operation will raise an error if we are trying to use port that is already allocated. Started environment is polled until Airflow scheduler starts. """ assert_image_exists(self.image_version) self.assert_requirements_exist() files.assert_dag_path_exists(self.dags_path) self.create_database_files() db_path = self.airflow_db if self.is_database_sqlite3 else self.airflow_db_folder files.fix_file_permissions( entrypoint=self.entrypoint_file, run=self.run_file, requirements=self.requirements_file, db_path=db_path, ) files.fix_line_endings( entrypoint=self.entrypoint_file, run=self.run_file, requirements=self.requirements_file, ) if not self.is_database_sqlite3: LOG.info(f"Database engine is selected as {self.database_engine}. The container will start before") db_container = self.start_container(self.db_container_name, False) self.wait_for_db_start() self.ensure_container_is_attached_to_network(db_container) LOG.info(f"Database started!") container = self.start_container(self.container_name, assert_not_running) self.ensure_container_is_attached_to_network(container) self.wait_for_start() self.print_start_message() def ensure_container_is_attached_to_network(self, container): network = self.get_docker_network() existing_containers = [c.name for c in network.containers] if container.name in existing_containers: network.disconnect(container.name) network.connect(container) def print_start_message(self): """Print the start message after the environment is up and ready.""" console.get_console().print( constants.START_MESSAGE.format( env_name=self.name, dags_path=self.dags_path, port=self.port, ) ) def logs(self, follow, max_lines): """ Fetch and print logs from the running composer local environment. Container `logs` method returns blocking generator if follow is True, and byte-decoded string if follow is False. That's why we need two methods of handling and decoding logs. """ log_lines = self.get_container(self.container_name).logs( timestamps=True, stream=follow, follow=follow, tail=max_lines, ) if follow: LOG.debug( "Printing previous %s lines and following output " "from the container logs:", max_lines, ) for line in log_lines: line = line.decode("utf-8").strip() console.get_console().print(line) else: LOG.debug( "Printing previous %s lines from container logs:", max_lines ) log_lines = log_lines.decode("utf-8") for line in log_lines.split("\n"): console.get_console().print(line) def stop(self, remove_container=False): """ Stops the local composer environment. By default container is not removed. """ with console.get_console().status( f"[bold green]Stopping composer local environment..." ): db_container = self.get_container(self.db_container_name, ignore_not_found=True) if db_container: db_container.stop() if remove_container: db_container.remove() container = self.get_container(self.container_name, ignore_not_found=True) if container: container.stop() if remove_container: container.remove() if remove_container: network = self.get_docker_network() network.remove() def restart(self): """ Restarts the local composer environment. This operation will stop and remove container if it is running. Then it will start it again. """ try: self.stop(remove_container=True) except errors.EnvironmentNotRunningError: pass self.start(assert_not_running=False) def status(self) -> str: """Get status of the local composer environment.""" try: return self.get_container(self.container_name).status except errors.EnvironmentNotRunningError: return "Not started" def run_airflow_command(self, command: List) -> None: """ Run command list in the environment container. The commands are prefixed with `airflow`. """ container = self.get_container(self.container_name, assert_running=True) command.insert(0, "airflow") command.insert(0, "/home/airflow/run_as_user.sh") result = container.exec_run(cmd=command) console.get_console().print(result.output.decode()) def get_host_port(self) -> int: """ Return port of the running environment. If it fails to retrieve it, return port from the environment configuration. """ try: return self.get_container(self.container_name).ports["8080/tcp"][0]["HostPort"] except (IndexError, KeyError): LOG.info(constants.FAILED_TO_GET_DOCKER_PORT_WARN) return self.port def prepare_env_description(self, env_status: str) -> str: """Prepare description of the local composer environment.""" if env_status == constants.ContainerStatus.RUNNING: port = self.get_host_port() web_url = constants.WEBSERVER_URL_MESSAGE.format(port=port) else: web_url = "" env_status = utils.wrap_status_in_color(env_status) return constants.DESCRIBE_ENV_MESSAGE.format( name=self.name, state=env_status, web_url=web_url, image_version=self.image_version, dags_path=self.dags_path, gcloud_path=utils.resolve_gcloud_config_path(), ) + (constants.KUBECONFIG_PATH_MESSAGE.format(kube_config_path=utils.resolve_kube_config_path()) if utils.resolve_kube_config_path() else "no file") + constants.FINAL_ENV_MESSAGE def describe(self) -> None: """Describe the local composer environment.""" env_status = self.status() desc = self.prepare_env_description(env_status) console.get_console().print(desc) def remove(self, force, force_error): containers = {self.container_name} if not self.is_database_sqlite3: containers.add(self.db_container_name) for container_name in containers: container = self.get_container(container_name, ignore_not_found=True) if container is not None: if container.status == constants.ContainerStatus.RUNNING: if not force: raise force_error container.stop() container.remove() network = self.get_docker_network() network.remove()