def dag_preprocess_tables( dag_id, schedule_interval, start_date, target_project_id, target_dataset_id, table_config, table_partition, ): dag = DAG(dag_id=dag_id, schedule_interval=schedule_interval, start_date=start_date) for table in table_config: start_check_tables_task = DummyOperator( task_id='%s-%s' % ("start_check_tables_task", table["name"]), dag=dag) check_if_table_exist = BranchPythonOperator( task_id='%s-%s' % (table["name"], "check_if_table_exist"), python_callable=if_tbl_exists, op_kwargs={ 'dataset': target_dataset_id, 'project': target_project_id, 'table_name': table["name"] }, dag=dag) table_exists = DummyOperator(task_id='%s-%s' % (table["name"], "table_exists"), dag=dag) table_does_not_exist = DummyOperator( task_id='%s-%s' % (table["name"], "table_does_not_exist"), dag=dag) # [start create equipped_item_reference if not exists] create_if_not_exists = BigQueryCreateEmptyTableOperator( task_id='%s-%s' % (table["name"], "create_if_not_exists"), project_id=target_project_id, dataset_id=target_dataset_id, table_id=table["name"], gcs_schema_object=table["schema_gcs_location"], time_partitioning=table_partition, trigger_rule=TriggerRule.ALL_SUCCESS, dag=dag) end_check_tables_task = DummyOperator( task_id='%s-%s' % ("end_check_tables_task", table["name"]), trigger_rule='none_failed_or_skipped', dag=dag) start_check_tables_task >> check_if_table_exist >> [ table_does_not_exist, table_exists ] table_does_not_exist >> create_if_not_exists >> end_check_tables_task table_exists >> end_check_tables_task return dag
def test_execute(self, mock_hook): operator = BigQueryCreateEmptyTableOperator(task_id=TASK_ID, dataset_id=TEST_DATASET, project_id=TEST_PROJECT_ID, table_id=TEST_TABLE_ID) operator.execute(None) mock_hook.return_value \ .get_conn() \ .cursor() \ .create_empty_table \ .assert_called_once_with( dataset_id=TEST_DATASET, project_id=TEST_PROJECT_ID, table_id=TEST_TABLE_ID, schema_fields=None, time_partitioning={} )
def test_execute(self, mock_hook): operator = BigQueryCreateEmptyTableOperator(task_id=TASK_ID, dataset_id=TEST_DATASET, project_id=TEST_PROJECT_ID, table_id=TEST_TABLE_ID) operator.execute(None) mock_hook.return_value \ .get_conn() \ .cursor() \ .create_empty_table \ .assert_called_once_with( dataset_id=TEST_DATASET, project_id=TEST_PROJECT_ID, table_id=TEST_TABLE_ID, schema_fields=None, time_partitioning={} )
def __init__(self, project_id, dataset_id, table_id, schema_fields, *args, **kwargs): BigQueryDropTableOperator.__init__(self, project_id=project_id, dataset_id=dataset_id, table_id=table_id, *args, **kwargs) BigQueryCreateEmptyTableOperator.__init__( self, task_id='drop_create_{}'.format(table_id), project_id=project_id, dataset_id=dataset_id, table_id=table_id, schema_fields=schema_fields, *args, **kwargs)
} dag = DAG( dag_id='create_and_insert_table', default_args=default_args, description='DAG to create empty table, dataset provided', schedule_interval=timedelta(days=1), catchup=False, ) create_table_task = BigQueryCreateEmptyTableOperator( project_id="dark-furnace-298806", dataset_id="shtest", task_id="create_empty_table", table_id="new_service_table", bigquery_conn_id="bigquery_default", google_cloud_storage_conn_id="google_cloud_default", dag=dag ) insert_task = BigQueryOperator( task_id="insert_into_table", bql=''' select * from bigquery-public-data.austin_311.311_service_requests limit 5 ''', use_legacy_sql=False, destination_dataset_table="dark-furnace-298806:shtest.new_service_table",
task_id="create-dataset", dataset_id=DATASET_NAME) create_dataset_with_location = BigQueryCreateEmptyDatasetOperator( task_id="create_dataset_with_location", dataset_id=LOCATION_DATASET_NAME, location=BQ_LOCATION) create_table = BigQueryCreateEmptyTableOperator( task_id="create-table", dataset_id=DATASET_NAME, table_id="test_table", schema_fields=[ { "name": "emp_name", "type": "STRING", "mode": "REQUIRED" }, { "name": "salary", "type": "INTEGER", "mode": "NULLABLE" }, ], ) create_table_with_location = BigQueryCreateEmptyTableOperator( task_id="create_table_with_location", dataset_id=LOCATION_DATASET_NAME, table_id="test_table", schema_fields=[ {
# Create table # CreateTable = BigQueryCreateEmptyTableOperator( # task_id='BigQueryCreateEmptyTableOperator_task', # dataset_id='ODS', # table_id='Employees', # project_id='internal-gcp-project', # gcs_schema_object='gs://schema-bucket/employee_schema.json', # bigquery_conn_id='airflow-service-account', # google_cloud_storage_conn_id='airflow-service-account' # ) CreateTable = BigQueryCreateEmptyTableOperator( task_id='BigQueryCreateEmptyTableOperator_task', dataset_id='Covid', table_id='GoogleTrend_test', project_id='covidproject-278521', gcs_schema_object='gs://testcovidlinh/googletrend_schema.json', # schema_fields=googletrend_schema, bigquery_conn_id=gcp_conn_id, google_cloud_storage_conn_id=gcp_conn_id) # Loading data from GCS to BigQuery gcs_to_bigquery = GoogleCloudStorageToBigQueryOperator( task_id='GCS_to_BigQuery', dag=dag, bucket='testcovidlinh', source_objects=source_objects, # schema_object = "/tmp/covidStatSchema.json", schema_object='googletrend_schema.json', # schema_fields = googletrend_schema, source_format='CSV',
# task2 create a new covid_19 partition table. task2 = BigQueryCreateEmptyTableOperator(task_id='bq_create_new_talble', dag=dag, dataset_id=BQ_DATASET, table_id=BQ_TABLE, project_id=BQ_PROJECT, bigquery_conn_id=BQ_CONN_ID, schema_fields=[{ "name": "DateStr", "type": "DATE", "mode": "REQUIRED" }, { "name": "State", "type": "STRING", "mode": "REQUIRED" }, { "name": "Count", "type": "INTEGER", "mode": "REQUIRED" }, { "name": "Status", "type": "STRING", "mode": "REQUIRED" }], time_partitioning={ "type": "DAY", "field": "DateStr", "expiration_ms": "5184000000" } # partition expiration 60 days )
dag = DAG('incremental_ingestion', catchup=False, default_args=default_args, schedule_interval=dt.timedelta(days=1)) load_config = PythonOperator(task_id='load_config', provide_context=True, python_callable=fetch_config, dag=dag) bq_create_staging = BigQueryCreateEmptyTableOperator( task_id='bq_create_staging', project_id= '{{(var.json|attr("config-{}".format(run_id)))["bigquery"]["project_id"]}}', dataset_id= '{{(var.json|attr("config-{}".format(run_id)))["bigquery"]["dataset_id"]}}', table_id= '{{(var.json|attr("config-{}".format(run_id)))["bigquery"]["staging_table_id"]}}', gcs_schema_object= '{{(var.json|attr("config-{}".format(run_id)))["bigquery"]["schema_file"]}}', dag=dag) stage_data = DataflowTemplateOperator( task_id='stage_data', template= '{{ (var.json|attr("config-{}".format(run_id)))["dataflow"]["template"] }}', dataflow_default_options={ 'project': '{{ (var.json|attr("config-{}".format(run_id)))["dataflow"]["options"]["project"] }}', 'region': '{{ (var.json|attr("config-{}".format(run_id)))["dataflow"]["options"]["region"] }}',
task_id='create_football_matches_dataset', project_id=project_id, dataset_id=dataset_id, bigquery_conn_id=gcp_conn, dag=dag ) create_games_table = BigQueryCreateEmptyTableOperator( task_id="create_games_table", project_id=project_id, dataset_id=dataset_id, bigquery_conn_id=gcp_conn, table_id="games", schema_fields=[{"name": "date", "type": "TIMESTAMP", "mode": "REQUIRED"}, {"name": "team_1", "type": "STRING", "mode": "REQUIRED"}, {"name": "team_2", "type": "STRING", "mode": "REQUIRED"}, {"name": "team_1_score", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "team_2_score", "type": "INTEGER", "mode": "REQUIRED"}, {"name": "home_team", "type": "STRING", "mode": "NULLABLE"}, {"name": "tournament", "type": "STRING", "mode": "NULLABLE"}, {"name": "country", "type": "STRING", "mode": "NULLABLE"}, {"name": "city", "type": "STRING", "mode": "NULLABLE"}, ], dag=dag ) etl_games_job_code_path = 'gs://' + gcs_football_bucket + '/spark_jobs/etl_games_to_bigquery.py' submit_etl_games_spark_job = DataProcPySparkOperator( task_id='submit_etl_games_spark_job', main=etl_games_job_code_path, cluster_name=cluster_name, job_name='etl_games_to_bigquery',
def execute(self, context): BigQueryDropTableOperator.execute(self, context) BigQueryCreateEmptyTableOperator.execute(self, context)