simple execute on session

Author: vgvoleg

examples/query-service/basic_example.py

@@ -0,0 +1,30 @@
+import ydb
+
+from ydb.query.session import QuerySessionSync
+
+
+def main():
+    driver_config = ydb.DriverConfig(
+        endpoint="grpc://localhost:2136",
+        database="/local",
+        # credentials=ydb.credentials_from_env_variables(),
+        # root_certificates=ydb.load_ydb_root_certificate(),
+    )
+    try:
+        driver = ydb.Driver(driver_config)
+        driver.wait(timeout=5)
+    except TimeoutError:
+        raise RuntimeError("Connect failed to YDB")
+
+    session = QuerySessionSync(driver)
+    session.create()
+
+    it = session.execute("select 1; select 2;")
+    for result_set in it:
+        print(f"columns: {str(result_set.columns)}")
+        print(f"rows: {str(result_set.rows)}")
+
+
+
+if __name__ == '__main__':
+    main()

ydb/_grpc/grpcwrapper/ydb_query.py

@@ -92,7 +92,6 @@ def to_proto(self) -> ydb_query_pb2.TransactionSettings:
             return ydb_query_pb2.TransactionSettings(online_read_only=self.tx_mode.to_proto())
         if self.tx_mode.name == 'stale_read_only':
             return ydb_query_pb2.TransactionSettings(stale_read_only=self.tx_mode.to_proto())
-        # TODO: add exception
 
 @dataclass
 class BeginTransactionRequest(IToProto):
@@ -102,7 +101,7 @@ class BeginTransactionRequest(IToProto):
     def to_proto(self) -> ydb_query_pb2.BeginTransactionRequest:
         return ydb_query_pb2.BeginTransactionRequest(
             session_id=self.session_id,
-            tx_settings=self.tx_settings
+            tx_settings=self.tx_settings.to_proto(),
             )
 
 @dataclass
@@ -120,7 +119,7 @@ def from_proto(msg: ydb_query_pb2.BeginTransactionResponse) -> "BeginTransaction
 @dataclass
 class QueryContent(IFromPublic, IToProto):
     text: str
-    syntax: Optional[str]
+    syntax: Optional[str] = None
 
     @staticmethod
     def from_public(query: str) -> "QueryContent":
@@ -132,9 +131,9 @@ def to_proto(self) -> ydb_query_pb2.QueryContent:
 
 @dataclass
 class TransactionControl(IToProto):
-    begin_tx: Optional[TransactionSettings]
-    commit_tx: Optional[bool]
-    tx_id: Optional[str]
+    begin_tx: Optional[TransactionSettings] = None
+    commit_tx: Optional[bool] = None
+    tx_id: Optional[str] = None
 
     def to_proto(self) -> ydb_query_pb2.TransactionControl:
         if self.tx_id:
@@ -143,19 +142,28 @@ def to_proto(self) -> ydb_query_pb2.TransactionControl:
                 commit_tx=self.commit_tx,
             )
         return ydb_query_pb2.TransactionControl(
-            begin_tx=self.begin_tx,
+            begin_tx=self.begin_tx.to_proto(),
             commit_tx=self.commit_tx
         )
 
 
 @dataclass
-class ExecuteQueryRequest:
-    exec_mode: Optional[str]
-    concurrent_result_sets: bool = False
-    parameters: Optional[dict]
-    query_content: QueryContent
+class ExecuteQueryRequest(IToProto):
     session_id: str
-    stats_mode: Optional[str]
+    query_content: QueryContent
     tx_control: TransactionControl
+    concurrent_result_sets: Optional[bool] = False
+    exec_mode: Optional[str] = None
+    parameters: Optional[dict] = None
+    stats_mode: Optional[str] = None
 
-
+    def to_proto(self) -> ydb_query_pb2.ExecuteQueryRequest:
+        return ydb_query_pb2.ExecuteQueryRequest(
+            session_id=self.session_id,
+            tx_control=self.tx_control.to_proto(),
+            query_content=self.query_content.to_proto(),
+            exec_mode=ydb_query_pb2.EXEC_MODE_EXECUTE,
+            stats_mode=self.stats_mode,
+            concurrent_result_sets=self.concurrent_result_sets,
+            parameters=self.parameters,
+        )

ydb/query/base.py

@@ -7,9 +7,12 @@
 from .._grpc.grpcwrapper.common_utils import (
     SupportedDriverType,
 )
-
-from .._grpc.grpcwrapper.ydb_query_public_types import BaseQueryTxMode
-
+from .._grpc.grpcwrapper import ydb_query
+from .._grpc.grpcwrapper.ydb_query_public_types import (
+    BaseQueryTxMode,
+    QuerySerializableReadWrite
+)
+from .. import convert
 
 class QueryClientSettings: ...
 
@@ -82,4 +85,22 @@ def session(self) -> IQuerySession:
         pass
 
 
+def create_execute_query_request(query: str, session_id: str, commit_tx: bool):
+    req = ydb_query.ExecuteQueryRequest(
+        session_id=session_id,
+        query_content=ydb_query.QueryContent.from_public(
+            query=query,
+        ),
+        tx_control=ydb_query.TransactionControl(
+            begin_tx=ydb_query.TransactionSettings(
+                tx_mode=QuerySerializableReadWrite(),
+            ),
+            commit_tx=commit_tx
+        ),
+    )
+
+    return req.to_proto()
+
+def wrap_execute_query_response(rpc_state, response_pb):
 
+    return convert.ResultSet.from_message(response_pb.result_set)

ydb/query/session.py

@@ -71,6 +71,19 @@ def _attach_call(self):
             _apis.QueryService.AttachSession,
         )
 
+    def _execute_call(self, query: str, commit_tx: bool):
+        request = base.create_execute_query_request(
+            query=query,
+            session_id=self._state.session_id,
+            commit_tx=commit_tx
+        )
+        return self._driver(
+            request,
+            _apis.QueryService.Stub,
+            _apis.QueryService.ExecuteQuery,
+            # base.wrap_execute_query_response
+        )
+
 class QuerySessionSync(BaseQuerySession):
     _stream = None
 
@@ -96,15 +109,13 @@ def _attach(self):
         ).start()
 
     def _chech_session_status_loop(self, status_stream):
-            print("CHECK STATUS")
             try:
                 for status in status_stream:
                     if status.status != issues.StatusCode.SUCCESS:
                         print("STATUS NOT SUCCESS")
                         self._state.reset(False)
             except Exception as e:
                 pass
-            print("CHECK STATUS STOP")
 
 
     def delete(self) -> None:
@@ -119,7 +130,20 @@ def create(self) -> None:
         self._create_call()
         self._attach()
 
-    def transaction(self, tx_mode: base.BaseQueryTxMode) -> base.IQueryTxContext:
+    def transaction(self, tx_mode: base.BaseQueryTxMode = None) -> base.IQueryTxContext:
         if not self._state.session_id:
             return
-        return BaseTxContext(tx_mode)
+        return BaseTxContext(
+            self._driver,
+            self._state,
+            self,
+            tx_mode,
+        )
+
+    def execute(self, query: str, commit_tx: bool = True):
+        stream_it = self._execute_call(query, commit_tx)
+
+        return _utilities.SyncResponseIterator(
+            stream_it,
+            lambda resp: base.wrap_execute_query_response(rpc_state=None, response_pb=resp),
+        )