new tests and connection fixes

.github/workflows/lint.yml

@@ -0,0 +1,47 @@
+name: Lint
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    concurrency:
+      group: unit-${{ github.ref }}-${{ matrix.python-version }}-${{ matrix.poetry-version }}
+      cancel-in-progress: true
+
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.9]
+        poetry-version: [1.8]
+
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          persist-credentials: false
+          fetch-depth: 0
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install Poetry ${{ matrix.poetry-version }}
+        run: |
+          pip install "poetry~=${{ matrix.poetry-version }}.0"
+
+          # Ensure that Poetry is not upgraded past the version we are testing
+          poetry add "poetry@~${{ matrix.poetry-version }}" --lock
+
+      - name: Install packages
+        run: |
+          poetry install
+
+      - name: Run tests
+        run: |
+          poetry run poe lint

tests/test_connection.py

@@ -0,0 +1,172 @@
+from contextlib import suppress
+
+import pytest
+import pytest_asyncio
+import ydb
+
+import ydb_dbapi as dbapi
+
+
+class BaseDBApiTestSuit:
+    async def _test_isolation_level_read_only(
+        self,
+        connection: dbapi.Connection,
+        isolation_level: str,
+        read_only: bool,
+    ):
+        await connection.cursor().execute(
+            "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
+        )
+        connection.set_isolation_level(isolation_level)
+
+        cursor = connection.cursor()
+
+        await connection.begin()
+
+        query = "UPSERT INTO foo(id) VALUES (1)"
+        if read_only:
+            with pytest.raises(dbapi.DatabaseError):
+                await cursor.execute(query)
+        else:
+            await cursor.execute(query)
+
+        await connection.rollback()
+
+        await connection.cursor().execute("DROP TABLE foo")
+        await connection.cursor().close()
+
+    async def _test_connection(self, connection: dbapi.Connection):
+        await connection.commit()
+        await connection.rollback()
+
+        cur = connection.cursor()
+        with suppress(dbapi.DatabaseError):
+            await cur.execute("DROP TABLE foo")
+
+        assert not await connection.check_exists("/local/foo")
+        with pytest.raises(dbapi.ProgrammingError):
+            await connection.describe("/local/foo")
+
+        await cur.execute(
+            "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
+        )
+
+        assert await connection.check_exists("/local/foo")
+
+        col = (await connection.describe("/local/foo")).columns[0]
+        assert col.name == "id"
+        assert col.type == ydb.PrimitiveType.Int64
+
+        await cur.execute("DROP TABLE foo")
+        await cur.close()
+
+    async def _test_cursor_raw_query(self, connection: dbapi.Connection):
+        cur = connection.cursor()
+        assert cur
+
+        with suppress(dbapi.DatabaseError):
+            await cur.execute("DROP TABLE test")
+
+        await cur.execute(
+            "CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
+        )
+
+        await cur.execute(
+            """
+            DECLARE $data AS List<Struct<id:Int64, text: Utf8>>;
+
+            INSERT INTO test SELECT id, text FROM AS_TABLE($data);
+            """,
+            {
+                "$data": ydb.TypedValue(
+                    [
+                        {"id": 17, "text": "seventeen"},
+                        {"id": 21, "text": "twenty one"},
+                    ],
+                    ydb.ListType(
+                        ydb.StructType()
+                        .add_member("id", ydb.PrimitiveType.Int64)
+                        .add_member("text", ydb.PrimitiveType.Utf8)
+                    ),
+                )
+            },
+        )
+
+        await cur.execute("DROP TABLE test")
+
+        await cur.close()
+
+    async def _test_errors(self, connection: dbapi.Connection):
+        with pytest.raises(dbapi.InterfaceError):
+            await dbapi.connect("localhost:2136", database="/local666")
+
+        cur = connection.cursor()
+
+        with suppress(dbapi.DatabaseError):
+            await cur.execute("DROP TABLE test")
+
+        with pytest.raises(dbapi.DataError):
+            await cur.execute("SELECT 18446744073709551616")
+
+        with pytest.raises(dbapi.DataError):
+            await cur.execute("SELECT * FROM 拉屎")
+
+        with pytest.raises(dbapi.DataError):
+            await cur.execute("SELECT floor(5 / 2)")
+
+        with pytest.raises(dbapi.ProgrammingError):
+            await cur.execute("SELECT * FROM test")
+
+        await cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
+
+        await cur.execute("INSERT INTO test(id) VALUES(1)")
+        with pytest.raises(dbapi.IntegrityError):
+            await cur.execute("INSERT INTO test(id) VALUES(1)")
+
+        await cur.execute("DROP TABLE test")
+        await cur.close()
+
+
+class TestAsyncConnection(BaseDBApiTestSuit):
+    @pytest_asyncio.fixture
+    async def connection(self, endpoint, database):
+        host, port = endpoint.split(":")
+        conn = await dbapi.connect(host=host, port=port, database=database)
+        try:
+            yield conn
+        finally:
+            await conn.close()
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "isolation_level, read_only",
+        [
+            (dbapi.IsolationLevel.SERIALIZABLE, False),
+            (dbapi.IsolationLevel.AUTOCOMMIT, False),
+            # (dbapi.IsolationLevel.ONLINE_READONLY, True),
+            # (dbapi.IsolationLevel.ONLINE_READONLY_INCONSISTENT, True),
+            # (dbapi.IsolationLevel.STALE_READONLY, True),
+            # (dbapi.IsolationLevel.SNAPSHOT_READONLY, True),
+        ],
+    )
+    async def test_isolation_level_read_only(
+        self,
+        isolation_level: str,
+        read_only: bool,
+        connection: dbapi.Connection,
+    ):
+        await self._test_isolation_level_read_only(
+            connection, isolation_level, read_only
+        )
+
+    @pytest.mark.asyncio
+    async def test_connection(self, connection: dbapi.Connection):
+        await self._test_connection(connection)
+
+    @pytest.mark.asyncio
+    async def test_cursor_raw_query(self, connection: dbapi.Connection):
+        await self._test_cursor_raw_query(connection)
+
+    @pytest.mark.asyncio
+    async def test_errors(self, connection: dbapi.Connection):
+        await self._test_errors(connection)

tests/test_cursor.py

@@ -1,11 +1,10 @@
 import pytest
 import ydb_dbapi
-import ydb_dbapi.cursors
 
 
 @pytest.mark.asyncio
 async def test_cursor_ddl(session_pool):
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
 
     yql = """
         CREATE TABLE table (
@@ -29,7 +28,7 @@ async def test_cursor_ddl(session_pool):
 
 @pytest.mark.asyncio
 async def test_cursor_dml(session_pool):
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -40,7 +39,7 @@ async def test_cursor_dml(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
 
     yql_text = """
     SELECT COUNT(*) FROM table as sum
@@ -55,7 +54,7 @@ async def test_cursor_dml(session_pool):
 
 @pytest.mark.asyncio
 async def test_cursor_fetch_one(session_pool):
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -65,7 +64,7 @@ async def test_cursor_fetch_one(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
 
     yql_text = """
     SELECT id, val FROM table
@@ -84,7 +83,7 @@ async def test_cursor_fetch_one(session_pool):
 
 @pytest.mark.asyncio
 async def test_cursor_fetch_many(session_pool):
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
     yql_text = """
     INSERT INTO table (id, val) VALUES
     (1, 1),
@@ -96,7 +95,7 @@ async def test_cursor_fetch_many(session_pool):
     await cursor.execute(query=yql_text)
     assert await cursor.fetchone() is None
 
-    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
 
     yql_text = """
     SELECT id, val FROM table
@@ -122,4 +121,58 @@ async def test_cursor_fetch_many(session_pool):
 
 @pytest.mark.asyncio
 async def test_cursor_fetch_all(session_pool):
-    pass
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+    yql_text = """
+    INSERT INTO table (id, val) VALUES
+    (1, 1),
+    (2, 2),
+    (3, 3)
+    """
+
+    await cursor.execute(query=yql_text)
+    assert await cursor.fetchone() is None
+
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+
+    yql_text = """
+    SELECT id, val FROM table
+    """
+
+    await cursor.execute(query=yql_text)
+
+    assert cursor.rowcount == 3
+
+    res = await cursor.fetchall()
+    assert len(res) == 3
+    assert res[0][0] == 1
+    assert res[1][0] == 2
+    assert res[2][0] == 3
+
+    assert await cursor.fetchall() is None
+
+
+@pytest.mark.asyncio
+async def test_cursor_next_set(session_pool):
+    cursor = ydb_dbapi.Cursor(session_pool=session_pool)
+    yql_text = """SELECT 1 as val; SELECT 2 as val;"""
+
+    await cursor.execute(query=yql_text)
+
+    res = await cursor.fetchall()
+    assert len(res) == 1
+    assert res[0][0] == 1
+
+    nextset = await cursor.nextset()
+    assert nextset
+
+    res = await cursor.fetchall()
+    assert len(res) == 1
+    assert res[0][0] == 2
+
+    nextset = await cursor.nextset()
+    assert nextset
+
+    assert await cursor.fetchall() is None
+
+    nextset = await cursor.nextset()
+    assert not nextset

ydb_dbapi/__init__.py

@@ -1 +1,3 @@
 from .errors import *  # noqa
+from .connection import Connection, IsolationLevel, connect  # noqa
+from .cursors import Cursor  # noqa