switch back to QuerySessionPool

Author: blinkov

.github/workflows/run-tests.yml

@@ -13,5 +13,13 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.11'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements-dev.txt
       - name: Run tests
         run: make test

pytest.ini

@@ -3,7 +3,6 @@ testpaths = tests
 python_files = test_*.py
 python_classes = Test*
 python_functions = test_*
-addopts = --cov=ydb_mcp --cov-report=term-missing --cov-report=xml --cov-report=html --no-cov-on-fail
 
 markers =
     unit: mark a test as a unit test
@@ -26,4 +25,6 @@ filterwarnings =
     ignore::RuntimeWarning:ydb_mcp.patches:
     ignore:Task was destroyed but it is pending:RuntimeWarning:asyncio.base_events
     ignore:Error handling discovery task:RuntimeWarning:tests.integration.conftest
-    ignore:Error stopping driver:RuntimeWarning:tests.integration.conftest
+    ignore:Error stopping driver:RuntimeWarning:tests.integration.conftest
+
+addopts = --cov=ydb_mcp --cov-report=term-missing --cov-report=xml --cov-report=html --no-cov-on-fail

tests/integration/conftest.py

@@ -42,9 +42,6 @@
 asyncio_logger = logging.getLogger("asyncio")
 asyncio_logger.setLevel(logging.ERROR)
 
-# Use pytest-asyncio's built-in event_loop fixture
-pytestmark = pytest.mark.asyncio(scope="session")
-
 
 async def cleanup_pending_tasks():
     """Clean up any pending tasks in the current event loop."""

tests/integration/test_authentication_integration.py

@@ -26,7 +26,7 @@
 logger = logging.getLogger(__name__)
 
 # Use loop_scope instead of scope for the asyncio marker
-pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")]
+pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
 
 
 async def test_login_password_authentication(mcp_server):

tests/integration/test_directory_operations.py

@@ -22,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 # Mark these tests as integration and asyncio tests
-pytestmark = [pytest.mark.integration, pytest.mark.asyncio(scope="session")]
+pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
 
 
 def parse_text_content(response):

tests/integration/test_mcp_server_integration.py

@@ -30,7 +30,7 @@
 logger = logging.getLogger(__name__)
 
 # Use loop_scope instead of scope for the asyncio marker
-pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")]
+pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
 
 
 async def test_simple_query(mcp_server):

tests/integration/test_table_operations.py

@@ -17,7 +17,7 @@
 logger = logging.getLogger(__name__)
 
 # Mark these tests as integration and asyncio tests
-pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")]
+pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
 
 # We'll use the root directory which should already exist
 ROOT_DIR = "/"

tests/test_connection.py

@@ -50,7 +50,7 @@ async def test_connect(self, mock_driver_class):
         )
         mock_driver_class.return_value = mock_driver
 
-        with patch("ydb.SessionPool") as mock_session_pool_class:
+        with patch("ydb.aio.QuerySessionPool") as mock_session_pool_class:
             # Setup session pool mock
             mock_session_pool = MagicMock()
             mock_session_pool_class.return_value = mock_session_pool
@@ -89,7 +89,7 @@ async def test_connect_with_database_in_endpoint(self, mock_driver_class):
         )
         mock_driver_class.return_value = mock_driver
 
-        with patch("ydb.SessionPool") as mock_session_pool_class:
+        with patch("ydb.aio.QuerySessionPool") as mock_session_pool_class:
             # Setup session pool mock
             mock_session_pool = MagicMock()
             mock_session_pool_class.return_value = mock_session_pool
@@ -130,7 +130,7 @@ async def test_connect_with_explicit_database(self, mock_driver_class):
         )
         mock_driver_class.return_value = mock_driver
 
-        with patch("ydb.SessionPool") as mock_session_pool_class:
+        with patch("ydb.aio.QuerySessionPool") as mock_session_pool_class:
             # Setup session pool mock
             mock_session_pool = MagicMock()
             mock_session_pool_class.return_value = mock_session_pool

ydb_mcp/connection.py

@@ -5,6 +5,7 @@
 from urllib.parse import urlparse
 
 import ydb
+from ydb.aio import QuerySessionPool
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +22,7 @@ def __init__(self, connection_string: str, database: str = None):
         """
         self.connection_string = connection_string
         self.driver: Optional[ydb.Driver] = None
-        self.session_pool: Optional[ydb.SessionPool] = None
+        self.session_pool: Optional[ydb.aio.QuerySessionPool] = None
         self._database = database
         self.last_error = None
 
@@ -68,7 +69,7 @@ def _parse_endpoint_and_database(self) -> Tuple[str, str]:
 
         return endpoint, database
 
-    async def connect(self) -> Tuple[ydb.Driver, ydb.SessionPool]:
+    async def connect(self) -> Tuple[ydb.Driver, ydb.aio.QuerySessionPool]:
         """Connect to YDB and setup session pool asynchronously.
 
         Returns:
@@ -104,7 +105,7 @@ async def connect(self) -> Tuple[ydb.Driver, ydb.SessionPool]:
             logger.info("Connected to YDB successfully")
 
             # Create session pool
-            self.session_pool = ydb.SessionPool(self.driver)
+            self.session_pool = ydb.aio.QuerySessionPool(self.driver)
 
             return self.driver, self.session_pool