Add unittests for BQ 2.0

Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/ce3f1726
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/ce3f1726
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/ce3f1726

Branch: refs/heads/python-sdk
Commit: ce3f1726ff9c69d718c41cae02cf2d4a370f5a22
Parents: 0a33965
Author: Sourabh Bajaj <sourabhba...@google.com>
Authored: Tue Nov 1 13:25:28 2016 -0700
Committer: Robert Bradshaw <rober...@google.com>
Committed: Thu Nov 3 10:38:44 2016 -0700

----------------------------------------------------------------------
 sdks/python/apache_beam/io/bigquery.py      |  2 +-
 sdks/python/apache_beam/io/bigquery_test.py | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/ce3f1726/sdks/python/apache_beam/io/bigquery.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/bigquery.py 
b/sdks/python/apache_beam/io/bigquery.py
index 41b1bdc..a3999af 100644
--- a/sdks/python/apache_beam/io/bigquery.py
+++ b/sdks/python/apache_beam/io/bigquery.py
@@ -315,7 +315,7 @@ class BigQuerySource(dataflow_io.NativeSource):
       useLegacySql: Specifies whether to use BigQuery's legacy
         SQL dialect for this query. The default value is true. If set to false,
         the query will use BigQuery's updated SQL dialect with improved
-        standards compliance. This parameter is forced to True for table 
inputs.
+        standards compliance. This parameter is ignored for table inputs.
 
     Raises:
       ValueError: if any of the following is true

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/ce3f1726/sdks/python/apache_beam/io/bigquery_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/bigquery_test.py 
b/sdks/python/apache_beam/io/bigquery_test.py
index 09e64ce..ac71880 100644
--- a/sdks/python/apache_beam/io/bigquery_test.py
+++ b/sdks/python/apache_beam/io/bigquery_test.py
@@ -130,6 +130,12 @@ class TestBigQuerySource(unittest.TestCase):
     source = beam.io.BigQuerySource(query='my_query')
     self.assertEqual(source.query, 'my_query')
     self.assertIsNone(source.table_reference)
+    self.assertTrue(source.use_legacy_sql)
+
+  def test_specify_query_sql_format(self):
+    source = beam.io.BigQuerySource(query='my_query', use_legacy_sql=False)
+    self.assertEqual(source.query, 'my_query')
+    self.assertFalse(source.use_legacy_sql)
 
 
 class TestBigQuerySink(unittest.TestCase):
@@ -238,6 +244,24 @@ class TestBigQueryReader(unittest.TestCase):
         actual_rows.append(row)
     self.assertEqual(actual_rows, expected_rows)
     self.assertEqual(schema, reader.schema)
+    self.assertTrue(reader.use_legacy_sql)
+
+  def test_read_from_query_sql_format(self):
+    client = mock.Mock()
+    client.jobs.Insert.return_value = bigquery.Job(
+        jobReference=bigquery.JobReference(
+            jobId='somejob'))
+    table_rows, schema, expected_rows = self.get_test_rows()
+    client.jobs.GetQueryResults.return_value = 
bigquery.GetQueryResultsResponse(
+        jobComplete=True, rows=table_rows, schema=schema)
+    actual_rows = []
+    with beam.io.BigQuerySource(
+        query='query', use_legacy_sql=False).reader(client) as reader:
+      for row in reader:
+        actual_rows.append(row)
+    self.assertEqual(actual_rows, expected_rows)
+    self.assertEqual(schema, reader.schema)
+    self.assertFalse(reader.use_legacy_sql)
 
   def test_using_both_query_and_table_fails(self):
     with self.assertRaises(ValueError) as exn:

Reply via email to