content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import uuid
  from datetime import datetime, timedelta
  
  from airflow import DAG  # noqa
  from airflow import macros  # noqa
  from airflow.operators.python_operator import PythonOperator  # noqa
  from elasticsearch import Elasticsearch
  from pyhocon import ConfigFactory
  
  from databuilder.extractor.athena_metadata_extractor import AthenaMetadataExtractor
  from databuilder.extractor.neo4j_extractor import Neo4jExtractor
  from databuilder.extractor.neo4j_search_data_extractor import Neo4jSearchDataExtractor
  from databuilder.extractor.sql_alchemy_extractor import SQLAlchemyExtractor
  from databuilder.job.job import DefaultJob
  from databuilder.loader.file_system_elasticsearch_json_loader import FSElasticsearchJSONLoader
  from databuilder.loader.file_system_neo4j_csv_loader import FsNeo4jCSVLoader
  from databuilder.publisher import neo4j_csv_publisher
  from databuilder.publisher.elasticsearch_publisher import ElasticsearchPublisher
  from databuilder.publisher.neo4j_csv_publisher import Neo4jCsvPublisher
  from databuilder.task.task import DefaultTask
  from databuilder.transformer.base_transformer import NoopTransformer
  
  dag_args = {
      'concurrency': 10,
      # One dagrun at a time
      'max_active_runs': 1,
      # 4AM, 4PM PST
      'schedule_interval': '0 11 * * *',
      'catchup': False
  }
  
  default_args = {
      'owner': 'amundsen',
      'start_date': datetime(2018, 6, 18),
      'depends_on_past': False,
      'email': [''],
      'email_on_failure': False,
      'email_on_retry': False,
      'retries': 3,
      'priority_weight': 10,
      'retry_delay': timedelta(minutes=5),
      'execution_timeout': timedelta(minutes=120)
  }
  
  # NEO4J cluster endpoints
  NEO4J_ENDPOINT = 'bolt://127.0.0.1:7687'
  
  neo4j_endpoint = NEO4J_ENDPOINT
  
  neo4j_user = 'neo4j'
  neo4j_password = 'test'
  
  es = Elasticsearch([
      {'host': '127.0.0.1'},
  ])
  
  # TODO: user provides a list of schema for indexing
  SUPPORTED_SCHEMAS = ['sampledb']
  # String format - ('schema1', schema2', .... 'schemaN')
  SUPPORTED_SCHEMA_SQL_IN_CLAUSE = "('{schemas}')".format(schemas="', '".join(SUPPORTED_SCHEMAS))
  
  OPTIONAL_TABLE_NAMES = ''
  AWS_ACCESS = 'YOUR_ACCESS_KEY'
  AWS_SECRET = 'YOUR_SECRET_KEY'
  
  
  def connection_string():
      access_key = AWS_ACCESS
      secret = AWS_SECRET
      host = 'athena.us-east-1.amazonaws.com'
      extras = 's3_staging_dir=s3://aws-athena-query-results-032106861074-us-east-1/'
      return "awsathena+rest://%s:%s@%s:443/?%s" % (access_key, secret, host, extras)
  
  
  def create_table_extract_job():
      where_clause_suffix = f"where table_schema in {SUPPORTED_SCHEMA_SQL_IN_CLAUSE}"
  
      tmp_folder = '/var/tmp/amundsen/table_metadata'
      node_files_folder = f'{tmp_folder}/nodes/'
      relationship_files_folder = f'{tmp_folder}/relationships/'
  
      job_config = ConfigFactory.from_dict({
          f'extractor.athena_metadata.{AthenaMetadataExtractor.WHERE_CLAUSE_SUFFIX_KEY}': where_clause_suffix,
          f'extractor.athena_metadata.extractor.sqlalchemy.{SQLAlchemyExtractor.CONN_STRING}': connection_string(),
          f'extractor.athena_metadata.{AthenaMetadataExtractor.CATALOG_KEY}': "'AwsDataCatalog'",
          f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.NODE_DIR_PATH}': node_files_folder,
          f'loader.filesystem_csv_neo4j.{FsNeo4jCSVLoader.RELATION_DIR_PATH}': relationship_files_folder,
          f'publisher.neo4j.{neo4j_csv_publisher.NODE_FILES_DIR}': node_files_folder,
          f'publisher.neo4j.{neo4j_csv_publisher.RELATION_FILES_DIR}': relationship_files_folder,
          f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_END_POINT_KEY}': neo4j_endpoint,
          f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_USER}': neo4j_user,
          f'publisher.neo4j.{neo4j_csv_publisher.NEO4J_PASSWORD}': neo4j_password,
          f'publisher.neo4j.{neo4j_csv_publisher.JOB_PUBLISH_TAG}': 'unique_tag',  # should use unique tag here like {ds}
      })
      job = DefaultJob(conf=job_config,
                       task=DefaultTask(extractor=AthenaMetadataExtractor(), loader=FsNeo4jCSVLoader(),
                                        transformer=NoopTransformer()),
                       publisher=Neo4jCsvPublisher())
      job.launch()
  
  
  def create_es_publisher_sample_job():
      # loader saves data to this location and publisher reads it from here
      extracted_search_data_path = '/var/tmp/amundsen/search_data.json'
  
      task = DefaultTask(loader=FSElasticsearchJSONLoader(),
                         extractor=Neo4jSearchDataExtractor(),
                         transformer=NoopTransformer())
  
      # elastic search client instance
      elasticsearch_client = es
      # unique name of new index in Elasticsearch
      elasticsearch_new_index_key = 'tables' + str(uuid.uuid4())
      # related to mapping type from /databuilder/publisher/elasticsearch_publisher.py#L38
      elasticsearch_new_index_key_type = 'table'
      # alias for Elasticsearch used in amundsensearchlibrary/search_service/config.py as an index
      elasticsearch_index_alias = 'table_search_index'
  
      job_config = ConfigFactory.from_dict({
          f'extractor.search_data.extractor.neo4j.{Neo4jExtractor.GRAPH_URL_CONFIG_KEY}': neo4j_endpoint,
          f'extractor.search_data.extractor.neo4j.{Neo4jExtractor.MODEL_CLASS_CONFIG_KEY}':
              'databuilder.models.table_elasticsearch_document.TableESDocument',
          f'extractor.search_data.extractor.neo4j.{Neo4jExtractor.NEO4J_AUTH_USER}': neo4j_user,
          f'extractor.search_data.extractor.neo4j.{Neo4jExtractor.NEO4J_AUTH_PW}': neo4j_password,
          f'loader.filesystem.elasticsearch.{FSElasticsearchJSONLoader.FILE_PATH_CONFIG_KEY}': extracted_search_data_path,
          f'loader.filesystem.elasticsearch.{FSElasticsearchJSONLoader.FILE_MODE_CONFIG_KEY}': 'w',
          f'publisher.elasticsearch.{ElasticsearchPublisher.FILE_PATH_CONFIG_KEY}': extracted_search_data_path,
          f'publisher.elasticsearch.{ElasticsearchPublisher.FILE_MODE_CONFIG_KEY}': 'r',
          f'publisher.elasticsearch.{ElasticsearchPublisher.ELASTICSEARCH_CLIENT_CONFIG_KEY}':
              elasticsearch_client,
          f'publisher.elasticsearch.{ElasticsearchPublisher.ELASTICSEARCH_NEW_INDEX_CONFIG_KEY}':
              elasticsearch_new_index_key,
          f'publisher.elasticsearch.{ElasticsearchPublisher.ELASTICSEARCH_DOC_TYPE_CONFIG_KEY}':
              elasticsearch_new_index_key_type,
          f'publisher.elasticsearch.{ElasticsearchPublisher.ELASTICSEARCH_ALIAS_CONFIG_KEY}':
              elasticsearch_index_alias
      })
  
      job = DefaultJob(conf=job_config,
                       task=task,
                       publisher=ElasticsearchPublisher())
      job.launch()
  
  
  with DAG('amundsen_databuilder', default_args=default_args, **dag_args) as dag:
      athena_table_extract_job = PythonOperator(
          task_id='athena_table_extract_job',
          python_callable=create_table_extract_job
      )
  
      athena_es_publisher_job = PythonOperator(
          task_id='athena_es_publisher_job',
          python_callable=create_es_publisher_sample_job
      )
  
      # elastic search update run after table metadata has been updated
      athena_table_extract_job >> athena_es_publisher_job
  