ability to execute ddl queries

Author: vgvoleg

examples/query-service/basic_example.py

@@ -14,15 +14,22 @@ def main():
     except TimeoutError:
         raise RuntimeError("Connect failed to YDB")
 
-    # client = ydb.QueryClientSync(driver)
-    # session = client.session().create()
     pool = ydb.QuerySessionPool(driver)
-    # with pool.checkout() as session:
+
+    print("=" * 50)
+    print("DELETE TABLE IF EXISTS")
+    pool.execute_with_retries("drop table if exists example;", ddl=True)
+
+    print("=" * 50)
+    print("CREATE TABLE")
+    pool.execute_with_retries("CREATE TABLE example(key UInt64, value String, PRIMARY KEY (key));", ddl=True)
 
     def callee(session):
         print("=" * 50)
+        session.execute("""delete from example;""")
+
         print("BEFORE ACTION")
-        it = session.execute("""SELECT COUNT(*) FROM example;""")
+        it = session.execute("""SELECT COUNT(*) as rows_count FROM example;""")
         for result_set in it:
             print(f"rows: {str(result_set.rows)}")
 
@@ -34,15 +41,15 @@ def callee(session):
 
         tx.execute("""INSERT INTO example (key, value) VALUES (0055, "onepieceisreal");""")
 
-        for result_set in tx.execute("""SELECT COUNT(*) FROM example;"""):
+        for result_set in tx.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
             print(f"rows: {str(result_set.rows)}")
 
         tx.commit()
 
         print("=" * 50)
         print("AFTER COMMIT TX")
 
-        for result_set in session.execute("""SELECT COUNT(*) FROM example;"""):
+        for result_set in session.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
             print(f"rows: {str(result_set.rows)}")
 
         print("=" * 50)
@@ -54,15 +61,15 @@ def callee(session):
 
         tx.execute("""INSERT INTO example (key, value) VALUES (0066, "onepieceisreal");""")
 
-        for result_set in tx.execute("""SELECT COUNT(*) FROM example;"""):
+        for result_set in tx.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
             print(f"rows: {str(result_set.rows)}")
 
         tx.rollback()
 
         print("=" * 50)
         print("AFTER ROLLBACK TX")
 
-        for result_set in session.execute("""SELECT COUNT(*) FROM example;"""):
+        for result_set in session.execute("""SELECT COUNT(*) as rows_count FROM example;"""):
             print(f"rows: {str(result_set.rows)}")
 
     pool.retry_operation_sync(callee)

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -161,16 +161,17 @@ def to_proto(self) -> ydb_query_pb2.TransactionControl:
 class ExecuteQueryRequest(IToProto):
     session_id: str
     query_content: QueryContent
-    tx_control: TransactionControl
+    tx_control: Optional[TransactionControl] = None
     concurrent_result_sets: Optional[bool] = False
     exec_mode: Optional[int] = None
     parameters: Optional[dict] = None
     stats_mode: Optional[int] = None
 
     def to_proto(self) -> ydb_query_pb2.ExecuteQueryRequest:
+        tx_control = self.tx_control.to_proto() if self.tx_control is not None else self.tx_control
         return ydb_query_pb2.ExecuteQueryRequest(
             session_id=self.session_id,
-            tx_control=self.tx_control.to_proto(),
+            tx_control=tx_control,
             query_content=self.query_content.to_proto(),
             exec_mode=self.exec_mode,
             stats_mode=self.stats_mode,

ydb/query/base.py

@@ -155,43 +155,40 @@ def create_execute_query_request(
     exec_mode: QueryExecMode = None,
     parameters: dict = None,
     concurrent_result_sets: bool = False,
+    empty_tx_control: bool = False,
 ):
     syntax = QuerySyntax.YQL_V1 if not syntax else syntax
     exec_mode = QueryExecMode.EXECUTE if not exec_mode else exec_mode
-    if tx_id:
-        req = ydb_query.ExecuteQueryRequest(
-            session_id=session_id,
-            query_content=ydb_query.QueryContent.from_public(
-                query=query,
-                syntax=syntax,
-            ),
-            tx_control=ydb_query.TransactionControl(
-                tx_id=tx_id,
-                commit_tx=commit_tx,
-            ),
-            exec_mode=exec_mode,
-            parameters=parameters,
-            concurrent_result_sets=concurrent_result_sets,
+
+    tx_control = None
+    if empty_tx_control:
+        tx_control = None
+    elif tx_id:
+        tx_control = ydb_query.TransactionControl(
+            tx_id=tx_id,
+            commit_tx=commit_tx,
         )
     else:
         tx_mode = tx_mode if tx_mode is not None else QuerySerializableReadWrite()
-        req = ydb_query.ExecuteQueryRequest(
-            session_id=session_id,
-            query_content=ydb_query.QueryContent.from_public(
-                query=query,
-                syntax=syntax,
+        tx_control = ydb_query.TransactionControl(
+            begin_tx=ydb_query.TransactionSettings(
+                tx_mode=tx_mode,
             ),
-            tx_control=ydb_query.TransactionControl(
-                begin_tx=ydb_query.TransactionSettings(
-                    tx_mode=tx_mode,
-                ),
-                commit_tx=commit_tx,
-            ),
-            exec_mode=exec_mode,
-            parameters=parameters,
-            concurrent_result_sets=concurrent_result_sets,
+            commit_tx=commit_tx,
         )
 
+    req = ydb_query.ExecuteQueryRequest(
+        session_id=session_id,
+        query_content=ydb_query.QueryContent.from_public(
+            query=query,
+            syntax=syntax,
+        ),
+        tx_control=tx_control,
+        exec_mode=exec_mode,
+        parameters=parameters,
+        concurrent_result_sets=concurrent_result_sets,
+    )
+
     return req.to_proto()
 
 

ydb/query/pool.py

@@ -91,8 +91,8 @@ def retry_operation_impl(callee: Callable, retry_settings: RetrySettings = None,
             # you should provide your own handler you want
             retry_settings.unknown_error_handler(e)
             raise
-
-    raise status
+    if status:
+        raise status
 
 
 class QuerySessionPool:
@@ -116,6 +116,20 @@ def wrapped_callee():
             else:
                 return next_opt.result
 
+    def execute_with_retries(self, query: str, ddl: bool = False, retry_settings: RetrySettings = None, *args, **kwargs):
+        retry_settings = RetrySettings() if retry_settings is None else retry_settings
+        with self.checkout() as session:
+            def wrapped_callee():
+                it = session.execute(query, empty_tx_control=ddl)
+                return [result_set for result_set in it]
+
+            opt_generator = retry_operation_impl(wrapped_callee, retry_settings, *args, **kwargs)
+            for next_opt in opt_generator:
+                if isinstance(next_opt, YdbRetryOperationSleepOpt):
+                    time.sleep(next_opt.timeout)
+                else:
+                    return next_opt.result
+
 
 class SimpleQuerySessionCheckout:
     def __init__(self, pool: QuerySessionPool):

ydb/query/session.py

@@ -150,11 +150,27 @@ def _attach_call(self):
             _apis.QueryService.AttachSession,
         )
 
-    def _execute_call(self, query: str, commit_tx: bool):
+    def _execute_call(
+        self,
+        query: str,
+        commit_tx: bool = False,
+        tx_mode: base.BaseQueryTxMode = None,
+        syntax: base.QuerySyntax = None,
+        exec_mode: base.QueryExecMode = None,
+        parameters: dict = None,
+        concurrent_result_sets: bool = False,
+        empty_tx_control: bool = False,
+    ):
         request = base.create_execute_query_request(
             query=query,
             session_id=self._state.session_id,
             commit_tx=commit_tx,
+            tx_mode=tx_mode,
+            syntax=syntax,
+            exec_mode=exec_mode,
+            parameters=parameters,
+            concurrent_result_sets=concurrent_result_sets,
+            empty_tx_control=empty_tx_control,
         )
 
         return self._driver(
@@ -225,10 +241,28 @@ def transaction(self, tx_mode: base.BaseQueryTxMode = None) -> base.IQueryTxCont
             tx_mode,
         )
 
-    def execute(self, query: str, parameters=None):
+    def execute(
+        self,
+        query: str,
+        tx_mode: base.BaseQueryTxMode = None,
+        syntax: base.QuerySyntax = None,
+        exec_mode: base.QueryExecMode = None,
+        parameters: dict = None,
+        concurrent_result_sets: bool = False,
+        empty_tx_control: bool = False
+    ):
         self._state._check_session_ready_to_use()
 
-        stream_it = self._execute_call(query, commit_tx=True)
+        stream_it = self._execute_call(
+            query=query,
+            commit_tx=True,
+            tx_mode=tx_mode,
+            syntax=syntax,
+            exec_mode=exec_mode,
+            parameters=parameters,
+            concurrent_result_sets=concurrent_result_sets,
+            empty_tx_control=empty_tx_control,
+        )
 
         return _utilities.SyncResponseIterator(
             stream_it,