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_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,28 @@ 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 +180,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

@@ -697,6 +697,9 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
         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):