add an option deferred_fetch to Cursor.execute()

Author: dungdm93

trino/client.py

@@ -775,7 +775,7 @@ def result(self):
     def info_uri(self):
         return self._info_uri
 
-    def execute(self, additional_http_headers=None) -> TrinoResult:
+    def execute(self, additional_http_headers=None, deferred_fetch: bool = False) -> TrinoResult:
         """Initiate a Trino query by sending the SQL statement
 
         This is the first HTTP request sent to the coordinator.
@@ -799,9 +799,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

@@ -377,7 +377,7 @@ 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)
+        cursor.execute(statement, parameters, **context.execution_options)
 
     def do_rollback(self, dbapi_connection: trino_dbapi.Connection):
         if dbapi_connection.transaction is not None: