用于导出、加载和解析以太坊区块链数据的 Airflow DAG_python_代码_下载

from __future__ import print_function import json import logging import os import time from datetime import datetime, timedelta from tempfile import TemporaryDirectory from airflow import models from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook from airflow.contrib.operators.bigquery_operator import BigQueryOperator from airflow.contrib.sensors.gcs_sensor import GoogleCloudStorageObjectSensor from airflow.operators import python_operator from airflow.operators.email_operator import EmailOperator from airflow.operators.python_operator import PythonOperator from google.cloud import bigquery from google.cloud.bigquery import TimePartitioning, SchemaField from ethereumetl_airflow.bigquery_utils import submit_bigquery_job from ethereumetl_airflow.build_export_dag import upload_to_gcs logging.basicConfig() logging.getLogger().setLevel(logging.DEBUG) def build_load_dag( dag_id, output_bucket, destination_dataset_project_id, chain='ethereum', notification_emails=None, load_start_date=datetime(2018, 7, 1), schedule_interval='0 0 * * *', load_all_partitions=True ): # The following datasets must be created in BigQuery: # - crypto_{chain}_raw # - crypto_{chain}_temp # - crypto_{chain} dataset_name = f'crypto_{chain}' dataset_name_raw = f'crypto_{chain}_raw' dataset_name_temp = f'crypto_{chain}_temp' if not destination_dataset_project_id: raise ValueError('destination_dataset_project_id is required') environment = { 'dataset_name': dataset_name, 'dataset_name_raw': dataset_name_raw, 'dataset_name_temp': dataset_name_temp, 'destination_dataset_project_id': destination_dataset_project_id, 'load_all_partitions': load_all_partitions } def read_bigquery_schema_from_file(filepath): result = [] file_content = read_file(filepath) json_content = json.loads(file_content) for field in json_content: result.append(bigquery.SchemaField( name=field.get('name'), field_type=field.get('type', 'STRING'), mode=field.get('mode', 'NULLABLE'), description=field.get('description'))) return result def read_file(filepath): with open(filepath) as file_handle: content = file_handle.read() return content default_dag_args = { 'depends_on_past': False, 'start_date': load_start_date, 'email_on_failure': True, 'email_on_retry': False, 'retries': 5, 'retry_delay': timedelta(minutes=5) } if notification_emails and len(notification_emails) > 0: default_dag_args['email'] = [email.strip() for email in notification_emails.split(',')] # Define a DAG (directed acyclic graph) of tasks. dag = models.DAG( dag_id, catchup=False, schedule_interval=schedule_interval, default_args=default_dag_args) dags_folder = os.environ.get('DAGS_FOLDER', '/home/airflow/gcs/dags') def add_load_tasks(task, file_format, allow_quoted_newlines=False): wait_sensor = GoogleCloudStorageObjectSensor( task_id='wait_latest_{task}'.format(task=task), timeout=60 * 60, poke_interval=60, bucket=output_bucket, object='export/{task}/block_date={datestamp}/{task}.{file_format}'.format( task=task, datestamp='{{ds}}', file_format=file_format), dag=dag ) def load_task(ds, **kwargs): client = bigquery.Client() job_config = bigquery.LoadJobConfig() schema_path = os.path.join(dags_folder, 'resources/stages/raw/schemas/{task}.json'.format(task=task)) schema = read_bigquery_schema_from_file(schema_path) schema = adjust_schema_for_kovan(dag_id, task, schema) job_config.schema = schema job_config.source_format = bigquery.SourceFormat.CSV if file_format == 'csv' else bigquery.SourceFormat.NEWLINE_DELIMITED_JSON if file_format == 'csv': job_config.skip_leading_rows = 1 job_config.write_disposition = 'WRITE_TRUNCATE' job_config.allow_quoted_newlines = allow_quoted_newlines job_config.ignore_unknown_values = True export_location_uri = 'gs://{bucket}/export'.format(bucket=output_bucket) if load_all_partitions: uri = '{export_location_uri}/{task}/*.{file_format}'.format( export_location_uri=export_location_uri, task=task, file_format=file_format) else: uri = '{export_location_uri}/{task}/block_date={ds}/*.{file_format}'.format( export_location_uri=export_location_uri, task=task, ds=ds, file_format=file_format) table_ref = client.dataset(dataset_name_raw).table(task) load_job = client.load_table_from_uri(uri, table_ref, job_config=job_config) submit_bigquery_job(load_job, job_config) assert load_job.state == 'DONE' load_operator = PythonOperator( task_id='load_{task}'.format(task=task), python_callable=load_task, provide_context=True, execution_timeout=timedelta(minutes=30), dag=dag ) wait_sensor >> load_operator return load_operator def add_enrich_tasks(task, time_partitioning_field='block_timestamp', dependencies=None, always_load_all_partitions=False): def enrich_task(ds, **kwargs): template_context = kwargs.copy() template_context['ds'] = ds template_context['params'] = environment client = bigquery.Client() # Need to use a temporary table because bq query sets field modes to NULLABLE and descriptions to null # when writeDisposition is WRITE_TRUNCATE # Create a temporary table temp_table_name = '{task}_{milliseconds}'.format(task=task, milliseconds=int(round(time.time() * 1000))) temp_table_ref = client.dataset(dataset_name_temp).table(temp_table_name) schema_path = os.path.join(dags_folder, 'resources/stages/enrich/schemas/{task}.json'.format(task=task)) schema = read_bigquery_schema_from_file(schema_path) schema = adjust_schema_for_kovan(dag_id, task, schema) table = bigquery.Table(temp_table_ref, schema=schema) description_path = os.path.join( dags_folder, 'resources/stages/enrich/descriptions/{task}.txt'.format(task=task)) table.description = read_file(description_path) if time_partitioning_field is not None: table.time_partitioning = TimePartitioning(field=time_partitioning_field) logging.info('Creating table: ' + json.dumps(table.to_api_repr())) table = client.create_table(table) assert table.table_id == temp_table_name # Query from raw to temporary table query_job_config = bigquery.QueryJobConfig() # Finishes faster, query limit for concurrent interactive queries is 50 query_job_config.priority = bigquery.QueryPriority.INTERACTIVE query_job_config.destination = temp_table_ref sql_path = os.path.join(dags_folder, 'resources/stages/enrich/sqls/{task}.sql'.format(task=task)) sql_template = read_file(sql_path) sql = kwargs['task'].render_template(sql_template, template_context) print('Enrichment sql:') print(sql) query_job = client.query(sql, location='US', job_config=query_job_config) submit_bigquery_job(query_job, query_job_config) assert query_job.state == 'DONE' if load_all_partitions or always_load_all_partitions: # Copy temporary table to destination co