basic pool & retries & example

vgvoleg · vgvoleg · commit 802c8414b734 · 2024-07-19T14:29:19.000+03:00
diff --git a/examples/query-service/basic_example.py b/examples/query-service/basic_example.py
@@ -1,7 +1,5 @@
 import ydb
 
-from ydb.query.session import QuerySessionSync
-
 
 def main():
     driver_config = ydb.DriverConfig(
@@ -16,14 +14,57 @@ def main():
     except TimeoutError:
         raise RuntimeError("Connect failed to YDB")
 
-    session = QuerySessionSync(driver)
-    session.create()
+    # client = ydb.QueryClientSync(driver)
+    # session = client.session().create()
+    pool = ydb.QuerySessionPool(driver)
+    # with pool.checkout() as session:
+    def callee(session):
+        print("="*50)
+        print("BEFORE ACTION")
+        it = session.execute("""SELECT COUNT(*) FROM example;""")
+        for result_set in it:
+            print(f"rows: {str(result_set.rows)}")
+
+        print("="*50)
+        print("INSERT WITH COMMIT TX")
+        tx = session.transaction()
+
+        tx.begin()
+
+        tx.execute("""INSERT INTO example (key, value) VALUES (0033, "onepieceisreal");""")
+
+        for result_set in tx.execute("""SELECT 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;"""):
+            print(f"rows: {str(result_set.rows)}")
+
+        print("="*50)
+        print("INSERT WITH ROLLBACK TX")
+
+        tx = session.transaction()
+
+        tx.begin()
+
+        tx.execute("""INSERT INTO example (key, value) VALUES (0044, "onepieceisreal");""")
+
+        for result_set in tx.execute("""SELECT 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;"""):
+            print(f"rows: {str(result_set.rows)}")
 
-    it = session.execute("select 1; select 2;", commit_tx=False)
-    for result_set in it:
-        # pass
-        print(f"columns: {str(result_set.columns)}")
-        print(f"rows: {str(result_set.rows)}")
+    pool.retry_operation_sync(callee)
 
 
 if __name__ == "__main__":
diff --git a/ydb/__init__.py b/ydb/__init__.py
@@ -19,6 +19,7 @@
 from .tracing import *  # noqa
 from .topic import *  # noqa
 from .draft import *  # noqa
+from .query import * # noqa
 
 try:
     import ydb.aio as aio  # noqa
diff --git a/ydb/_grpc/grpcwrapper/ydb_query.py b/ydb/_grpc/grpcwrapper/ydb_query.py
@@ -129,11 +129,11 @@ def from_proto(msg: ydb_query_pb2.RollbackTransactionResponse) -> "RollbackTrans
 @dataclass
 class QueryContent(IFromPublic, IToProto):
     text: str
-    syntax: Optional[str] = None
+    syntax: Optional[int] = None
 
     @staticmethod
-    def from_public(query: str) -> "QueryContent":
-        return QueryContent(text=query)
+    def from_public(query: str, syntax: int = None) -> "QueryContent":
+        return QueryContent(text=query, syntax=syntax)
 
     def to_proto(self) -> ydb_query_pb2.QueryContent:
         return ydb_query_pb2.QueryContent(text=self.text, syntax=self.syntax)
@@ -163,16 +163,16 @@ class ExecuteQueryRequest(IToProto):
     query_content: QueryContent
     tx_control: TransactionControl
     concurrent_result_sets: Optional[bool] = False
-    exec_mode: Optional[str] = None
+    exec_mode: Optional[int] = None
     parameters: Optional[dict] = None
-    stats_mode: Optional[str] = None
+    stats_mode: Optional[int] = 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,
+            exec_mode=self.exec_mode,
             stats_mode=self.stats_mode,
             concurrent_result_sets=self.concurrent_result_sets,
             parameters=self.parameters,
diff --git a/ydb/query/__init__.py b/ydb/query/__init__.py
@@ -0,0 +1,25 @@
+from .base import (
+    IQueryClient,
+    SupportedDriverType,
+    QueryClientSettings,
+)
+
+from .session import QuerySessionSync
+
+from .._grpc.grpcwrapper.ydb_query_public_types import (
+    QueryOnlineReadOnly,
+    QuerySerializableReadWrite,
+    QuerySnapshotReadOnly,
+    QueryStaleReadOnly,
+)
+
+from .pool import QuerySessionPool
+
+
+class QueryClientSync(IQueryClient):
+    def __init__(self, driver: SupportedDriverType, query_client_settings: QueryClientSettings = None):
+        self._driver = driver
+        self._settings = query_client_settings
+
+    def session(self) -> QuerySessionSync:
+        return QuerySessionSync(self._driver, self._settings)
diff --git a/ydb/query/base.py b/ydb/query/base.py
@@ -1,4 +1,5 @@
 import abc
+import enum
 import functools
 
 from typing import (
@@ -63,7 +64,7 @@ def __init__(self, driver: SupportedDriverType, settings: QueryClientSettings =
         pass
 
     @abc.abstractmethod
-    def create(self) -> None:
+    def create(self) -> "IQuerySession":
         pass
 
     @abc.abstractmethod
@@ -117,7 +118,7 @@ def rollback(settings: QueryClientSettings = None):
         pass
 
     @abc.abstractmethod
-    def execute(query: str):
+    def execute(query: str, commit_tx=False):
         pass
 
 
@@ -130,39 +131,73 @@ def session(self) -> IQuerySession:
         pass
 
 
+class QuerySyntax(enum.IntEnum):
+    UNSPECIFIED = 0
+    YQL_V1 = 1
+    PG = 2
+
+
+class QueryExecMode(enum.IntEnum):
+    UNSPECIFIED = 0
+    PARSE = 10
+    VALIDATE = 20
+    EXPLAIN = 30
+    EXECUTE = 50
+
+
 def create_execute_query_request(
-    query: str, session_id: str, tx_id: str = None, commit_tx: bool = False, tx_mode: BaseQueryTxMode = None
+    query: str,
+    session_id: str,
+    tx_id: str = None,
+    commit_tx: bool = False,
+    tx_mode: BaseQueryTxMode = None,
+    syntax: QuerySyntax = None,
+    exec_mode: QueryExecMode = None,
+    parameters: dict = None,
+    concurrent_result_sets: 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,
         )
     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,
                 ),
                 commit_tx=commit_tx,
             ),
+            exec_mode=exec_mode,
+            parameters=parameters,
+            concurrent_result_sets=concurrent_result_sets,
         )
 
     return req.to_proto()
 
 
 def wrap_execute_query_response(rpc_state, response_pb):
+    issues._process_response(response_pb)
     return convert.ResultSet.from_message(response_pb.result_set)
 
 
diff --git a/ydb/query/pool.py b/ydb/query/pool.py
@@ -0,0 +1,132 @@
+import abc
+import time
+from typing import Callable
+
+from . import base
+from .session import (
+    QuerySessionSync,
+    BaseQuerySession,
+)
+from .. import issues
+from .._errors import check_retriable_error
+
+
+class RetrySettings(object):
+    def __init__(
+        self,
+        max_retries: int = 10,
+        max_session_acquire_timeout: int = None,
+        on_ydb_error_callback: Callable = None,
+        idempotent: bool = False,
+    ):
+        self.max_retries = max_retries
+        self.max_session_acquire_timeout = max_session_acquire_timeout
+        self.on_ydb_error_callback = (lambda e: None) if on_ydb_error_callback is None else on_ydb_error_callback
+        self.retry_not_found = True
+        self.idempotent = idempotent
+        self.retry_internal_error = True
+        self.unknown_error_handler = lambda e: None
+
+
+class YdbRetryOperationSleepOpt:
+    def __init__(self, timeout):
+        self.timeout = timeout
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.timeout == other.timeout
+
+    def __repr__(self):
+        return "YdbRetryOperationSleepOpt(%s)" % self.timeout
+
+
+class YdbRetryOperationFinalResult:
+    def __init__(self, result):
+        self.result = result
+        self.exc = None
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.result == other.result and self.exc == other.exc
+
+    def __repr__(self):
+        return "YdbRetryOperationFinalResult(%s, exc=%s)" % (self.result, self.exc)
+
+    def set_exception(self, exc):
+        self.exc = exc
+
+
+def retry_operation_impl(callee: Callable, retry_settings: RetrySettings = None, *args, **kwargs):
+    retry_settings = RetrySettings() if retry_settings is None else retry_settings
+    status = None
+
+    for attempt in range(retry_settings.max_retries + 1):
+        try:
+            result = YdbRetryOperationFinalResult(callee(*args, **kwargs))
+            yield result
+
+            if result.exc is not None:
+                raise result.exc
+
+        except issues.Error as e:
+            status = e
+            retry_settings.on_ydb_error_callback(e)
+
+            retriable_info = check_retriable_error(e, retry_settings, attempt)
+            if not retriable_info.is_retriable:
+                raise
+
+            skip_yield_error_types = [
+                issues.Aborted,
+                issues.BadSession,
+                issues.NotFound,
+                issues.InternalError,
+            ]
+
+            yield_sleep = True
+            for t in skip_yield_error_types:
+                if isinstance(e, t):
+                    yield_sleep = False
+
+            if yield_sleep:
+                yield YdbRetryOperationSleepOpt(retriable_info.sleep_timeout_seconds)
+
+        except Exception as e:
+            # you should provide your own handler you want
+            retry_settings.unknown_error_handler(e)
+            raise
+
+    raise status
+
+
+class QuerySessionPool:
+    def __init__(self, driver: base.SupportedDriverType):
+        self._driver = driver
+
+    def checkout(self):
+        return SimpleQuerySessionCheckout(self)
+
+    def retry_operation_sync(self, callee: Callable, retry_settings: RetrySettings = None, *args, **kwargs):
+        retry_settings = RetrySettings() if retry_settings is None else retry_settings
+
+        def wrapped_callee():
+            with self.checkout() as session:
+                return callee(session, *args, **kwargs)
+
+        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):
+        self._pool = pool
+        self._session = QuerySessionSync(pool._driver)
+
+    def __enter__(self):
+        self._session.create()
+        return self._session
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._session.delete()
diff --git a/ydb/query/session.py b/ydb/query/session.py
@@ -205,14 +205,16 @@ def delete(self) -> None:
         self._delete_call()
         self._stream.cancel()
 
-    def create(self) -> None:
+    def create(self) -> "QuerySessionSync":
         if self._state._already_in(QuerySessionStateEnum.CREATED):
             return
 
         self._state._check_invalid_transition(QuerySessionStateEnum.CREATED)
         self._create_call()
         self._attach()
 
+        return self
+
     def transaction(self, tx_mode: base.BaseQueryTxMode = None) -> base.IQueryTxContext:
         self._state._check_session_ready_to_use()
 
diff --git a/ydb/query/transaction.py b/ydb/query/transaction.py
