content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  """
  This is a example script demonstrating how to load data into Neo4j and
  Elasticsearch without using an Airflow DAG.
  It contains several jobs:
  - `run_csv_job`: runs a job that extracts table data from a CSV, loads (writes)
    this into a different local directory as a csv, then publishes this data to
    neo4j.
  - `run_table_column_job`: does the same thing as `run_csv_job`, but with a csv
    containing column data.
  - `create_last_updated_job`: creates a job that gets the current time, dumps it
    into a predefined model schema, and publishes this to neo4j.
  - `create_es_publisher_sample_job`: creates a job that extracts data from neo4j
    and pubishes it into elasticsearch.
  For other available extractors, please take a look at
  https://github.com/amundsen-io/amundsendatabuilder#list-of-extractors
  """
  
  import json
  import logging
  import os
  import sys
  import uuid
  
  from elasticsearch import Elasticsearch
  from pyhocon import ConfigFactory
  from sqlalchemy.ext.declarative import declarative_base
  
  from databuilder.extractor.dbt_extractor import DbtExtractor
  from databuilder.extractor.neo4j_search_data_extractor import Neo4jSearchDataExtractor
  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.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
  
  es_host = os.getenv('CREDENTIALS_ELASTICSEARCH_PROXY_HOST', 'localhost')
  neo_host = os.getenv('CREDENTIALS_NEO4J_PROXY_HOST', 'localhost')
  
  es_port = os.getenv('CREDENTIALS_ELASTICSEARCH_PROXY_PORT', 9200)
  neo_port = os.getenv('CREDENTIALS_NEO4J_PROXY_PORT', 7687)
  if len(sys.argv) > 1:
      es_host = sys.argv[1]
  if len(sys.argv) > 2:
      neo_host = sys.argv[2]
  
  es = Elasticsearch([
      {'host': es_host, 'port': es_port},
  ])
  
  Base = declarative_base()
  
  NEO4J_ENDPOINT = f'bolt://{neo_host}:{neo_port}'
  
  neo4j_endpoint = NEO4J_ENDPOINT
  
  neo4j_user = 'neo4j'
  neo4j_password = 'test'
  
  LOGGER = logging.getLogger(__name__)
  
  
  def run_dbt_job(database_name, catalog_file_loc, manifest_file_loc, source_url=None):
      tmp_folder = '/var/tmp/amundsen/dbt_run'
      node_files_folder = f'{tmp_folder}/nodes'
      relationship_files_folder = f'{tmp_folder}/relationships'
  
      dbt_extractor = DbtExtractor()
      csv_loader = FsNeo4jCSVLoader()
  
      task = DefaultTask(extractor=dbt_extractor,
                         loader=csv_loader,
                         transformer=NoopTransformer())
  
      # Catalop and manifest files can be passed in as file locations or a valid python
      # dict, allowing you to retrieve the files from S3 or another source and pass it in
      with open(manifest_file_loc, 'rb') as f:
          manifest_data = json.load(f)
  
      job_config = ConfigFactory.from_dict({
          'extractor.dbt.database_name': database_name,
          'extractor.dbt.catalog_json': catalog_file_loc,  # File
          'extractor.dbt.manifest_json': json.dumps(manifest_data),  # JSON Dumped objecy
          'extractor.dbt.source_url': source_url,
          'loader.filesystem_csv_neo4j.node_dir_path': node_files_folder,
          'loader.filesystem_csv_neo4j.relationship_dir_path': relationship_files_folder,
          'loader.filesystem_csv_neo4j.delete_created_directories': True,
          'publisher.neo4j.node_files_directory': node_files_folder,
          'publisher.neo4j.relation_files_directory': relationship_files_folder,
          'publisher.neo4j.neo4j_endpoint': neo4j_endpoint,
          'publisher.neo4j.neo4j_user': neo4j_user,
          'publisher.neo4j.neo4j_password': neo4j_password,
          'publisher.neo4j.neo4j_encrypted': False,
          'publisher.neo4j.job_publish_tag': 'unique_tag',  # should use unique tag here like {ds}
      })
  
      DefaultJob(conf=job_config,
                 task=task,
                 publisher=Neo4jCsvPublisher()).launch()
  
  
  def create_es_publisher_sample_job(elasticsearch_index_alias='table_search_index',
                                     elasticsearch_doc_type_key='table',
                                     model_name='databuilder.models.table_elasticsearch_document.TableESDocument',
                                     entity_type='table',
                                     elasticsearch_mapping=None):
      """
      :param elasticsearch_index_alias:  alias for Elasticsearch used in
                                         amundsensearchlibrary/search_service/config.py as an index
      :param elasticsearch_doc_type_key: name the ElasticSearch index is prepended with. Defaults to `table` resulting in
                                         `table_{uuid}`
      :param model_name:                 the Databuilder model class used in transporting between Extractor and Loader
      :param entity_type:                Entity type handed to the `Neo4jSearchDataExtractor` class, used to determine
                                         Cypher query to extract data from Neo4j. Defaults to `table`.
      :param elasticsearch_mapping:      Elasticsearch field mapping "DDL" handed to the `ElasticsearchPublisher` class,
                                         if None is given (default) it uses the `Table` query baked into the Publisher
      """
      # 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 = f'{elasticsearch_doc_type_key}_{uuid.uuid4()}'
  
      job_config = ConfigFactory.from_dict({
          'extractor.search_data.entity_type': entity_type,
          'extractor.search_data.extractor.neo4j.graph_url': neo4j_endpoint,
          'extractor.search_data.extractor.neo4j.model_class': model_name,
          'extractor.search_data.extractor.neo4j.neo4j_auth_user': neo4j_user,
          'extractor.search_data.extractor.neo4j.neo4j_auth_pw': neo4j_password,
          'extractor.search_data.extractor.neo4j.neo4j_encrypted': False,
          'loader.filesystem.elasticsearch.file_path': extracted_search_data_path,
          'loader.filesystem.elasticsearch.mode': 'w',
          'publisher.elasticsearch.file_path': extracted_search_data_path,
          'publisher.elasticsearch.mode': 'r',
          'publisher.elasticsearch.client': elasticsearch_client,
          'publisher.elasticsearch.new_index': elasticsearch_new_index_key,
          'publisher.elasticsearch.doc_type': elasticsearch_doc_type_key,
          'publisher.elasticsearch.alias': elasticsearch_index_alias,
      })
  
      # only optionally add these keys, so need to dynamically `put` them
      if elasticsearch_mapping:
          job_config.put(f'publisher.elasticsearch.{ElasticsearchPublisher.ELASTICSEARCH_MAPPING_CONFIG_KEY}',
                         elasticsearch_mapping)
  
      job = DefaultJob(conf=job_config,
                       task=task,
                       publisher=ElasticsearchPublisher())
      return job
  
  
  if __name__ == "__main__":
      # Uncomment next line to get INFO level logging
      # logging.basicConfig(level=logging.INFO)
  
      run_dbt_job(
          database_name='snowflake',
          catalog_file_loc='example/sample_data/dbt/catalog.json',
          manifest_file_loc='example/sample_data/dbt/manifest.json',
          source_url='https://github.com/your-company/your-repo/tree/main'
      )
  
      job_es_table = create_es_publisher_sample_job(
          elasticsearch_index_alias='table_search_index',
          elasticsearch_doc_type_key='table',
          entity_type='table',
          model_name='databuilder.models.table_elasticsearch_document.TableESDocument')
      job_es_table.launch()
  