Merge pull request tylertreat#101 from Jonovono/use-legacy-sql

tylertreat · tylertreat · commit 35ec0f0c33e7 · 2016-05-24T16:39:05.000-05:00
Add support for UseLegacySQL boolean in order to use standard SQL
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -262,7 +262,7 @@ def _insert_job(self, body_object):
             body=body_object
         ).execute()
 
-    def query(self, query, max_results=None, timeout=0, dry_run=False):
+    def query(self, query, max_results=None, timeout=0, dry_run=False, use_legacy_sql=None):
         """Submit a query to BigQuery.
 
         Parameters
@@ -278,6 +278,9 @@ def query(self, query, max_results=None, timeout=0, dry_run=False):
             If True, the query isn't actually run. A valid query will return an
             empty response, while an invalid one will return the same error
             message it would if it wasn't a dry run.
+        use_legacy_sql : bool, optional. Default True.
+            If False, the query will use BigQuery's standard SQL (https://cloud.google.com/bigquery/sql-reference/)
+
 
         Returns
         -------
@@ -298,8 +301,12 @@ def query(self, query, max_results=None, timeout=0, dry_run=False):
             'query': query,
             'timeoutMs': timeout * 1000,
             'dryRun': dry_run,
-            'maxResults': max_results,
+            'maxResults': max_results
         }
+
+        if use_legacy_sql is not None:
+            query_data['useLegacySql'] = use_legacy_sql
+
         return self._submit_query_job(query_data)
 
     def get_query_schema(self, job_id):
@@ -1027,6 +1034,7 @@ def write_to_table(
             priority=None,
             create_disposition=None,
             write_disposition=None,
+            use_legacy_sql=None
     ):
         """
         Write query result to table. If dataset or table is not provided,
@@ -1055,6 +1063,9 @@ def write_to_table(
             One of the JOB_CREATE_* constants
         write_disposition : str, optional
             One of the JOB_WRITE_* constants
+        use_legacy_sql:
+            If False, the query will use BigQuery's standard SQL (https://cloud.google.com/bigquery/sql-reference/)
+
 
         Returns
         -------
@@ -1084,6 +1095,9 @@ def write_to_table(
         if use_query_cache is not None:
             configuration['useQueryCache'] = use_query_cache
 
+        if use_legacy_sql is not None:
+            configuration['useLegacySql'] = use_legacy_sql
+
         if priority:
             configuration['priority'] = priority
 
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -253,6 +253,7 @@ def test_query(self):
         self.assertEquals(job_id, 'spiderman')
         self.assertEquals(results, [])
 
+
     def test_query_max_results_set(self):
         """Ensure that we retrieve the job id from the query and the maxResults
         parameter is set.
@@ -418,6 +419,30 @@ def test_query_with_results(self):
         self.assertEquals(job_id, 'spiderman')
         self.assertEquals(results, [{'foo': 10}])
 
+    def test_query_with_using_legacy_sql(self):
+        """Ensure that use_legacy_sql bool gets used"""
+
+        mock_query_job = mock.Mock()
+        expected_job_id = 'spiderman'
+        expected_job_ref = {'jobId': expected_job_id}
+
+        mock_query_job.execute.return_value = {
+            'jobReference': expected_job_ref,
+            'jobComplete': True
+        }
+
+        self.mock_job_collection.query.return_value = mock_query_job
+
+        job_id, results = self.client.query(self.query, use_legacy_sql=False)
+
+        self.mock_job_collection.query.assert_called_once_with(
+            projectId=self.project_id,
+            body={'query': self.query, 'timeoutMs': 0, 'dryRun': False,
+                  'maxResults': None, 'useLegacySql': False}
+        )
+        self.assertEquals(job_id, 'spiderman')
+        self.assertEquals(results, [])
+
 
 class TestGetQueryResults(unittest.TestCase):
 
