Query service client support

docker-compose.yml

@@ -1,7 +1,7 @@
 version: "3.3"
 services:
   ydb:
-    image: cr.yandex/yc/yandex-docker-local-ydb:latest
+    image: ydbplatform/local-ydb:trunk
     restart: always
     ports:
       - 2136:2136

examples/query-service/basic_example.py

@@ -0,0 +1,87 @@
+import ydb
+
+
+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")
+
+    pool = ydb.QuerySessionPool(driver)
+
+    print("=" * 50)
+    print("DELETE TABLE IF EXISTS")
+    pool.execute_with_retries("drop table if exists example")
+
+    print("=" * 50)
+    print("CREATE TABLE")
+    pool.execute_with_retries("CREATE TABLE example(key UInt64, value String, PRIMARY KEY (key))")
+
+    pool.execute_with_retries("INSERT INTO example (key, value) VALUES (1, 'onepieceisreal')")
+
+    def callee(session):
+        print("=" * 50)
+        with session.execute("delete from example"):
+            pass
+
+        print("BEFORE ACTION")
+        with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+            for result_set in results:
+                print(f"rows: {str(result_set.rows)}")
+
+        print("=" * 50)
+        print("INSERT WITH COMMIT TX")
+
+        with session.transaction() as tx:
+            tx.begin()
+
+            with tx.execute("INSERT INTO example (key, value) VALUES (1, 'onepieceisreal')"):
+                pass
+
+            with tx.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+                for result_set in results:
+                    print(f"rows: {str(result_set.rows)}")
+
+            tx.commit()
+
+        print("=" * 50)
+        print("AFTER COMMIT TX")
+
+        with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+            for result_set in results:
+                print(f"rows: {str(result_set.rows)}")
+
+        print("=" * 50)
+        print("INSERT WITH ROLLBACK TX")
+
+        with session.transaction() as tx:
+            tx.begin()
+
+            with tx.execute("INSERT INTO example (key, value) VALUES (2, 'onepieceisreal')"):
+                pass
+
+            with tx.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+                for result_set in results:
+                    print(f"rows: {str(result_set.rows)}")
+
+            tx.rollback()
+
+        print("=" * 50)
+        print("AFTER ROLLBACK TX")
+
+        with session.execute("SELECT COUNT(*) as rows_count FROM example") as results:
+            for result_set in results:
+                print(f"rows: {str(result_set.rows)}")
+
+    pool.retry_operation_sync(callee)
+
+
+if __name__ == "__main__":
+    main()

tests/query/conftest.py

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

tests/query/test_query_session.py

@@ -0,0 +1,91 @@
+import pytest
+
+from ydb.query.session import QuerySessionSync
+
+
+def _check_session_state_empty(session: QuerySessionSync):
+    assert session._state.session_id is None
+    assert session._state.node_id is None
+    assert not session._state.attached
+
+
+def _check_session_state_full(session: QuerySessionSync):
+    assert session._state.session_id is not None
+    assert session._state.node_id is not None
+    assert session._state.attached
+
+
+class TestQuerySession:
+    def test_session_normal_lifecycle(self, session: QuerySessionSync):
+        _check_session_state_empty(session)
+
+        session.create()
+        _check_session_state_full(session)
+
+        session.delete()
+        _check_session_state_empty(session)
+
+    def test_second_create_do_nothing(self, session: QuerySessionSync):
+        session.create()
+        _check_session_state_full(session)
+
+        session_id_before = session._state.session_id
+        node_id_before = session._state.node_id
+
+        session.create()
+        _check_session_state_full(session)
+
+        assert session._state.session_id == session_id_before
+        assert session._state.node_id == node_id_before
+
+    def test_second_delete_do_nothing(self, session: QuerySessionSync):
+        session.create()
+
+        session.delete()
+        session.delete()
+
+    def test_delete_before_create_not_possible(self, session: QuerySessionSync):
+        with pytest.raises(RuntimeError):
+            session.delete()
+
+    def test_create_after_delete_not_possible(self, session: QuerySessionSync):
+        session.create()
+        session.delete()
+        with pytest.raises(RuntimeError):
+            session.create()
+
+    def test_transaction_before_create_raises(self, session: QuerySessionSync):
+        with pytest.raises(RuntimeError):
+            session.transaction()
+
+    def test_transaction_after_delete_raises(self, session: QuerySessionSync):
+        session.create()
+
+        session.delete()
+
+        with pytest.raises(RuntimeError):
+            session.transaction()
+
+    def test_transaction_after_create_not_raises(self, session: QuerySessionSync):
+        session.create()
+        session.transaction()
+
+    def test_execute_before_create_raises(self, session: QuerySessionSync):
+        with pytest.raises(RuntimeError):
+            session.execute("select 1;")
+
+    def test_execute_after_delete_raises(self, session: QuerySessionSync):
+        session.create()
+        session.delete()
+        with pytest.raises(RuntimeError):
+            session.execute("select 1;")
+
+    def test_basic_execute(self, session: QuerySessionSync):
+        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]

tests/query/test_query_session_pool.py

@@ -0,0 +1,49 @@
+import pytest
+import ydb
+from ydb.query.pool import QuerySessionPool
+from ydb.query.session import QuerySessionSync, QuerySessionStateEnum
+
+
+class TestQuerySessionPool:
+    def test_checkout_provides_created_session(self, pool: QuerySessionPool):
+        with pool.checkout() as session:
+            assert session._state._state == QuerySessionStateEnum.CREATED
+
+        assert session._state._state == QuerySessionStateEnum.CLOSED
+
+    def test_oneshot_query_normal(self, pool: QuerySessionPool):
+        res = pool.execute_with_retries("select 1;")
+        assert len(res) == 1
+
+    def test_oneshot_ddl_query(self, pool: QuerySessionPool):
+        pool.execute_with_retries("create table Queen(key UInt64, PRIMARY KEY (key));")
+        pool.execute_with_retries("drop table Queen;")
+
+    def test_oneshot_query_raises(self, pool: QuerySessionPool):
+        with pytest.raises(ydb.GenericError):
+            pool.execute_with_retries("Is this the real life? Is this just fantasy?")
+
+    def test_retry_op_uses_created_session(self, pool: QuerySessionPool):
+        def callee(session: QuerySessionSync):
+            assert session._state._state == QuerySessionStateEnum.CREATED
+
+        pool.retry_operation_sync(callee)
+
+    def test_retry_op_normal(self, pool: QuerySessionPool):
+        def callee(session: QuerySessionSync):
+            with session.transaction() as tx:
+                iterator = tx.execute("select 1;", commit_tx=True)
+                return [result_set for result_set in iterator]
+
+        res = pool.retry_operation_sync(callee)
+        assert len(res) == 1
+
+    def test_retry_op_raises(self, pool: QuerySessionPool):
+        class CustomException(Exception):
+            pass
+
+        def callee(session: QuerySessionSync):
+            raise CustomException()
+
+        with pytest.raises(CustomException):
+            pool.retry_operation_sync(callee)

tests/query/test_query_transaction.py

@@ -0,0 +1,81 @@
+import pytest
+
+from ydb.query.transaction import BaseTxContext
+from ydb.query.transaction import QueryTxStateEnum
+
+
+class TestQueryTransaction:
+    def test_tx_begin(self, tx: BaseTxContext):
+        assert tx.tx_id is None
+
+        tx.begin()
+        assert tx.tx_id is not None
+
+    def test_tx_allow_double_commit(self, tx: BaseTxContext):
+        tx.begin()
+        tx.commit()
+        tx.commit()
+
+    def test_tx_allow_double_rollback(self, tx: BaseTxContext):
+        tx.begin()
+        tx.rollback()
+        tx.rollback()
+
+    def test_tx_commit_before_begin(self, tx: BaseTxContext):
+        tx.commit()
+        assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
+
+    def test_tx_rollback_before_begin(self, tx: BaseTxContext):
+        tx.rollback()
+        assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
+
+    def test_tx_first_execute_begins_tx(self, tx: BaseTxContext):
+        tx.execute("select 1;")
+        tx.commit()
+
+    def test_interactive_tx_commit(self, tx: BaseTxContext):
+        tx.execute("select 1;", commit_tx=True)
+        with pytest.raises(RuntimeError):
+            tx.execute("select 1;")
+
+    def test_tx_execute_raises_after_commit(self, tx: BaseTxContext):
+        tx.begin()
+        tx.commit()
+        with pytest.raises(RuntimeError):
+            tx.execute("select 1;")
+
+    def test_tx_execute_raises_after_rollback(self, tx: BaseTxContext):
+        tx.begin()
+        tx.rollback()
+        with pytest.raises(RuntimeError):
+            tx.execute("select 1;")
+
+    def test_context_manager_rollbacks_tx(self, tx: BaseTxContext):
+        with tx:
+            tx.begin()
+
+        assert tx._tx_state._state == QueryTxStateEnum.ROLLBACKED
+
+    def test_context_manager_normal_flow(self, tx: BaseTxContext):
+        with tx:
+            tx.begin()
+            tx.execute("select 1;")
+            tx.commit()
+
+        assert tx._tx_state._state == QueryTxStateEnum.COMMITTED
+
+    def test_context_manager_does_not_hide_exceptions(self, tx: BaseTxContext):
+        class CustomException(Exception):
+            pass
+
+        with pytest.raises(CustomException):
+            with tx:
+                raise CustomException()
+
+    def test_execute_as_context_manager(self, tx: BaseTxContext):
+        tx.begin()
+
+        with tx.execute("select 1;") as results:
+            res = [result_set for result_set in results]
+
+        assert len(res) == 1

ydb/__init__.py

@@ -19,6 +19,8 @@
 from .tracing import *  # noqa
 from .topic import *  # noqa
 from .draft import *  # noqa
+from .query import *  # noqa
+from .retries import *  # noqa
 
 try:
     import ydb.aio as aio  # noqa