[AIRFLOW-6855]: Escape project_dataset_table in SQL query in gcs to bq … (#7475)

mtagle · web-flow · commit 0de0347b27a9 · 2020-03-17T11:00:46.000-07:00
* [AIRFLOW-6855]: Escape project_dataset_table in SQL query in gcs to bq operator Without escaping, if the project is specified in project_dataset_table and contains a -, the query will fail with an error. * Make string formatting in gcs_to_bigquery f-strings, add unit tests. * pylint appeasement, hopefully. * More not understanding how mocking works in python... * Add task ids to tests, remove tst_gcs_to_bigquery.py as a missing test file. * maybe more correct mocking of class vairables.
diff --git a/airflow/providers/google/cloud/operators/gcs_to_bigquery.py b/airflow/providers/google/cloud/operators/gcs_to_bigquery.py
@@ -298,10 +298,15 @@ def execute(self, context):
                 cluster_fields=self.cluster_fields,
                 encryption_configuration=self.encryption_configuration)
 
+        if cursor.use_legacy_sql:
+            escaped_table_name = f'[{self.destination_project_dataset_table}]'
+        else:
+            escaped_table_name = f'`{self.destination_project_dataset_table}`'
+
         if self.max_id_key:
             cursor.execute('SELECT MAX({}) FROM {}'.format(
                 self.max_id_key,
-                self.destination_project_dataset_table))
+                escaped_table_name))
             row = cursor.fetchone()
             max_id = row[0] if row[0] else 0
             self.log.info(
diff --git a/tests/providers/google/cloud/operators/test_gcs_to_bigquery.py b/tests/providers/google/cloud/operators/test_gcs_to_bigquery.py
@@ -0,0 +1,70 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import unittest
+
+import mock
+
+from airflow.providers.google.cloud.operators.gcs_to_bigquery import GCSToBigQueryOperator
+
+TASK_ID = 'test-gcs-to-bq-operator'
+TEST_EXPLICIT_DEST = 'test-project.dataset.table'
+TEST_BUCKET = 'test-bucket'
+MAX_ID_KEY = 'id'
+TEST_SOURCE_OBJECTS = ['test/objects/*']
+
+
+class TestGoogleCloudStorageToBigQueryOperator(unittest.TestCase):
+
+    @mock.patch('airflow.providers.google.cloud.operators.gcs_to_bigquery.BigQueryHook')
+    def test_execute_explicit_project_legacy(self, bq_hook):
+        operator = GCSToBigQueryOperator(task_id=TASK_ID,
+                                         bucket=TEST_BUCKET,
+                                         source_objects=TEST_SOURCE_OBJECTS,
+                                         destination_project_dataset_table=TEST_EXPLICIT_DEST,
+                                         max_id_key=MAX_ID_KEY)
+
+        # using legacy SQL
+        bq_hook.return_value.get_conn.return_value.cursor.return_value.use_legacy_sql = True
+
+        operator.execute(None)
+
+        bq_hook.return_value \
+            .get_conn.return_value \
+            .cursor.return_value \
+            .execute \
+            .assert_called_once_with("SELECT MAX(id) FROM [test-project.dataset.table]")
+
+    @mock.patch('airflow.providers.google.cloud.operators.gcs_to_bigquery.BigQueryHook')
+    def test_execute_explicit_project(self, bq_hook):
+        operator = GCSToBigQueryOperator(task_id=TASK_ID,
+                                         bucket=TEST_BUCKET,
+                                         source_objects=TEST_SOURCE_OBJECTS,
+                                         destination_project_dataset_table=TEST_EXPLICIT_DEST,
+                                         max_id_key=MAX_ID_KEY)
+
+        # using non-legacy SQL
+        bq_hook.return_value.get_conn.return_value.cursor.return_value.use_legacy_sql = False
+
+        operator.execute(None)
+
+        bq_hook.return_value \
+            .get_conn.return_value \
+            .cursor.return_value \
+            .execute \
+            .assert_called_once_with("SELECT MAX(id) FROM `test-project.dataset.table`")
diff --git a/tests/test_project_structure.py b/tests/test_project_structure.py
@@ -35,7 +35,6 @@
     'tests/providers/apache/spark/hooks/test_spark_jdbc_script.py',
     'tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py',
     'tests/providers/google/cloud/operators/test_datastore.py',
-    'tests/providers/google/cloud/operators/test_gcs_to_bigquery.py',
     'tests/providers/google/cloud/operators/test_sql_to_gcs.py',
     'tests/providers/google/cloud/sensors/test_bigquery.py',
     'tests/providers/google/cloud/utils/test_field_sanitizer.py',
