add an option to deferred fetch result in Cursor.execute()

tests/unit/conftest.py

@@ -14,8 +14,10 @@
 
 import pytest
 
+from tests.unit.oauth_test_utils import SERVER_ADDRESS
 
-@pytest.fixture(scope="session")
+
+@pytest.fixture
 def sample_post_response_data():
  """
  This is the response to the first HTTP request (a POST) from an actual
@@ -38,10 +40,10 @@ def sample_post_response_data():
  """
 
  yield {
- "nextUri": "https://coordinator:8080/v1/statement/20210817_140827_00000_arvdv/1",
+ "nextUri": f"{SERVER_ADDRESS}:8080/v1/statement/20210817_140827_00000_arvdv/1",
  "id": "20210817_140827_00000_arvdv",
  "taskDownloadUris": [],
- "infoUri": "https://coordinator:8080/query.html?20210817_140827_00000_arvdv",
+ "infoUri": f"{SERVER_ADDRESS}:8080/query.html?20210817_140827_00000_arvdv",
  "stats": {
  "scheduled": False,
  "runningSplits": 0,
@@ -60,7 +62,7 @@ def sample_post_response_data():
  }
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def sample_get_response_data():
  """
  This is the response to the second HTTP request (a GET) from an actual
@@ -73,7 +75,7 @@ def sample_get_response_data():
  """
  yield {
  "id": "20210817_140827_00000_arvdv",
- "nextUri": "coordinator:8080/v1/statement/20210817_140827_00000_arvdv/2",
+ "nextUri": f"{SERVER_ADDRESS}:8080/v1/statement/20210817_140827_00000_arvdv/2",
  "data": [
  ["UUID-0", "http://worker0:8080", "0.157", False, "active"],
  ["UUID-1", "http://worker1:8080", "0.157", False, "active"],
@@ -132,7 +134,7 @@ def sample_get_response_data():
  },
  ],
  "taskDownloadUris": [],
- "partialCancelUri": "http://localhost:8080/v1/stage/20210817_140827_00000_arvdv.0", # NOQA: E501
+ "partialCancelUri": f"{SERVER_ADDRESS}:8080/v1/stage/20210817_140827_00000_arvdv.0", # NOQA: E501
  "stats": {
  "nodes": 2,
  "processedBytes": 880,
@@ -181,11 +183,11 @@ def sample_get_response_data():
  "queuedSplits": 0,
  "wallTimeMillis": 36,
  },
- "infoUri": "http://coordinator:8080/query.html?20210817_140827_00000_arvdv", # NOQA: E501
+ "infoUri": f"{SERVER_ADDRESS}:8080/query.html?20210817_140827_00000_arvdv", # NOQA: E501
  }
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def sample_get_error_response_data():
  yield {
  "error": {
@@ -195,8 +197,7 @@ def sample_get_error_response_data():
  "errorType": "USER_ERROR",
  "failureInfo": {
  "errorLocation": {"columnNumber": 15, "lineNumber": 1},
- "message": "line 1:15: Schema must be specified "
- "when session schema is not set",
+ "message": "line 1:15: Schema must be specified when session schema is not set",
  "stack": [
  "io.trino.sql.analyzer.SemanticExceptions.semanticException(SemanticExceptions.java:48)",
  "io.trino.sql.analyzer.SemanticExceptions.semanticException(SemanticExceptions.java:43)",
@@ -241,7 +242,7 @@ def sample_get_error_response_data():
  "message": "line 1:15: Schema must be specified when session schema is not set",
  },
  "id": "20210817_140827_00000_arvdv",
- "infoUri": "http://localhost:8080/query.html?20210817_140827_00000_arvdv",
+ "infoUri": f"{SERVER_ADDRESS}:8080/query.html?20210817_140827_00000_arvdv",
  "stats": {
  "completedSplits": 0,
  "cpuTimeMillis": 0,

tests/unit/test_client.py

@@ -1018,6 +1018,57 @@ def json(self):
  assert isinstance(result, TrinoResult)
 
 
+def test_trino_query_deferred_fetch(sample_get_response_data):
+ """
+ Validates that the `TrinoQuery.execute` function deferred_fetch and non-block execution
+ """
+
+ class MockResponse(mock.Mock):
+ # Fake response class
+ @property
+ def headers(self):
+ return {
+ 'X-Trino-Fake-1': 'one',
+ 'X-Trino-Fake-2': 'two',
+ }
+
+ def json(self):
+ return sample_get_response_data
+
+ rows = sample_get_response_data['data']
+ sample_get_response_data['data'] = []
+ sql = 'SELECT 1'
+ request = TrinoRequest(
+ host="coordinator",
+ port=8080,
+ client_session=ClientSession(
+ user="test",
+ source="test",
+ catalog="test",
+ schema="test",
+ properties={},
+ ),
+ http_scheme="http",
+ )
+ query = TrinoQuery(
+ request=request,
+ query=sql
+ )
+
+ with \
+ mock.patch.object(request, 'post', return_value=MockResponse()), \
+ mock.patch.object(query, 'fetch', return_value=rows) as mock_fetch:
+ result = query.execute()
+ mock_fetch.assert_called_once()
+ assert result.rows == rows
+
+ with \
+ mock.patch.object(request, 'post', return_value=MockResponse()), \
+ mock.patch.object(query, 'fetch', return_value=rows) as mock_fetch:
+ result = query.execute(deferred_fetch=True)
+ mock_fetch.assert_not_called()
+
+
 def test_delay_exponential_without_jitter():
  max_delay = 1200.0
  get_delay = _DelayExponential(base=5, jitter=False, max_delay=max_delay)

tests/unit/test_dbapi.py

@@ -14,7 +14,6 @@
 from unittest.mock import patch
 
 import httpretty
-from httpretty import httprettified
 from requests import Session
 
 from tests.unit.oauth_test_utils import (
@@ -58,7 +57,7 @@ def test_http_session_is_defaulted_when_not_specified(mock_client):
  assert mock_client.TrinoRequest.http.Session.return_value in request_args
 
 
-@httprettified
+@httpretty.activate
 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())
@@ -73,13 +72,15 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl
  httpretty.register_uri(
  method=httpretty.POST,
  uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
- body=post_statement_callback)
+ 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)
+ body=get_statement_callback
+ )
 
  # bind get token
  get_token_callback = GetTokenCallback(token_server, token)
@@ -122,7 +123,7 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl
  assert len(_get_token_requests(challenge_id)) == 2
 
 
-@httprettified
+@httpretty.activate
 def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post_response_data,
  sample_get_response_data):
  token = str(uuid.uuid4())
@@ -188,7 +189,7 @@ def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post
  assert len(_get_token_requests(challenge_id)) == 1
 
 
-@httprettified
+@httpretty.activate
 def test_token_retrieved_once_when_multithreaded(sample_post_response_data, sample_get_response_data):
  token = str(uuid.uuid4())
  challenge_id = str(uuid.uuid4())

trino/auth.py

@@ -395,7 +395,7 @@ def _determine_host(url: Optional[str]) -> Any:
 
 
 class OAuth2Authentication(Authentication):
- def __init__(self, redirect_auth_url_handler: CompositeRedirectHandler = CompositeRedirectHandler([
+ def __init__(self, redirect_auth_url_handler: RedirectHandler = CompositeRedirectHandler([
  WebBrowserRedirectHandler(),
  ConsoleRedirectHandler()
  ])):

trino/client.py

@@ -688,11 +688,11 @@ def __init__(self, query, rows: List[Any]):
  self._rownumber = 0
 
  @property
- def rows(self):
+ def rows(self) -> List[Any]:
  return self._rows
 
  @rows.setter
- def rows(self, rows):
+ def rows(self, rows: List[Any]):
  self._rows = rows
 
  @property
@@ -702,14 +702,13 @@ def rownumber(self) -> int:
  def __iter__(self):
  # A query only transitions to a FINISHED state when the results are fully consumed:
  # The reception of the data is acknowledged by calling the next_uri before exposing the data through dbapi.
- while not self._query.finished or self._rows is not None:
- next_rows = self._query.fetch() if not self._query.finished else None
+ while not self._query.finished or self._rows:
  for row in self._rows:
  self._rownumber += 1
  logger.debug("row %s", row)
  yield row
 
- self._rows = next_rows
+ self._rows = self._query.fetch() if not self._query.finished else []
 
 
 class TrinoQuery(object):
@@ -778,13 +777,18 @@ def result(self):
  def info_uri(self):
  return self._info_uri
 
- def execute(self, additional_http_headers=None) -> TrinoResult:
- """Initiate a Trino query by sending the SQL statement
-
- This is the first HTTP request sent to the coordinator.
- It sets the query_id and returns a Result object used to
- track the rows returned by the query. To fetch all rows,
- call fetch() until finished is true.
+ def execute(
+ self,
+ additional_http_headers: Optional[Dict[str, Any]] = None,
+ deferred_fetch: bool = False,
+ ) -> TrinoResult:
+ """Initiate a Trino query by sending the SQL statement to the coordinator.
+ To fetch all rows, call fetch() until finished is true.
+
+ Parameters:
+ additional_http_headers: extra headers send to the Trino server.
+ deferred_fetch: By default, the execution is blocked until at least one row is received
+ or query is finished or cancelled. To continue without waiting the result, set deferred_fetch=True.
  """
  if self.cancelled:
  raise exceptions.TrinoUserError("Query has been cancelled", self.query_id)
@@ -805,9 +809,11 @@ 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 or query is finished or cancelled
- while not self.finished and not self.cancelled and len(self._result.rows) == 0:
- self._result.rows += self.fetch()
+ if not deferred_fetch:
+ # Execute should block until at least one row is received or query is finished or cancelled
+ 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):

trino/dbapi.py

@@ -558,7 +558,10 @@ def _deallocate_prepared_statement(self, statement_name: str) -> None:
  def _generate_unique_statement_name(self):
  return 'st_' + uuid.uuid4().hex.replace('-', '')
 
- def execute(self, operation, params=None):
+ def execute(self, operation, params=None, **kwargs: Any):
+ additional_http_headers = kwargs.get("additional_http_headers", None)
+ deferred_fetch = kwargs.get("deferred_fetch", False)
+
  if params:
  assert isinstance(params, (list, tuple)), (
  'params must be a list or tuple containing the query '
@@ -575,7 +578,10 @@ def execute(self, operation, params=None):
  self._query = self._execute_prepared_statement(
  statement_name, params
  )
- self._iterator = iter(self._query.execute())
+ self._iterator = iter(self._query.execute(
+ additional_http_headers=additional_http_headers,
+ deferred_fetch=deferred_fetch,
+ ))
  finally:
  # Send deallocate statement
  # At this point the query can be deallocated since it has already
@@ -584,12 +590,18 @@ def execute(self, operation, params=None):
  self._deallocate_prepared_statement(statement_name)
  else:
  self._query = self._execute_immediate_statement(operation, params)
- self._iterator = iter(self._query.execute())
+ self._iterator = iter(self._query.execute(
+ additional_http_headers=additional_http_headers,
+ deferred_fetch=deferred_fetch,
+ ))
 
  else:
  self._query = trino.client.TrinoQuery(self._request, query=operation,
  legacy_primitive_types=self._legacy_primitive_types)
- self._iterator = iter(self._query.execute())
+ self._iterator = iter(self._query.execute(
+ additional_http_headers=additional_http_headers,
+ deferred_fetch=deferred_fetch,
+ ))
  return self
 
  def executemany(self, operation, seq_of_params):

trino/sqlalchemy/dialect.py

@@ -396,7 +396,8 @@ def _get_default_schema_name(self, connection: Connection) -> Optional[str]:
  def do_execute(
  self, cursor: Cursor, statement: str, parameters: Tuple[Any, ...], context: DefaultExecutionContext = None
  ):
- cursor.execute(statement, parameters)
+ execution_options = (context.execution_options or {}) if context else {}
+ cursor.execute(statement, parameters, **execution_options)
 
  def do_rollback(self, dbapi_connection: trino_dbapi.Connection):
  if dbapi_connection.transaction is not None: