From 8985dc09bcbc5d419aab5ed19d3b406f4e70b8f0 Mon Sep 17 00:00:00 2001 From: Yu ISHIKAWA Date: Fri, 21 Jul 2017 20:48:59 -0700 Subject: [PATCH 1/2] [AIRFLOW-1389] BigQueryOperator should support `createDisposition` --- airflow/contrib/hooks/bigquery_hook.py | 5 +++++ airflow/contrib/operators/bigquery_operator.py | 12 ++++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py index dc98b896fdb97..a8fba8a4b4a6f 100644 --- a/airflow/contrib/hooks/bigquery_hook.py +++ b/airflow/contrib/hooks/bigquery_hook.py @@ -192,6 +192,7 @@ def __init__(self, service, project_id): def run_query( self, bql, destination_dataset_table = False, write_disposition = 'WRITE_EMPTY', + create_disposition='CREATE_IF_NEEDED', allow_large_results=False, udf_config = False, use_legacy_sql=True, @@ -210,6 +211,9 @@ def run_query( BigQuery table to save the query results. :param write_disposition: What to do if the table already exists in BigQuery. + :type write_disposition: string + :param create_disposition: Specifies whether the job is allowed to create new tables. + :type create_disposition: string :param allow_large_results: Whether to allow large results. :type allow_large_results: boolean :param udf_config: The User Defined Function configuration for the query. @@ -238,6 +242,7 @@ def run_query( configuration['query'].update({ 'allowLargeResults': allow_large_results, 'writeDisposition': write_disposition, + 'createDisposition': create_disposition, 'destinationTable': { 'projectId': destination_project, 'datasetId': destination_dataset, diff --git a/airflow/contrib/operators/bigquery_operator.py b/airflow/contrib/operators/bigquery_operator.py index 5e8a0d5eb56f9..ecb0b9c884f61 100644 --- a/airflow/contrib/operators/bigquery_operator.py +++ b/airflow/contrib/operators/bigquery_operator.py @@ -31,6 +31,12 @@ class BigQueryOperator(BaseOperator): (.|:). that, if set, will store the results of the query. :type destination_dataset_table: string + :param write_disposition: Specifies the action that occurs if the destination table + already exists. (default: 'WRITE_EMPTY') + :type write_disposition: string + :param create_disposition: Specifies whether the job is allowed to create new tables. + (default: 'CREATE_IF_NEEDED') + :type create_disposition: string :param bigquery_conn_id: reference to a specific BigQuery hook. :type bigquery_conn_id: string :param delegate_to: The account to impersonate, if any. @@ -55,6 +61,7 @@ def __init__(self, bql, destination_dataset_table=False, write_disposition='WRITE_EMPTY', + create_disposition='CREATE_IF_NEEDED', allow_large_results=False, bigquery_conn_id='bigquery_default', delegate_to=None, @@ -67,6 +74,7 @@ def __init__(self, self.bql = bql self.destination_dataset_table = destination_dataset_table self.write_disposition = write_disposition + self.create_disposition = create_disposition self.allow_large_results = allow_large_results self.bigquery_conn_id = bigquery_conn_id self.delegate_to = delegate_to @@ -81,5 +89,5 @@ def execute(self, context): conn = hook.get_conn() cursor = conn.cursor() cursor.run_query(self.bql, self.destination_dataset_table, self.write_disposition, - self.allow_large_results, self.udf_config, self.use_legacy_sql, - self.maximum_billing_tier) + self.create_disposition, self.allow_large_results, self.udf_config, + self.use_legacy_sql, self.maximum_billing_tier) From b35f887ad73dd2b22f182e7d99d53f3c9aabea50 Mon Sep 17 00:00:00 2001 From: Yu ISHIKAWA Date: Thu, 27 Jul 2017 14:21:55 -0700 Subject: [PATCH 2/2] Change the order of parameters --- airflow/contrib/hooks/bigquery_hook.py | 4 ++-- airflow/contrib/operators/bigquery_operator.py | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py index a8fba8a4b4a6f..73e0a43887950 100644 --- a/airflow/contrib/hooks/bigquery_hook.py +++ b/airflow/contrib/hooks/bigquery_hook.py @@ -192,11 +192,11 @@ def __init__(self, service, project_id): def run_query( self, bql, destination_dataset_table = False, write_disposition = 'WRITE_EMPTY', - create_disposition='CREATE_IF_NEEDED', allow_large_results=False, udf_config = False, use_legacy_sql=True, - maximum_billing_tier=None): + maximum_billing_tier=None, + create_disposition='CREATE_IF_NEEDED'): """ Executes a BigQuery SQL query. Optionally persists results in a BigQuery table. See here: diff --git a/airflow/contrib/operators/bigquery_operator.py b/airflow/contrib/operators/bigquery_operator.py index ecb0b9c884f61..aaffc2ef4d4fa 100644 --- a/airflow/contrib/operators/bigquery_operator.py +++ b/airflow/contrib/operators/bigquery_operator.py @@ -61,13 +61,13 @@ def __init__(self, bql, destination_dataset_table=False, write_disposition='WRITE_EMPTY', - create_disposition='CREATE_IF_NEEDED', allow_large_results=False, bigquery_conn_id='bigquery_default', delegate_to=None, udf_config=False, use_legacy_sql=True, maximum_billing_tier=None, + create_disposition='CREATE_IF_NEEDED', *args, **kwargs): super(BigQueryOperator, self).__init__(*args, **kwargs) @@ -89,5 +89,6 @@ def execute(self, context): conn = hook.get_conn() cursor = conn.cursor() cursor.run_query(self.bql, self.destination_dataset_table, self.write_disposition, - self.create_disposition, self.allow_large_results, self.udf_config, - self.use_legacy_sql, self.maximum_billing_tier) + self.allow_large_results, self.udf_config, + self.use_legacy_sql, self.maximum_billing_tier, + self.create_disposition)