refactor session

Author: vgvoleg

tests/query/conftest.py

@@ -0,0 +1,27 @@
+import pytest
+from ydb.query.session import QuerySessionSync
+
+
+@pytest.fixture
+def session(driver_sync):
+    session = QuerySessionSync(driver_sync)
+
+    yield session
+
+    try:
+        session.delete()
+    except BaseException:
+        pass
+
+@pytest.fixture
+def transaction(session):
+    session.create()
+    transaction = session.transaction()
+
+    yield transaction
+
+    try:
+        transaction.rollback()
+    except BaseException:
+        pass
+

tests/query/test_query_session.py

@@ -1,7 +1,5 @@
 import pytest
 
-import ydb.query.session
-
 def _check_session_state_empty(session):
     assert session._state.session_id == None
     assert session._state.node_id == None
@@ -13,8 +11,7 @@ def _check_session_state_full(session):
     assert session._state.attached == True
 
 class TestQuerySession:
-    def test_session_normal_lifecycle(self, driver_sync):
-        session = ydb.query.session.QuerySessionSync(driver_sync)
+    def test_session_normal_lifecycle(self, session):
         _check_session_state_empty(session)
 
         session.create()
@@ -23,8 +20,7 @@ def test_session_normal_lifecycle(self, driver_sync):
         session.delete()
         _check_session_state_empty(session)
 
-    def test_second_create_do_nothing(self, driver_sync):
-        session = ydb.query.session.QuerySessionSync(driver_sync)
+    def test_second_create_do_nothing(self, session):
         session.create()
         _check_session_state_full(session)
 
@@ -36,3 +32,55 @@ def test_second_create_do_nothing(self, driver_sync):
 
         assert session._state.session_id == session_id_before
         assert session._state.node_id == node_id_before
+
+    def test_second_delete_do_nothing(self, session):
+        session.create()
+
+        session.delete()
+        session.delete()
+
+    def test_delete_before_create_not_possible(self, session):
+        with pytest.raises(RuntimeError):
+            session.delete()
+
+    def test_create_after_delete_not_possible(self, session):
+        session.create()
+        session.delete()
+        with pytest.raises(RuntimeError):
+            session.create()
+
+    def test_transaction_before_create_raises(self, session):
+        with pytest.raises(RuntimeError):
+            session.transaction()
+
+    def test_transaction_after_delete_raises(self, session):
+        session.create()
+
+        session.delete()
+
+        with pytest.raises(RuntimeError):
+            session.transaction()
+
+    def test_transaction_after_create_not_raises(self, session):
+        session.create()
+        session.transaction()
+
+    def test_execute_before_create_raises(self, session):
+        with pytest.raises(RuntimeError):
+            session.execute("select 1;")
+
+    def test_execute_after_delete_raises(self, session):
+        session.create()
+        session.delete()
+        with pytest.raises(RuntimeError):
+            session.execute("select 1;")
+
+    def test_basic_execute(self, session):
+        session.create()
+        it = session.execute("select 1;")
+        result_sets = [result_set for result_set in it]
+
+        assert len(result_sets) == 1
+        assert len(result_sets[0].rows) == 1
+        assert len(result_sets[0].columns) == 1
+        assert list(result_sets[0].rows[0].values()) == [1]

ydb/query/base.py

@@ -20,83 +20,40 @@
 class QueryClientSettings: ...
 
 
-class QuerySessionStateEnum(enum.Enum):
-    NOT_INITIALIZED = "NOT_INITIALIZED"
-    CREATED = "CREATED"
-    CLOSED = "CLOSED"
-
-
-class QuerySessionStateHelper(abc.ABC):
-    _VALID_TRANSITIONS = {
-        QuerySessionStateEnum.NOT_INITIALIZED: [QuerySessionStateEnum.CREATED],
-        QuerySessionStateEnum.CREATED: [QuerySessionStateEnum.CLOSED],
-        QuerySessionStateEnum.CLOSED: []
-    }
-
-    @classmethod
-    def valid_transition(cls, before: QuerySessionStateEnum, after: QuerySessionStateEnum) -> bool:
-        return after in cls._VALID_TRANSITIONS[before]
-
-
-class QueryTxStateEnum(enum.Enum):
-    NOT_INITIALIZED = "NOT_INITIALIZED"
-    BEGINED = "BEGINED"
-    COMMITTED = "COMMITTED"
-    ROLLBACKED = "ROLLBACKED"
-    DEAD = "DEAD"
-
-
-class QueryTxStateHelper(abc.ABC):
-    _VALID_TRANSITIONS = {
-        QueryTxStateEnum.NOT_INITIALIZED: [QueryTxStateEnum.BEGINED, QueryTxStateEnum.DEAD],
-        QueryTxStateEnum.BEGINED: [QueryTxStateEnum.COMMITTED, QueryTxStateEnum.ROLLBACKED, QueryTxStateEnum.DEAD],
-        QueryTxStateEnum.COMMITTED: [],
-        QueryTxStateEnum.ROLLBACKED: [],
-        QueryTxStateEnum.DEAD: [],
-    }
-
-    @classmethod
-    def valid_transition(cls, before: QueryTxStateEnum, after: QueryTxStateEnum) -> bool:
-        return after in cls._VALID_TRANSITIONS[before]
-
-
-class QuerySessionState:
-    _session_id: Optional[str]
-    _node_id: Optional[int]
-    _attached: bool = False
-    _settings: Optional[QueryClientSettings]
-
+class IQuerySessionState(abc.ABC):
     def __init__(self, settings: QueryClientSettings = None):
-        self._settings = settings
-        self.reset()
+        pass
 
+    @abc.abstractmethod
     def reset(self) -> None:
-        self._session_id = None
-        self._node_id = None
-        self._attached = False
+        pass
 
     @property
+    @abc.abstractmethod
     def session_id(self) -> Optional[str]:
-        return self._session_id
+        pass
 
-    def set_session_id(self, session_id: str) -> "QuerySessionState":
-        self._session_id = session_id
-        return self
+    @abc.abstractmethod
+    def set_session_id(self, session_id: str) -> "IQuerySessionState":
+        pass
 
     @property
+    @abc.abstractmethod
     def node_id(self) -> Optional[int]:
-        return self._node_id
+        pass
 
-    def set_node_id(self, node_id: int) -> "QuerySessionState":
-        self._node_id = node_id
-        return self
+    @abc.abstractmethod
+    def set_node_id(self, node_id: int) -> "IQuerySessionState":
+        pass
 
     @property
+    @abc.abstractmethod
     def attached(self) -> bool:
-        return self._attached
+        pass
 
-    def set_attached(self, attached: bool) -> None:
-        self._attached = attached
+    @abc.abstractmethod
+    def set_attached(self, attached: bool) -> "IQuerySessionState":
+        pass
 
 
 class IQuerySession(abc.ABC):
@@ -120,7 +77,7 @@ def transaction(self, tx_mode: BaseQueryTxMode) -> "IQueryTxContext":
 class IQueryTxContext(abc.ABC):
 
     @abc.abstractmethod
-    def __init__(self, driver: SupportedDriverType, session_state: QuerySessionState, session: IQuerySession, tx_mode: BaseQueryTxMode = None):
+    def __init__(self, driver: SupportedDriverType, session_state: IQuerySessionState, session: IQuerySession, tx_mode: BaseQueryTxMode = None):
         pass
 
     @abc.abstractmethod
@@ -184,17 +141,17 @@ def create_execute_query_request(query: str, session_id: str, commit_tx: bool):
     return req.to_proto()
 
 def wrap_execute_query_response(rpc_state, response_pb):
-
     return convert.ResultSet.from_message(response_pb.result_set)
 
+
 X_YDB_SERVER_HINTS = "x-ydb-server-hints"
 X_YDB_SESSION_CLOSE = "session-close"
 
 
 def _check_session_is_closing(rpc_state, session_state):
     metadata = rpc_state.trailing_metadata()
     if X_YDB_SESSION_CLOSE in metadata.get(X_YDB_SERVER_HINTS, []):
-        session_state.set_closing()
+        session_state.set_closing() # TODO: clarify & implement
 
 
 def bad_session_handler(func):