From b899ad12e17cb87c58d3ae46b4388d917c5743f2 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Thu, 12 Nov 2020 17:22:09 -0600
Subject: [PATCH] fix(dbapi): allow rows to be fetched from scripts (#387)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The `is_dml` logic is not needed now that we moved to `getQueryResults` instead of `tabledata.list` (https://github.com/googleapis/python-bigquery/pull/375).

Previously, the destination table of a DML query would return a non-null value that was unreadable or would return nonsense with DML (and some DDL) queries.

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Towards #377  🦕
---
 google/cloud/bigquery/dbapi/cursor.py |  8 -----
 tests/system.py                       | 43 ++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 9 deletions(-)

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`"
     )