def set_settings_from_remote(remote: str): if remote: if remote not in Remotes.remotes_config.keys(): print( f'Remote {remote} is not defined in .typhoonremotes. Found : {list(Remotes.remotes_config.keys())}', file=sys.stderr) sys.exit(-1) Settings.metadata_db_url = Remotes.metadata_db_url(remote) if Settings.metadata_store( Remotes.aws_profile(remote)).name == 'airflow': Settings.fernet_key = Remotes.fernet_key(remote) if Remotes.use_name_as_suffix(remote): Settings.metadata_suffix = remote
def add_variable(remote: Optional[str], var_id: str, var_type: str, contents): """Add variable to the metadata store""" set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) var = Variable(var_id, VariableType[var_type.upper()], contents) metadata_store.set_variable(var) print(f'Variable {var_id} added')
def list_dags(remote: Optional[str], long: bool, json_output): set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) if long: header = ['DAG_NAME', 'DEPLOYMENT_DATE'] table_body = [[x.dag_name, x.deployment_date.isoformat()] for x in metadata_store.get_dag_deployments()] print(tabulate(table_body, header, 'plain')) if not remote: dag_errors = get_dag_errors() if dag_errors: header = ['DAG_NAME', 'ERROR LOCATION', 'ERROR MESSAGE'] table_body = [[dag_name, error[0]['loc'], error[0]['msg']] for dag_name, error in dag_errors.items()] print(colored(tabulate(table_body, header, 'plain'), 'red'), file=sys.stderr) else: dag_names = sorted( set(x.dag_name for x in metadata_store.get_dag_deployments())) if json_output: print(json.dumps(dag_names)) else: for dag_name in dag_names: print(dag_name) if not remote: for dag_name, _ in get_dag_errors().items(): print(colored(dag_name, 'red'), file=sys.stderr)
def load_variable(remote: Optional[str], file: str): """Read variable from file and add it to the metadata store""" var = Variable.from_file(file) set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) metadata_store.set_variable(var) print(f'Variable {var.id} added')
def remote_list(long: bool): """List configured Typhoon remotes""" if long: header = [ 'REMOTE_NAME', 'AWS_PROFILE', 'USE_NAME_AS_SUFFIX', 'METADATA_DB_URL' ] table_body = [[ remote, Remotes.aws_profile(remote), Remotes.use_name_as_suffix(remote), Remotes.metadata_db_url(remote) ] for remote in Remotes.remote_names] print(tabulate(table_body, header, 'plain')) else: for remote in Remotes.remote_names: print(remote)
def add_connection(remote: Optional[str], conn_id: str, conn_env: str): """Add connection to the metadata store""" set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) conn_params = connections.get_connection_local(conn_id, conn_env) metadata_store.set_connection( Connection(conn_id=conn_id, **asdict(conn_params))) print(f'Connection {conn_id} added')
def get_undefined_variables_in_metadata_db(remote: Optional[str], var_ids: List[str]): undefined_variables = [] for var_id in var_ids: try: Settings.metadata_store( Remotes.aws_profile(remote)).get_variable(var_id) except MetadataObjectNotFound: undefined_variables.append(var_id) return undefined_variables
def get_undefined_connections_in_metadata_db(remote: Optional[str], conn_ids: List[str]): undefined_connections = [] for conn_id in conn_ids: try: Settings.metadata_store( Remotes.aws_profile(remote)).get_connection(conn_id) except MetadataObjectNotFound: undefined_connections.append(conn_id) return undefined_connections
def list_connections(remote: Optional[str], long: bool): """List connections in the metadata store""" set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) if long: header = ['CONN_ID', 'TYPE', 'HOST', 'PORT', 'SCHEMA'] table_body = [[ conn.conn_id, conn.conn_type, conn.host, conn.port, conn.schema ] for conn in metadata_store.get_connections()] print(tabulate(table_body, header, 'plain')) else: for conn in metadata_store.get_connections(): print(conn.conn_id)
def generate_terraform_files(remote: Optional[str] = None, force: bool = False, minimal: bool = False): terraform_dest_folder = (Settings.typhoon_home / 'terraform') if not force and terraform_dest_folder.exists(): print( f'Cannot generate terraform files because the folder exists already {terraform_dest_folder}' ) print('Run with -f/--force to force overwrite') exit(1) main_tf_file = TERRAFORM_FOLDER_PATH / 'main.tf' env_name = remote or 'test' tfvars_template = jinja2.Template( (TERRAFORM_FOLDER_PATH / 'env.tfvars.j2').read_text()) rendered_tfvars = tfvars_template.render( dict( env=env_name, runtime='python{}.{}'.format(*sys.version_info), connections_table=None if minimal else Settings.connections_table_name, variables_table=None if minimal else Settings.variables_table_name, dag_deployments_table=None if minimal else Settings.dag_deployments_table_name, metadata_db_url=None if minimal else Remotes.metadata_db_url(remote), metadata_suffix=None if minimal else Settings.metadata_suffix, s3_bucket=Remotes.s3_bucket(remote) if remote else '', project_name=Settings.project_name, )) terraform_dest_folder.mkdir(exist_ok=True) shutil.copy(str(main_tf_file), str(terraform_dest_folder)) if not minimal: metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) metadata_store_tf_file = TERRAFORM_FOLDER_PATH / f'metadata_stores/{metadata_store.name}.tf' shutil.copy(str(metadata_store_tf_file), str(terraform_dest_folder)) (terraform_dest_folder / f'{env_name}.tfvars').write_text(rendered_tfvars)
def dags_without_deploy(remote: Optional[str]) -> List[str]: add_yaml_constructors() undeployed_dags = [] for dag_code in get_dags_contents(Settings.dags_directory): loaded_dag = yaml.load(dag_code, yaml.FullLoader) dag_deployment = DagDeployment(dag_name=loaded_dag['name'], deployment_date=datetime.utcnow(), dag_code=dag_code) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) if loaded_dag.get('active', True): try: _ = metadata_store.get_dag_deployment( dag_deployment.deployment_hash) except MetadataObjectNotFound: undeployed_dags.append(dag_deployment.dag_name) return undeployed_dags
def list_variables(remote: Optional[str], long: bool): """List variables in the metadata store""" def var_contents(var: Variable) -> str: if var.type == VariableType.NUMBER: return var.contents else: return f'"{var.contents}"' if len( var.contents ) < max_len_var else f'"{var.contents[:max_len_var]}"...' set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) if long: max_len_var = 40 header = ['VAR_ID', 'TYPE', 'CONTENT'] table_body = [[var.id, var.type, var_contents(var)] for var in metadata_store.get_variables()] print(tabulate(table_body, header, 'plain')) else: for var in metadata_store.get_variables(): print(var.id)
def remote_add(remote: str): """Remove remote""" Remotes.remove_remote(remote) print(f'Removed remote {remote}')
def remove_connection(remote: Optional[str], conn_id: str): """Remove connection from the metadata store""" set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) metadata_store.delete_connection(conn_id) print(f'Connection {conn_id} deleted')
def remove_variable(remote: Optional[str], var_id: str): """Remove connection from the metadata store""" set_settings_from_remote(remote) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) metadata_store.delete_variable(var_id) print(f'Variable {var_id} deleted')
def status(remote: Optional[str]): """Information on project status""" set_settings_from_remote(remote) print(colored(ascii_art_logo, 'cyan')) if not Settings.typhoon_home: print(colored(f'FATAL: typhoon home not found...', 'red')) return else: print(colored('• Typhoon home defined as', 'green'), colored(Settings.typhoon_home, 'blue')) metadata_store = Settings.metadata_store(Remotes.aws_profile(remote)) if metadata_store.exists(): print(colored('• Metadata database found in', 'green'), colored(Settings.metadata_db_url, 'blue')) check_connections_yaml(remote) check_connections_dags(remote) check_variables_dags(remote) elif isinstance(metadata_store, SQLiteMetadataStore): print(colored('• Metadata store not found for', 'yellow'), colored(Settings.metadata_db_url, 'blue')) print( colored( ' - It will be created upon use, or create by running (idempotent) command', color=None), colored(f'typhoon migrate{" " + remote if remote else ""}', 'blue')) print(colored(' Skipping connections and variables checks...', 'red')) else: print(colored('• Metadata store not found or incomplete for', 'red'), colored(Settings.metadata_db_url, 'blue')) print( colored(' - Fix by running (idempotent) command', color=None), colored( f'typhoon metadata migrate{" " + remote if remote else ""}', 'blue')) print(colored(' Skipping connections and variables checks...', 'red')) if not remote: changed_dags = dags_with_changes() if changed_dags: print( colored('• Unbuilt changes in DAGs... To rebuild run', 'yellow'), colored( f'typhoon dag build{" " + remote if remote else ""} --all [--debug]', 'blue')) for dag in changed_dags: print(colored(f' - {dag}', 'blue')) else: print(colored('• DAGs up to date', 'green')) else: undeployed_dags = dags_without_deploy(remote) if undeployed_dags: print( colored('• Undeployed changes in DAGs... To deploy run', 'yellow'), colored( f'typhoon dag push {remote} --all [--build-dependencies]', 'blue')) for dag in undeployed_dags: print(colored(f' - {dag}', 'blue')) else: print(colored('• DAGs up to date', 'green'))
def remote_add(remote: str, aws_profile: str, metadata_db_url: str, use_name_as_suffix: bool): """Add a remote for deployments and management""" Remotes.add_remote(remote, aws_profile, metadata_db_url, use_name_as_suffix) print(f'Added remote {remote}')
def push_dags(remote: Optional[str], dag_name: Optional[str], all_: bool, code: bool, build_deps_locally: bool): """Build code for dags in $TYPHOON_HOME/out/""" set_settings_from_remote(remote) if dag_name and all_: raise click.UsageError( f'Illegal usage: DAG_NAME is mutually exclusive with --all') elif dag_name is None and not all_: raise click.UsageError(f'Illegal usage: Need either DAG_NAME or --all') if all_: dag_errors = get_dag_errors() if dag_errors: print(f'Found errors in the following DAGs:') for dag_name in dag_errors.keys(): print(f' - {dag_name}\trun typhoon dag build {dag_name}') sys.exit(-1) raise NotImplementedError() else: if Settings.deploy_target == 'typhoon': build_all_dags(remote=remote, matching=dag_name) dag_out_path = Settings.out_directory / dag_name builds_path = Settings.out_directory / 'builds' if builds_path.exists(): shutil.rmtree(str(builds_path), ignore_errors=True) builds_path.mkdir() (builds_path / dag_name).mkdir() shutil.make_archive(str(builds_path / dag_name / 'lambda'), 'zip', root_dir=str(dag_out_path)) print('Updating lambda code...') import boto3 session = boto3.session.Session( profile_name=Remotes.aws_profile(remote)) s3r = session.resource('s3') s3_key_lambda_zip = f'typhoon_dag_builds/{dag_name}/lambda.zip' s3r.Object(Remotes.s3_bucket(remote), s3_key_lambda_zip).put( Body=(builds_path / f'{dag_name}/lambda.zip').open('rb')) lambdac = session.client('lambda') try: lambdac.update_function_code( FunctionName=f'{dag_name}_{remote}', S3Bucket=Remotes.s3_bucket(remote), S3Key=s3_key_lambda_zip, ) except lambdac.exceptions.ResourceNotFoundException: print( f'Lambda for {dag_name} does not exist yet. Create it with terraform after this command' ) if not code: print('Building lambda dependencies...') if typhoon_version_is_local(): p = local_typhoon_path().rstrip('typhoon') local_typhoon_volume = ['-v', f'{p}:{p}'] else: local_typhoon_volume = [] if build_deps_locally: pip_command = f'{sys.executable} -m pip install -r requirements.txt --target .layer/python/' print(pip_command) p = subprocess.run(args=pip_command, cwd=str(dag_out_path), shell=True) print( colored( f'ERROR: Pip install command exited with return code {p.returncode}', 'red')) sys.exit(1) else: docker_pip_command = [ 'docker', 'run', '--rm', '-v', f'{dag_out_path}:/var/task', *local_typhoon_volume, 'lambci/lambda:build-python{}.{}'.format( *sys.version_info), 'pip', 'install', '-r', 'requirements.txt', '--target', './layer/python/', ] docker_pip_command = ' '.join(docker_pip_command) print(docker_pip_command) p = subprocess.run(args=docker_pip_command, cwd=str(dag_out_path), shell=True) if p.returncode != 0: print( colored( f'ERROR: Docker command exited with return code {p.returncode}', 'red')) print( colored( 'Check that docker is installed and running on your machine', 'red')) print( colored( 'Alternatively run this command with --build-deps-locally to skip docker, but resulting binaries may be incompatible with the lambda runtime.', 'red')) sys.exit(1) shutil.make_archive(str(builds_path / dag_name / 'layer'), 'zip', root_dir=str(dag_out_path / 'layer')) s3_key_layer_zip = f'typhoon_dag_builds/{dag_name}/layer.zip' s3r.Object(Remotes.s3_bucket(remote), s3_key_layer_zip).put( Body=(builds_path / f'{dag_name}/layer.zip').open('rb')) print('Publishing dependencies as layer...') layer_name = f'{dag_name}_dependencies' response = lambdac.publish_layer_version( LayerName=layer_name, Description='Dependencies needed for DAG', Content={ 'S3Bucket': Remotes.s3_bucket(remote), 'S3Key': str(s3_key_layer_zip), }, CompatibleRuntimes=[ 'python{}.{}'.format(*sys.version_info) ]) try: lambdac.update_function_configuration( FunctionName=f'{dag_name}_{remote}', Layers=[response['LayerVersionArn']], ) except lambdac.exceptions.ResourceNotFoundException: print( colored( f'WARNING: Lambda for {dag_name} does not exist yet. Create it with terraform after this command', 'yellow')) else: raise NotImplementedError()
def migrate(remote: str): """Create the necessary metadata tables""" set_settings_from_remote(remote) print(f'Migrating {Settings.metadata_db_url}...') Settings.metadata_store(aws_profile=Remotes.aws_profile(remote)).migrate()
def generate_settings_file(remote: str): metadata_db_url = Remotes.metadata_db_url(remote) metadata_suffix = Settings.metadata_suffix return f"""\