diff --git a/google/cloud/bigquery/dbapi/cursor.py b/google/cloud/bigquery/dbapi/cursor.py index 74f8aec4e..f48b47c12 100644 --- a/google/cloud/bigquery/dbapi/cursor.py +++ b/google/cloud/bigquery/dbapi/cursor.py @@ -220,14 +220,6 @@ def _try_fetch(self, size=None): self._query_data = iter([]) return - is_dml = ( - self._query_job.statement_type - and self._query_job.statement_type.upper() != "SELECT" - ) - if is_dml: - self._query_data = iter([]) - return - if self._query_data is None: bqstorage_client = self.connection._bqstorage_client diff --git a/tests/system.py b/tests/system.py index 68fcb918c..51a47c0b7 100644 --- a/tests/system.py +++ b/tests/system.py @@ -180,6 +180,7 @@ class Config(object): CLIENT = None CURSOR = None + DATASET = None def setUpModule(): @@ -189,7 +190,9 @@ def setUpModule(): class TestBigQuery(unittest.TestCase): def setUp(self): - self.to_delete = [] + Config.DATASET = _make_dataset_id("bq_system_tests") + dataset = Config.CLIENT.create_dataset(Config.DATASET) + self.to_delete = [dataset] def tearDown(self): def _still_in_use(bad_request): @@ -1790,6 +1793,44 @@ def test_dbapi_fetchall(self): row_tuples = [r.values() for r in rows] self.assertEqual(row_tuples, [(1, 2), (3, 4), (5, 6)]) + def test_dbapi_fetchall_from_script(self): + query = """ + CREATE TEMP TABLE Example + ( + x INT64, + y STRING + ); + + INSERT INTO Example + VALUES (5, 'foo'), + (6, 'bar'), + (7, 'baz'); + + SELECT * + FROM Example + ORDER BY x ASC; + """ + + Config.CURSOR.execute(query) + self.assertEqual(Config.CURSOR.rowcount, 3, "expected 3 rows") + rows = Config.CURSOR.fetchall() + row_tuples = [r.values() for r in rows] + self.assertEqual(row_tuples, [(5, "foo"), (6, "bar"), (7, "baz")]) + + def test_dbapi_create_view(self): + + query = """ + CREATE VIEW {}.dbapi_create_view + AS SELECT name, SUM(number) AS total + FROM `bigquery-public-data.usa_names.usa_1910_2013` + GROUP BY name; + """.format( + Config.DATASET + ) + + Config.CURSOR.execute(query) + self.assertEqual(Config.CURSOR.rowcount, 0, "expected 0 rows") + @unittest.skipIf( bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" )