Notes

"""
### Copy cloudsql tables to bigquery
"""
import os
import logging
from airflow import DAG
from datetime import datetime, timedelta
from airflow.contrib.operators.postgres_to_gcs import PostgresToGoogleCloudStorageOperator
from airflow.contrib.operators.gcs_to_bq import GoogleCloudStorageToBigQueryOperator

default_args = {
'owner': 'rohit',
'depends_on_past': False,
'start_date': '2020-10-28', # Change this to a literal date to make scheduler work.
'email': ['[email protected]'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 2,
'retry_delay': timedelta(minutes=1),
}

dag = DAG('etl_cloudsql_to_bigquery',
default_args=default_args,
schedule_interval='0 * * * *')
dag.doc_md = __doc__


class TableConfig:
def __init__(self,
cloud_sql_instance,
export_bucket,
export_database,
export_table,
export_query,
gcp_project,
stage_dataset,
stage_table,
stage_final_query,
bq_location
):

self.params = {
'export_table': export_table,
'export_bucket': export_bucket,
'export_database': export_database,
'export_query': export_query,
'gcp_project': gcp_project,
'stage_dataset': stage_dataset,
'stage_table': stage_table or export_table,
'stage_final_query': stage_final_query,
'cloud_sql_instance': cloud_sql_instance,
'bq_location': bq_location or "EU",
}

def get_tables():
"""
return a list of tables that should go from cloud sql to bigquery
All tables reside in the same cloud sql instance.
:return:
"""
dim_tables = ["company_stages_enum", "connection_status_enum", "component_type_enum", "lifecycle_phase_enum",
"manufacturing_methods_enum", "notification_action_types", "notification_entity_types", "organization_roles_enum",
"payment_terms_enum", "proposal_status_enum"]
fact_tables = ["addresses", "company_addresses", "component_files", "components", "connections",
"manufacturer_organizations", "innovator_organizations", "notifications", "orders", "organizations",
"pricing_information", "project_rfps", "projects", "rfps", "user_notifications", "user_projects", "users"]
export_tables = dim_tables + fact_tables
tables = []
for dim in export_tables:
cfg = TableConfig(cloud_sql_instance='CLOUD_SQL_INSTANCE_NAME',
export_table=dim.split(".")[-1],
export_bucket='YOUR_STAGING_BUCKET',
export_database=dim.split(".")[0],
export_query="SELECT * from {}".format(dim),
gcp_project="YOUR_PROJECT_ID",
stage_dataset="YOUR_STAGING_DATASET",
stage_table=None,
stage_final_query=None,
bq_location="EU")
tables.append(cfg)
return tables

def gen_export_table_task(table_config):
export_task = PostgresToGoogleCloudStorageOperator(task_id='export_{}'.format(table_config.params['export_table']),
dag=dag,
sql=table_config.params['export_query'],
bucket=table_config.params['export_bucket'],
filename="cloudsql_to_bigquery/{}/{}".format(table_config.params['export_database'],
table_config.params['export_table']) + "_{}",
schema_filename="cloudsql_to_bigquery/schema/{}/schema_raw".format(table_config.params['export_table']),
mysql_conn_id="gcp_dvh_cloudsql")
export_task.doc_md = """
#### Export table from cloudsql to cloud storage
task documentation
"""
return export_task


def gen_import_table_task(table_config):
import_task = GoogleCloudStorageToBigQueryOperator(
task_id='{}_to_bigquery'.format(table_config.params['export_table']),
bucket=table_config.params['export_bucket'],
source_objects=["cloudsql_to_bigquery/{}/{}*".format(table_config.params['export_database'],
table_config.params['export_table'])],
destination_project_dataset_table="{}.{}.{}".format(table_config.params['gcp_project'],
table_config.params['stage_dataset'],
table_config.params['stage_table']),
schema_object="cloudsql_to_bigquery/schema/{}/schema_raw".format(table_config.params['export_table']),
write_disposition='WRITE_TRUNCATE',
source_format="NEWLINE_DELIMITED_JSON",
dag=dag)

import_task.doc_md = """
#### Import table from storage to bigquery
task documentation
"""
return import_task


"""
The code that follows setups the dependencies between the tasks
"""

for table_config in get_tables():
export_script = gen_export_table_task(table_config)
import_script = gen_import_table_task(table_config)

export_script >> import_script
129,1 Bot