def test_should_pass_argument_to_hook(self, mock_hook):
task = GoogleCloudStorageObjectSensor(
task_id="task-id",
bucket=TEST_BUCKET,
object=TEST_OBJECT,
google_cloud_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
mock_hook.return_value.exists.return_value = True
result = task.poke(mock.MagicMock())
self.assertEqual(True, result)
mock_hook.assert_called_once_with(
delegate_to=TEST_DELEGATE_TO,
google_cloud_storage_conn_id=TEST_GCP_CONN_ID
)
mock_hook.return_value.exists.assert_called_once_with(TEST_BUCKET, TEST_OBJECT)
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 = Client()
job_config = LoadJobConfig()
schema_path = os.path.join(
dags_folder,
'resources/stages/raw/schemas/{task}.json'.format(task=task))
job_config.schema = read_bigquery_schema_from_file(schema_path)
job_config.source_format = SourceFormat.CSV if file_format == 'csv' else 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
# Load from
export_location_uri = 'gs://{bucket}/export'.format(
bucket=output_bucket)
date_glob = '*'
uri = '{export_location_uri}/{task}/{date_glob}.{file_format}'.format(
export_location_uri=export_location_uri,
task=task,
date_glob=date_glob,
file_format=file_format)
logging.info('Load from uri: ' + uri)
# Table name
table_name = task
logging.info('Table name: ' + table_name)
table_ref = client.dataset(dataset_name_raw).table(table_name)
# Load job
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_load_tasks(task, file_format, allow_quoted_newlines=False):
output_bucket = os.environ.get('OUTPUT_BUCKET')
if output_bucket is None:
raise ValueError('You must set OUTPUT_BUCKET environment variable')
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():
client = bigquery.Client()
job_config = bigquery.LoadJobConfig()
schema_path = os.path.join(
dags_folder,
'resources/stages/raw/schemas/{task}.json'.format(task=task))
job_config.schema = read_bigquery_schema_from_file(schema_path)
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
export_location_uri = 'gs://{bucket}/export'.format(
bucket=output_bucket)
uri = '{export_location_uri}/{task}/*.{file_format}'.format(
export_location_uri=export_location_uri,
task=task,
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,
execution_timeout=timedelta(minutes=30),
dag=dag)
wait_sensor >> load_operator
return load_operator
def add_load_tasks(task, time_partitioning_field='timestamp'):
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}.json'.format(
task=task, datestamp='{{ds}}'),
dag=dag)
def load_task():
client = bigquery.Client()
job_config = bigquery.LoadJobConfig()
schema_path = os.path.join(
dags_folder,
'resources/stages/load/schemas/{task}.json'.format(task=task))
job_config.schema = read_bigquery_schema_from_file(schema_path)
job_config.source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
job_config.write_disposition = 'WRITE_TRUNCATE'
job_config.ignore_unknown_values = True
job_config.time_partitioning = TimePartitioning(
field=time_partitioning_field)
export_location_uri = 'gs://{bucket}/export'.format(
bucket=output_bucket)
uri = '{export_location_uri}/{task}/*.json'.format(
export_location_uri=export_location_uri, task=task)
table_ref = create_dataset(
client, dataset_name,
destination_dataset_project_id).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,
execution_timeout=timedelta(minutes=30),
dag=dag)
wait_sensor >> load_operator
return load_operator
'retries': 1,
'retry_delay': timedelta(minutes=5)
}
# write dag
with DAG(dag_id='blog',
default_args=default_args,
schedule_interval=schedule_interval,
catchup=False) as dag:
t1 = BigQueryCheckOperator(task_id='check_bq_data_exists',
sql=queries.check_bq_data_exists,
use_legacy_sql=False)
t2 = GoogleCloudStorageObjectSensor(task_id='check_gcs_file_exists',
bucket=cfg.BUCKET,
object=cfg.SOURCE_OBJECT)
t3 = BigQueryOperator(task_id='write_weight_data_to_bq',
sql=queries.write_weight_data_to_bq,
destination_dataset_table=cfg.BQ_TABLE_WEIGHT,
create_disposition='CREATE_IF_NEEDED',
write_disposition='WRITE_TRUNCATE',
use_legacy_sql=False)
t4 = GoogleCloudStorageToBigQueryOperator(
task_id='write_route_data_to_bq',
bucket=cfg.BUCKET,
source_objects=[cfg.SOURCE_OBJECT],
field_delimiter=';',
destination_project_dataset_table=cfg.BQ_TABLE_ROUTE,
'retries': 1,
'retry_delay': datetime.timedelta(minutes=3),
'project_id': models.Variable.get('gcp_project')
}
with models.DAG(
'bdpc_lab8_orchestration',
schedule_interval='@hourly',
default_args=default_dag_args,
catchup=False) as dag:
# Sensor to check file in GCS -> gc://${bucket_name}/flights/${yyyy}/${MM}/${dd}/${HH}/_SUCCESS .
gcs_file_sensor = GoogleCloudStorageObjectSensor(
task_id='check_gcs_file_sensor',
timeout=120,
bucket='barinov_bdpc',
soft_fail=True, # This will skip the DAG if file not found, but alternatively this could be set to False to fail the DAG.
object='flights/{{ execution_date.format("%Y/%m/%d/%H") }}/_SUCCESS')
# Create a Cloud Dataproc cluster.
create_dataproc_cluster = dataproc_operator.DataprocClusterCreateOperator(
task_id='create_dataproc_cluster',
cluster_name='lab8-work-cluster-{{ ds_nodash }}',
num_workers=2,
zone=models.Variable.get('gce_zone'),
master_machine_type='n1-standard-1',
worker_machine_type='n1-standard-1')
# Run Spark job - Popular airports
run_dataproc_spark_t1 = dataproc_operator.DataProcSparkOperator(
task_id='run_dataproc_spark_task1',
'owner': 'Airflow',
#'start_date': airflow.utils.dates.days_ago(2),
'start_date': datetime.datetime(2020, 1, 27),
}
dag = DAG(
dag_id='exercise_gcp_dag',
default_args=args,
schedule_interval="0 0 * * *",
dagrun_timeout=datetime.timedelta(minutes=60),
)
sensor = GoogleCloudStorageObjectSensor(
task_id='sensor',
bucket='test_bucket312312',
timeout=3600,
object='test/{{ ds_nodash }}.csv',
dag=dag,
)
copy_to_bq = GoogleCloudStorageToBigQueryOperator(
task_id='copy_to_bq',
dag=dag,
bucket='test_bucket312312',
source_objects=['test/{{ ds_nodash }}.csv'],
destination_project_dataset_table=
'airflowbolcom-jan2829-2ad52563.test_dataset.test_table',
skip_leading_rows=1,
write_disposition='WRITE_TRUNCATE',
autodetect=True,
)
task_trigger_report_email = EmailOperator(
task_id='send_report_email_' + rpt_folder_name,
to=rpt_rcpnt,
cc=rpt_pic,
subject='Report ' + rpt_folder_name + "_" + "{{ ds }}",
html_content=EmailTemplate,
dag=dag)
rpt_cnt = 0
for itr_rpt in report_config[rpt_folder_cnt]['reports']:
rpt_name = report_config[rpt_folder_cnt]['reports'][rpt_cnt][
'report_name']
task_gcs_sensor = GoogleCloudStorageObjectSensor(
task_id='gcs_sensor_' + rpt_name,
bucket=rpt_bucket_name,
object='%s/%s_%s.txt' % (rpt_folder_name, "{{ ds }}", rpt_name),
google_cloud_conn_id='gcp_project_deng',
dag=dag)
task_gcs_give_acl = GoogleCloudStorageObjectCreateAclEntryOperator(
task_id='giveaclaccess_gcs_' + rpt_name,
object_name='%s/%s_%s.txt' %
(rpt_folder_name, "{{ ds }}", rpt_name),
entity='allUsers',
role='READER',
bucket='rpt_bucket_name',
google_cloud_storage_conn_id='gcp_project_deng',
dag=dag)
task_start >> task_start_report >> task_gcs_sensor >> task_gcs_give_acl
rpt_cnt += 1
def build_sessions_dag(dag_id,
output_bucket,
sql_dir,
source_project_id,
source_dataset_name,
destination_project_id,
destination_dataset_name,
temp_dataset_name,
notification_emails=None,
schedule_interval='0 14 * * *',
start_date=datetime(2015, 7, 30),
environment='prod'):
default_dag_args = {
'depends_on_past': False,
'start_date': 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=True,
schedule_interval=schedule_interval,
max_active_runs=1,
default_args=default_dag_args)
def read_file(filepath):
with open(filepath) as file_handle:
content = file_handle.read()
return content
def add_sessions_task(task, dependencies=None):
def sessions_task(ds, **kwargs):
client = bigquery.Client()
sql_path = os.path.join(sql_dir, '{task}.sql'.format(task=task))
sql_template = read_file(sql_path)
ds_no_dashes = ds.replace('-', '')
sql = sql_template.format(
ds=ds,
ds_no_dashes=ds_no_dashes,
source_project_id=source_project_id,
source_dataset_name=source_dataset_name,
destination_project_id=destination_project_id,
destination_dataset_name=destination_dataset_name,
temp_dataset_name=temp_dataset_name,
)
print(sql)
query_job = client.query(sql)
result = query_job.result()
logging.info(result)
sessions_operator = PythonOperator(
task_id=f'{task}',
# Necessary because we use traces overlapping the previous execution date.
depends_on_past=True,
wait_for_downstream=True,
python_callable=sessions_task,
provide_context=True,
execution_timeout=timedelta(minutes=60),
dag=dag)
if dependencies is not None and len(dependencies) > 0:
for dependency in dependencies:
dependency >> sessions_operator
return sessions_operator
stage_root_call_traces_task = add_sessions_task('root_call_traces')
upsert_sessions_task = add_sessions_task('sessions')
# Dummy task indicating successful DAG completion.
done_task = BashOperator(task_id='done', bash_command='echo done', dag=dag)
#
# Task sensor is enabled only in production for now because our load DAG is
# not running in the lower environments.
#
if environment == 'prod':
wait_for_ethereum_load_dag_task = GoogleCloudStorageObjectSensor(
task_id='wait_for_ethereum_load_dag',
timeout=60 * 60 * 12,
poke_interval=5 * 60,
bucket=output_bucket,
object=
"checkpoint/block_date={block_date}/load_complete_checkpoint.txt".
format(block_date='{{ds}}'),
dag=dag)
wait_for_ethereum_load_dag_task >> stage_root_call_traces_task
stage_root_call_traces_task >> upsert_sessions_task
upsert_sessions_task >> done_task
return dag
region='us-west1',
job_name=cleaned_dag_id + 'calc_retention_day1',
cluster_name=
'{{ ti.xcom_pull(key=unique_cluster_name, task_ids="generate_unique_cluster_name") }}'
+ '2',
execution_timeout=timedelta(hours=2),
arguments=args,
dag=dag)
# ==================
# == gcs sensors ===
# ==================
sensor_task = GoogleCloudStorageObjectSensor(task_id='sensor_task',
bucket='exampleBucket',
object='output/user/_SUCCESS',
poke_interval=30,
timeout=2700,
dag=dag)
# =======================
# == load to bigquery ===
# =======================
bq_load_user = GoogleCloudStorageToBigQueryOperator(
task_id='bq_load_user',
bucket='exampleBucket',
source_objects=["obj_folder/*"],
source_format='NEWLINE_DELIMITED_JSON',
destination_project_dataset_table='exampleSchema.exampleTable0',
schema_fields=[
{
'project_id': var_project_id
}
# [START composer_hadoop_schedule]
with models.DAG(
'airflow',
# Continue to run DAG once per day
schedule_interval=datetime.timedelta(hours=1),
default_args=default_dag_args,
catchup=False) as dag:
# [END composer_hadoop_schedule]
# Sensor
gcs_file_sensor = GoogleCloudStorageObjectSensor(
task_id='waiting_file_sensor',
timeout=120,
bucket=var_bucket_name,
soft_fail=True,
object='flights/{{ execution_date.format("%Y/%m/%d/%H") }}/_SUCCESS')
# Create a Cloud Dataproc cluster.
create_dataproc_cluster = dataproc_operator.DataprocClusterCreateOperator(
task_id='create_dataproc_cluster',
# Give the cluster a unique name by appending the date scheduled.
# See https://airflow.apache.org/code.html#default-variables
cluster_name=naming_cluster + '-{{ ds_nodash }}',
num_workers=2,
zone=var_zone,
master_machine_type='n1-standard-1',
worker_machine_type='n1-standard-1')
# Run the Hadoop wordcount example installed on the Cloud Dataproc cluster
'retries': 1,
'retry_delay': datetime.timedelta(minutes=5),
'project_id': models.Variable.get('gcp_project')
}
with models.DAG(
'composer_spark_job',
# Continue to run DAG once per day
schedule_interval='0 *,* * * *',
default_args=default_dag_args,
catchup=False) as dag:
gcs_file_sensor = GoogleCloudStorageObjectSensor(
task_id='gcs_file_sensor_task',
bucket=bucket,
object=full_path,
timeout=3600)
# Create a Cloud Dataproc cluster.
create_dataproc_cluster = dataproc_operator.DataprocClusterCreateOperator(
task_id='create_dataproc_cluster',
# Give the cluster a unique name by appending the date scheduled.
# See https://airflow.apache.org/code.html#default-variables
cluster_name='composer-hadoop-tutorial-cluster-{{ ds_nodash }}',
num_workers=2,
zone=models.Variable.get('gce_zone'),
master_machine_type='n1-standard-1',
worker_machine_type='n1-standard-1')
DEFAULT_ARGS = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': datetime(2019, 8, 18),
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=5),
}
DAG_ID = DAG(dag_id=DAG_NAME,
default_args=DEFAULT_ARGS,
schedule_interval='@daily')
SENSE = GoogleCloudStorageObjectSensor(
task_id='get_file_sample',
dag=DAG_ID,
bucket=BUCKET,
object=OBJECT,
google_cloud_conn_id=GCS_CONN,
)
PRINT = BashOperator(task_id='print',
bash_command='echo "file arrived"',
params={'my_param': 'dummy param'},
dag=DAG_ID)
SENSE > PRINT
def build_load_worldpop_dag(
dag_id,
output_bucket,
countries,
large_countries,
destination_dataset_project_id,
destination_dataset_name,
destination_table_name,
staging_dataset_project_id,
staging_dataset_name,
dataflow_template_path,
dataflow_environment,
notification_emails=None,
load_start_date=datetime(2000, 1, 1),
load_schedule_interval="0 0 * * *",
load_max_active_runs=None,
load_concurrency=None,
load_retries=5,
load_retry_delay=300,
output_path_prefix="export",
**kwargs,
):
if not output_bucket:
raise ValueError("output_bucket is required")
if not destination_dataset_project_id:
raise ValueError("destination_dataset_project_id is required")
if not destination_dataset_name:
raise ValueError("destination_dataset_name is required")
if not destination_table_name:
raise ValueError("destination_table_name is required")
default_dag_args = {
"depends_on_past": False,
"start_date": load_start_date,
"end_date": None,
"email_on_failure": True,
"email_on_retry": False,
"retries": load_retries,
"retry_delay": timedelta(seconds=load_retry_delay),
}
if notification_emails and len(notification_emails) > 0:
default_dag_args["email"] = [
email.strip() for email in notification_emails.split(",")
]
if load_max_active_runs is None:
load_max_active_runs = configuration.conf.getint(
"core", "max_active_runs_per_dag")
dag = models.DAG(
dag_id,
schedule_interval=load_schedule_interval,
max_active_runs=load_max_active_runs,
concurrency=load_concurrency,
default_args=default_dag_args,
is_paused_upon_creation=True,
)
dags_folder = os.environ.get("DAGS_FOLDER", "/home/airflow/gcs/dags")
def read_bigquery_schema(schema):
schema_path = os.path.join(
dags_folder,
"resources/stages/load/schemas/{schema}.json".format(
schema=schema),
)
return read_bigquery_schema_from_file(schema_path)
def load_task(country, **context):
client = bigquery.Client()
job_config = bigquery.LoadJobConfig()
job_config.schema = read_bigquery_schema("world_pop")
job_config.source_format = bigquery.SourceFormat.PARQUET
job_config.write_disposition = "WRITE_TRUNCATE"
job_config.ignore_unknown_values = True
job_config.range_partitioning = RangePartitioning(
field="year",
range_=PartitionRange(start=1900, end=2100, interval=1),
)
execution_date = context["execution_date"]
load_table_name = "{table}_{country}_{year}".format(
table=destination_table_name,
country=country,
year=execution_date.strftime("%Y"),
)
table_ref = create_dataset(
client,
staging_dataset_name,
project=staging_dataset_project_id,
).table(load_table_name)
load_uri = "gs://{bucket}/{prefix}/world_pop/year={year}/parquet/{country}_{year}.parquet".format(
bucket=output_bucket,
prefix=output_path_prefix,
country=country,
year=execution_date.strftime("%Y"),
)
load_job = client.load_table_from_uri(
load_uri,
table_ref,
job_config=job_config,
)
submit_bigquery_job(load_job, job_config)
assert load_job.state == "DONE"
def merge_task(country, **context):
client = bigquery.Client()
table_ref = create_dataset(
client,
destination_dataset_name,
project=destination_dataset_project_id,
).table(destination_table_name)
if not does_table_exist(client, table_ref):
table = bigquery.Table(table_ref,
schema=read_bigquery_schema("world_pop"))
table.range_partitioning = RangePartitioning(
field="year",
range_=PartitionRange(start=1900, end=2100, interval=1),
)
table.clustering_fields = [
"geography",
"geography_polygon",
"country",
]
client.create_table(table)
job_config = bigquery.QueryJobConfig()
job_config.priority = bigquery.QueryPriority.INTERACTIVE
sql_path = os.path.join(
dags_folder, "resources/stages/load/sqls/merge_worldpop.sql")
sql_template = read_file(sql_path)
execution_date = context["execution_date"]
year = execution_date.strftime("%Y")
staging_table_name = "{table}_{country}_{year}".format(
table=destination_table_name, country=country, year=year)
template_context = {
"year": year,
"country": country,
"source_table": staging_table_name,
"source_project_id": staging_dataset_project_id,
"source_dataset_name": staging_dataset_name,
"destination_table": destination_table_name,
"destination_dataset_project_id": destination_dataset_project_id,
"destination_dataset_name": destination_dataset_name,
}
sql = context["task"].render_template(sql_template, template_context)
job = client.query(sql, location="US", job_config=job_config)
submit_bigquery_job(job, job_config)
assert job.state == "DONE"
priority = len(countries)
for country in countries.split(","):
c = country.lower()
wait_uri = (
"{prefix}/world_pop/year={year}/parquet/{country}_{year}.parquet".
format(
prefix=output_path_prefix,
country=country,
year='{{execution_date.strftime("%Y")}}',
))
wait_gcs = GoogleCloudStorageObjectSensor(
task_id=f"wait_{c}",
timeout=60 * 60,
poke_interval=60,
bucket=output_bucket,
object=wait_uri,
weight_rule="upstream",
priority_weight=priority,
dag=dag,
)
if country in large_countries:
input_file = "gs://{bucket}/{prefix}/world_pop/year={year}/parquet/{country}_{year}.parquet".format(
bucket=output_bucket,
prefix=output_path_prefix,
country=country,
year='{{ execution_date.strftime("%Y") }}',
)
output_table = "{table}_{country}_{year}".format(
table=destination_table_name,
country=country,
year='{{ execution_date.strftime("%Y") }}',
)
load_operator = DataflowStartFlexTemplateOperator(
task_id=f"run_dataflow_load_{c}",
body={
"launchParameter": {
"containerSpecGcsPath":
f"{dataflow_template_path}/load-parquet-0.1.0.json",
"jobName": "load-parquet-{country}".format(country=c) +
'-{{ execution_date.strftime("%Y%m%d-%H%M%S") }}',
"parameters": {
"input-file":
input_file,
"output-table":
f"{staging_dataset_project_id}:{staging_dataset_name}.{output_table}",
"output-schema":
"/dataflow/template/schemas/world_pop.json",
},
"environment": dataflow_environment,
}
},
location="us-central1",
wait_until_finished=True,
dag=dag,
)
else:
load_operator = PythonOperator(
task_id=f"load_{c}",
python_callable=load_task,
execution_timeout=timedelta(minutes=600),
provide_context=True,
op_kwargs={"country": country},
retries=1,
retry_delay=timedelta(seconds=300),
weight_rule="upstream",
priority_weight=priority,
dag=dag,
)
merge_operator = PythonOperator(
task_id=f"merge_{c}",
python_callable=merge_task,
execution_timeout=timedelta(minutes=600),
provide_context=True,
op_kwargs={"country": country},
retries=1,
retry_delay=timedelta(seconds=300),
weight_rule="upstream",
priority_weight=priority,
dag=dag,
)
priority -= 1
wait_gcs >> load_operator >> merge_operator
return dag
bike_share_retention_d7 = DataProcSparkOperator(
task_id='bike_share_retention_d7',
dataproc_spark_jars=[
'gs://jiuzhangsuanfa/jar/CohortProject-assembly-0.1.jar'
],
main_class='com.cohort.process.RetentionProcess',
region='us-west1',
job_name=dag_name + 'bike_share_retention_d7',
cluster_name='{{ ti.xcom_pull(key="cluster_name", task_ids="push-cluster-name") }}'
+ '4',
execution_timeout=timedelta(minutes=180),
arguments=args)
unique_user_sensor = GoogleCloudStorageObjectSensor(
task_id='unique_user_sensor',
bucket='jiuzhangsuanfa',
object='bike/unique-user/_SUCCESS',
poke_interval=30,
timeout=2700)
unique_user.set_upstream(dataproc_create_cluster_1)
unique_user.set_downstream(bike_share_aggregator)
bike_share_aggregator.set_downstream(dataproc_destroy_cluster_1)
bike_share_retention_d1.set_upstream(dataproc_create_cluster_2)
bike_share_retention_d1.set_downstream(dataproc_destroy_cluster_2)
bike_share_retention_d3.set_upstream(dataproc_create_cluster_3)
# print("initest : " + str(context.get("execution_date")))
# print("initest : " + type(context.get("execution_date")))
# print("initest : " + type(context.get("dag").schedule_interval))
test = pendulum.datetime(2019, 2, 16, 10, 0, 0)
print("initest Execution_date: " + str(context['execution_date']) +
" Schedule_interval: " + str(context['dag'].schedule_interval))
return context.get(
"execution_date") + test #context['dag'].schedule_interval
# def execute(self, context):
# execution_date = context.get("execution_date")
t1 = GoogleCloudStorageObjectSensor(
task_id='gcs_sensor_trigger',
bucket='dataflow-results-ini',
object='gs://dataflow-results-ini/airflow-tesing/test.csv',
google_cloud_conn_id='gcs_conn',
dag=dag)
t2 = GoogleCloudStorageToBigQueryOperator(
task_id='gcs_to_bq',
bucket='dataflow-results-ini',
source_objects=['airflow-testing/test.csv'],
schema_object='airflow-testing/test.json',
destination_project_dataset_table="dataflowtesting-218212.testing.newone",
write_disposition='WRITE_APPEND',
autodetect=True,
bigquery_conn_id=BQ_CONN_ID,
dag=dag)
t1 >> t2
# --------------------------------------------------------------------------------
dag = DAG('DAG_3_GCS_To_BigQuery',
default_args=default_args,
schedule_interval='@daily')
start = DummyOperator(
task_id='start',
trigger_rule='all_success',
dag=dag
)
GCS_landing_sensor = GoogleCloudStorageObjectSensor(
task_id='GCS_landing_sensor',
bucket=gcs_bucket,
object=os.path.join(gcs_dir, success_file),
google_cloud_conn_id=google_cloud_conn_id,
dag=dag
)
GCS_to_BigQuery = GoogleCloudStorageToBigQueryOperator(
task_id='GCS_to_BigQuery',
destination_project_dataset_table=bq_destination_dataset_table,
bucket=gcs_bucket,
source_format="avro",
source_objects=[os.path.join(gcs_dir, '*.avro')],
create_disposition='CREATE_IF_NEEDED',
# The following values are supported for `create_disposition`:
# CREATE_IF_NEEDED: If the table does not exist, BigQuery creates the table.
# CREATE_NEVER: The table must already exist. If it does not, a 'notFound' error is returned in the job result.
# The default value is CREATE_IF_NEEDED.
dataproc_spark_jars=['gs://path/jar/CohortAnalysis.jar'],
main_class='com.makoto.spark.process.RetentionComputeProcess',
region='us-west1',
job_name=dag_name + 'bike_share_retention_d7',
cluster_name=
'{{ ti.xcom_pull(key="cluster_name", task_ids="push-cluster-name") }}'
+ '4',
execution_timeout=timedelta(minutes=180),
arguments=args)
# check for every 1 min = poke_interval if _SUCCESS, the status of computed unique user list json exist
# timeout be 1 hour
# if completed user list generation, trigger down stream jobs
unique_user_sensor = GoogleCloudStorageObjectSensor(
task_id='unique_user_sensor',
bucket='bucketname',
object='bike/unique-user/_SUCCESS',
poke_interval=30,
timeout=6000)
unique_user.set_upstream(dataproc_create_cluster_1)
unique_user.set_downstream(bike_share_aggregator)
bike_share_aggregator.set_downstream(dataproc_destroy_cluster_1)
bike_share_retention_d1.set_upstream(dataproc_create_cluster_2)
bike_share_retention_d1.set_downstream(dataproc_destroy_cluster_2)
bike_share_retention_d3.set_upstream(dataproc_create_cluster_3)
def build_load_npp_dag(dag_id,
output_bucket,
destination_dataset_project_id,
destination_dataset_name,
destination_table_name,
notification_emails=None,
load_start_date=datetime(2000, 1, 1),
load_schedule_interval=None,
load_max_active_runs=None,
load_concurrency=None,
load_retries=5,
load_retry_delay=300,
output_path_prefix="export",
**kwargs):
if not output_bucket:
raise ValueError("output_bucket is required")
if not destination_dataset_project_id:
raise ValueError("destination_dataset_project_id is required")
if not destination_dataset_name:
raise ValueError("destination_dataset_name is required")
if not destination_table_name:
raise ValueError("destination_table_name is required")
default_dag_args = {
"depends_on_past": False,
"start_date": load_start_date,
"end_date": None,
"email_on_failure": True,
"email_on_retry": False,
"retries": load_retries,
"retry_delay": timedelta(seconds=load_retry_delay),
}
if notification_emails and len(notification_emails) > 0:
default_dag_args["email"] = [
email.strip() for email in notification_emails.split(",")
]
if load_max_active_runs is None:
load_max_active_runs = configuration.conf.getint(
"core", "max_active_runs_per_dag")
dag = models.DAG(
dag_id,
schedule_interval=load_schedule_interval,
max_active_runs=load_max_active_runs,
concurrency=load_concurrency,
default_args=default_dag_args,
is_paused_upon_creation=True,
)
def load_task(**context):
dags_folder = os.environ.get("DAGS_FOLDER", "/home/airflow/gcs/dags")
schema_path = os.path.join(
dags_folder,
"resources/stages/load/schemas/{schema}.json".format(
schema="annual_npp"),
)
client = bigquery.Client()
job_config = bigquery.LoadJobConfig()
job_config.schema = read_bigquery_schema_from_file(schema_path)
job_config.source_format = bigquery.SourceFormat.PARQUET
job_config.write_disposition = "WRITE_TRUNCATE"
job_config.ignore_unknown_values = True
job_config.clustering_fields = [
"geography",
"geography_polygon",
]
job_config.range_partitioning = RangePartitioning(
field="year",
range_=PartitionRange(start=1900, end=2100, interval=1),
)
execution_date = context["execution_date"]
load_table_name = "{table}${partition}".format(
table=destination_table_name,
partition=execution_date.strftime("%Y"))
table_ref = create_dataset(
client,
destination_dataset_name,
project=destination_dataset_project_id,
).table(load_table_name)
load_uri = "gs://{bucket}/{prefix}/annual_npp/parquet/{date}.parquet".format(
bucket=output_bucket,
prefix=output_path_prefix,
date=execution_date.strftime("%Y_%m_%d"),
)
load_job = client.load_table_from_uri(
load_uri,
table_ref,
job_config=job_config,
)
submit_bigquery_job(load_job, job_config)
assert load_job.state == "DONE"
wait_uri = "{prefix}/annual_npp/parquet/{date}.parquet".format(
prefix=output_path_prefix,
date='{{ execution_date.strftime("%Y_%m_%d") }}')
wait_gcs = GoogleCloudStorageObjectSensor(
task_id="wait_gcs",
timeout=60 * 60,
poke_interval=60,
bucket=output_bucket,
object=wait_uri,
dag=dag,
)
load_operator = PythonOperator(
task_id="load_to_bigquery",
python_callable=load_task,
execution_timeout=timedelta(minutes=600),
provide_context=True,
retries=1,
retry_delay=timedelta(seconds=300),
dag=dag,
)
wait_gcs >> load_operator
return dag
file_suffix = "_file.csv"
file_date = today.strftime('%Y-%m-%d')
full_path_today = file_prefix + file_date + file_suffix
file_date_yesterday = yesterday.strftime('%Y-%m-%d')
full_path_yesterday = file_prefix + file_date_yesterday + file_suffix
with airflow.DAG("file_sensor_example",
default_args=default_args,
schedule_interval="@once") as dag:
start_task = DummyOperator(task_id="start")
stop_task = DummyOperator(task_id="stop")
sensor_task = FileSensor(task_id="file_sensor_task",
poke_interval=30,
filepath="/tmp/")
#we expect yesterday to exist
gcs_file_sensor_yesterday = GoogleCloudStorageObjectSensor(
task_id='gcs_file_sensor_yesterday_task',
bucket='myBucketName',
object=full_path_yesterday)
#for this example we expect today not to exist, keep running until 120 timeout, checkout docs for more options like mode and soft_fail
gcs_file_sensor_today = GoogleCloudStorageObjectSensor(
task_id='gcs_file_sensor_today_task',
bucket='myBucketName',
object=full_path_today,
timeout=120)
start_task >> sensor_task >> gcs_file_sensor_yesterday >> gcs_file_sensor_today >> stop_task
'start_date': datetime(2019, 6, 17),
'email': ['you@your_org.org'],
'email_on_failure': True,
'email_on_retry': True,
'provide_context': True,
'owner': 'airflow',
'depends_on_past': True
}
dag = DAG(
dag_id='gcs_sensor',
default_args=args,
schedule_interval=None,
)
gcs_sensor = GoogleCloudStorageObjectSensor(
task_id="sense",
bucket='bucket_of_stuff',
object='sense_for_me.txt',
google_cloud_storage_conn_id='google_cloud_storage_default',
soft_fail=True,
poke_interval=5,
timeout=15,
mode="poke",
dag=dag)
start = DummyOperator(task_id="start", dag=dag)
end = DummyOperator(task_id="end", dag=dag)
start >> gcs_sensor >> end
seven_days_ago = datetime.combine(datetime.today() - timedelta(7),
datetime.min.time())
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': seven_days_ago,
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=30),
}
with DAG('v1_8_cloud_storage', schedule_interval=timedelta(days=1),
default_args=default_args) as dag:
sens_object_create = GoogleCloudStorageObjectSensor(
task_id='sens_object_create',
bucket='{{var.value.gcs_bucket}}',
object='{{var.value.gcs_root}}/gcp_smoke_bq/bq_to_gcp_avro/{{ ds_nodash }}/part-000000000000.avro',
google_cloud_conn_id='gcp_smoke'
)
sens_object_update = GoogleCloudStorageObjectUpdatedSensor(
task_id='sens_object_update',
bucket='{{var.value.gcs_bucket}}',
object='{{var.value.gcs_root}}/gcp_smoke_bq/bq_to_gcp_avro/99999999/part-000000000000.avro',
google_cloud_conn_id='gcp_smoke'
)
