Execute should block until at least one row is received

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -153,8 +153,8 @@ def test_execute_many_without_params(trino_connection):
     cur = trino_connection.cursor()
     cur.execute("CREATE TABLE memory.default.test_execute_many_without_param (value varchar)")
     cur.fetchall()
-    cur.executemany("INSERT INTO memory.default.test_execute_many_without_param (value) VALUES (?)", [])
     with pytest.raises(TrinoUserError) as e:
+        cur.executemany("INSERT INTO memory.default.test_execute_many_without_param (value) VALUES (?)", [])
         cur.fetchall()
     assert "Incorrect number of parameters: expected 1 but found 0" in str(e.value)
 
@@ -883,13 +883,12 @@ def test_transaction_autocommit(trino_connection_in_autocommit):
     with trino_connection_in_autocommit as connection:
         connection.start_transaction()
         cur = connection.cursor()
-        cur.execute(
-            """
-            CREATE TABLE memory.default.nation
-            AS SELECT * from tpch.tiny.nation
-            """)
-
         with pytest.raises(TrinoUserError) as transaction_error:
+            cur.execute(
+                """
+                CREATE TABLE memory.default.nation
+                AS SELECT * from tpch.tiny.nation
+                """)
             cur.fetchall()
         assert "Catalog only supports writes using autocommit: memory" \
                in str(transaction_error.value)

tests/unit/conftest.py

@@ -37,10 +37,10 @@ def sample_post_response_data():
     """
 
     yield {
-        "nextUri": "coordinator:8080/v1/statement/20210817_140827_00000_arvdv/1",
+        "nextUri": "https://coordinator:8080/v1/statement/20210817_140827_00000_arvdv/1",
         "id": "20210817_140827_00000_arvdv",
         "taskDownloadUris": [],
-        "infoUri": "http://coordinator:8080/query.html?20210817_140827_00000_arvdv",
+        "infoUri": "https://coordinator:8080/query.html?20210817_140827_00000_arvdv",
         "stats": {
             "scheduled": False,
             "runningSplits": 0,

tests/unit/test_client.py

@@ -892,9 +892,7 @@ def test_trino_result_response_headers():
         'X-Trino-Fake-2': 'two',
     })
 
-    result = TrinoResult(
-        query=mock_trino_query,
-    )
+    result = TrinoResult(query=mock_trino_query, rows=[])
     assert result.response_headers == mock_trino_query.response_headers
 
 

tests/unit/test_dbapi.py

@@ -51,21 +51,28 @@ def test_http_session_is_defaulted_when_not_specified(mock_client):
 
 
 @httprettified
-def test_token_retrieved_once_per_auth_instance(sample_post_response_data):
+def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sample_get_response_data):
     token = str(uuid.uuid4())
     challenge_id = str(uuid.uuid4())
 
     redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
     token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
 
     post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
+    get_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_get_response_data)
 
-    # bind post statement
+    # bind post statement to submit query
     httpretty.register_uri(
         method=httpretty.POST,
         uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
         body=post_statement_callback)
 
+    # bind get statement for result retrieval
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}/20210817_140827_00000_arvdv/1",
+        body=get_statement_callback)
+
     # bind get token
     get_token_callback = GetTokenCallback(token_server, token)
     httpretty.register_uri(
@@ -108,21 +115,29 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data):
 
 
 @httprettified
-def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post_response_data):
+def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post_response_data,
+                                                                     sample_get_response_data):
     token = str(uuid.uuid4())
     challenge_id = str(uuid.uuid4())
 
     redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
     token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
 
     post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
+    get_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_get_response_data)
 
-    # bind post statement
+    # bind post statement to submit query
     httpretty.register_uri(
         method=httpretty.POST,
         uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
         body=post_statement_callback)
 
+    # bind get statement for result retrieval
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}/20210817_140827_00000_arvdv/1",
+        body=get_statement_callback)
+
     # bind get token
     get_token_callback = GetTokenCallback(token_server, token)
     httpretty.register_uri(
@@ -166,21 +181,28 @@ def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post
 
 
 @httprettified
-def test_token_retrieved_once_when_multithreaded(sample_post_response_data):
+def test_token_retrieved_once_when_multithreaded(sample_post_response_data, sample_get_response_data):
     token = str(uuid.uuid4())
     challenge_id = str(uuid.uuid4())
 
     redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
     token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
 
     post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
+    get_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_get_response_data)
 
-    # bind post statement
+    # bind post statement to submit query
     httpretty.register_uri(
         method=httpretty.POST,
         uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
         body=post_statement_callback)
 
+    # bind get statement for result retrieval
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}/20210817_140827_00000_arvdv/1",
+        body=get_statement_callback)
+
     # bind get token
     get_token_callback = GetTokenCallback(token_server, token)
     httpretty.register_uri(

trino/client.py

@@ -592,12 +592,20 @@ class TrinoResult(object):
     https://docs.python.org/3/library/stdtypes.html#generator-types
     """
 
-    def __init__(self, query, rows=None):
+    def __init__(self, query, rows: List[Any]):
         self._query = query
         # Initial rows from the first POST request
         self._rows = rows
         self._rownumber = 0
 
+    @property
+    def rows(self):
+        return self._rows
+
+    @rows.setter
+    def rows(self, rows):
+        self._rows = rows
+
     @property
     def rownumber(self) -> int:
         return self._rownumber
@@ -650,7 +658,7 @@ def columns(self):
             while not self._columns and not self.finished and not self.cancelled:
                 # Columns are not returned immediately after query is submitted.
                 # Continue fetching data until columns information is available and push fetched rows into buffer.
-                self._result._rows += self.fetch()
+                self._result.rows += self.fetch()
         return self._columns
 
     @property
@@ -695,8 +703,11 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
             self._finished = True
 
         rows = self._row_mapper.map(status.rows) if self._row_mapper else status.rows
-
         self._result = TrinoResult(self, rows)
+
+        # Execute should block until at least one row is received
+        while not self.finished and not self.cancelled and len(self._result.rows) == 0:
+            self._result.rows += self.fetch()
         return self._result
 
     def _update_state(self, status):