diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py index dc98b896fdb97..73e0a43887950 100644 --- a/airflow/contrib/hooks/bigquery_hook.py +++ b/airflow/contrib/hooks/bigquery_hook.py @@ -195,7 +195,8 @@ def run_query( 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: @@ -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..aaffc2ef4d4fa 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. @@ -61,12 +67,14 @@ def __init__(self, udf_config=False, use_legacy_sql=True, maximum_billing_tier=None, + create_disposition='CREATE_IF_NEEDED', *args, **kwargs): super(BigQueryOperator, self).__init__(*args, **kwargs) 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,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.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)