Add ydb_root_directory parameter

Author: LuckySting

test/test_core.py

@@ -995,3 +995,58 @@ def test_index_with_join_usage(self, connection: sa.Connection, metadata: sa.Met
 
         cursor = connection.execute(select_stmt)
         assert cursor.one() == ("Sarah Connor", "wanted")
+
+
+class TestRootDirectory(TablesTest):
+    __backend__ = True
+
+    @classmethod
+    def define_tables(cls, metadata: sa.MetaData):
+        Table("some_dir/nested_dir/table", metadata, sa.Column("id", sa.Integer, primary_key=True))
+        Table("table", metadata, sa.Column("id", sa.Integer, primary_key=True))
+
+    @classmethod
+    def insert_data(cls, connection: sa.Connection):
+        table = cls.tables["some_dir/nested_dir/table"]
+        root_table = cls.tables["table"]
+
+        connection.execute(sa.insert(table).values({"id": 1}))
+        connection.execute(sa.insert(root_table).values({"id": 2}))
+
+    def test_root_directory(self):
+        engine = sa.create_engine(config.db_url, connect_args={"ydb_root_directory": "/local/some_dir/nested_dir"})
+        rel_table = Table("table", sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True))
+        abs_table = Table("/local/table", sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True))
+
+        with engine.connect() as conn:
+            result1 = conn.execute(sa.select(rel_table)).scalar()
+            result2 = conn.execute(sa.select(abs_table)).scalar()
+
+        assert result1 == 1
+        assert result2 == 2
+
+    def test_two_engines(self):
+        create_engine = sa.create_engine(config.db_url, connect_args={"ydb_root_directory": "/local/two/engines/test"})
+        select_engine = sa.create_engine(config.db_url, connect_args={"ydb_root_directory": "/local/two"})
+        table_to_create = Table("table", sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True))
+        table_to_select = Table("engines/test/table", sa.MetaData(), sa.Column("id", sa.Integer, primary_key=True))
+
+        table_to_create.create(create_engine)
+        try:
+            with create_engine.begin() as conn:
+                conn.execute(sa.insert(table_to_create).values({"id": 42}))
+
+            with select_engine.begin() as conn:
+                result = conn.execute(sa.select(table_to_select)).scalar()
+        finally:
+            table_to_create.drop(create_engine)
+
+        assert result == 42
+
+    def test_reflection(self):
+        reflection_engine = sa.create_engine(config.db_url, connect_args={"ydb_root_directory": "/local/some_dir"})
+        metadata = sa.MetaData()
+
+        metadata.reflect(reflection_engine)
+
+        assert "nested_dir/table" in metadata.tables

test/test_orm.py

@@ -1,6 +1,7 @@
+from types import MethodType
+
 import pytest
 import sqlalchemy as sa
-from types import MethodType
 from sqlalchemy import Column, Integer, Unicode
 from sqlalchemy.orm import declarative_base, sessionmaker
 from sqlalchemy.testing.fixtures import TablesTest, config

ydb_sqlalchemy/dbapi/connection.py

@@ -38,6 +38,7 @@ def __init__(
         self.database = database
         self.conn_kwargs = conn_kwargs
         self.credentials = self.conn_kwargs.pop("credentials", None)
+        self.root_directory = self.conn_kwargs.pop("ydb_root_directory", "")
 
         if "ydb_session_pool" in self.conn_kwargs:  # Use session pool managed manually
             self._shared_session_pool = True
@@ -58,20 +59,22 @@ def __init__(
         self.tx_context: Optional[ydb.TxContext] = None
 
     def cursor(self):
-        return self._cursor_class(self.session_pool, self.tx_mode, self.tx_context)
+        return self._cursor_class(self.session_pool, self.tx_mode, self.tx_context, self.root_directory)
 
     def describe(self, table_path: str) -> ydb.TableDescription:
-        abs_table_path = posixpath.join(self.database, table_path)
+        abs_table_path = posixpath.join(self.database, self.root_directory, table_path)
         cursor = self.cursor()
         return cursor.describe_table(abs_table_path)
 
     def check_exists(self, table_path: str) -> ydb.SchemeEntry:
+        abs_table_path = posixpath.join(self.database, self.root_directory, table_path)
         cursor = self.cursor()
-        return cursor.check_exists(table_path)
+        return cursor.check_exists(abs_table_path)
 
     def get_table_names(self) -> List[str]:
+        abs_dir_path = posixpath.join(self.database, self.root_directory)
         cursor = self.cursor()
-        return cursor.get_table_names()
+        return [posixpath.relpath(path, abs_dir_path) for path in cursor.get_table_names(abs_dir_path)]
 
     def set_isolation_level(self, isolation_level: str):
         class IsolationSettings(NamedTuple):

ydb_sqlalchemy/dbapi/cursor.py

@@ -3,6 +3,7 @@
 import functools
 import itertools
 import logging
+import posixpath
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
 
 import ydb
@@ -78,6 +79,7 @@ def __init__(
         session_pool: Union[ydb.SessionPool, ydb.aio.SessionPool],
         tx_mode: ydb.AbstractTransactionModeBuilder,
         tx_context: Optional[ydb.BaseTxContext] = None,
+        root_directory: str = "",
     ):
         self.session_pool = session_pool
         self.tx_mode = tx_mode
@@ -86,33 +88,36 @@ def __init__(
         self.arraysize = 1
         self.rows = None
         self._rows_prefetched = None
+        self.root_directory = root_directory
 
     @_handle_ydb_errors
     def describe_table(self, abs_table_path: str) -> ydb.TableDescription:
         return self._retry_operation_in_pool(self._describe_table, abs_table_path)
 
-    def check_exists(self, table_path: str) -> bool:
+    def check_exists(self, abs_table_path: str) -> bool:
         try:
-            self._retry_operation_in_pool(self._describe_path, table_path)
+            self._retry_operation_in_pool(self._describe_path, abs_table_path)
             return True
         except ydb.SchemeError:
             return False
 
     @_handle_ydb_errors
-    def get_table_names(self) -> List[str]:
-        directory: ydb.Directory = self._retry_operation_in_pool(self._list_directory)
-        return [child.name for child in directory.children if child.is_table()]
+    def get_table_names(self, abs_dir_path: str) -> List[str]:
+        directory: ydb.Directory = self._retry_operation_in_pool(self._list_directory, abs_dir_path)
+        result = []
+        for child in directory.children:
+            child_abs_path = posixpath.join(abs_dir_path, child.name)
+            if child.is_table():
+                result.append(child_abs_path)
+            elif child.is_directory() and not child.name.startswith("."):
+                result.extend(self.get_table_names(child_abs_path))
+        return result
 
     def execute(self, operation: YdbQuery, parameters: Optional[Mapping[str, Any]] = None):
-        if operation.is_ddl or not operation.parameters_types:
-            query = operation.yql_text
-            is_ddl = operation.is_ddl
-        else:
-            query = ydb.DataQuery(operation.yql_text, operation.parameters_types)
-            is_ddl = operation.is_ddl
+        query = self._get_ydb_query(operation)
 
         logger.info("execute sql: %s, params: %s", query, parameters)
-        if is_ddl:
+        if operation.is_ddl:
             chunks = self._execute_ddl(query)
         else:
             chunks = self._execute_dml(query, parameters)
@@ -130,6 +135,15 @@ def execute(self, operation: YdbQuery, parameters: Optional[Mapping[str, Any]] =
 
         self.rows = rows
 
+    def _get_ydb_query(self, operation: YdbQuery) -> Union[ydb.DataQuery, str]:
+        pragma = ""
+        if self.root_directory:
+            pragma = f'PRAGMA TablePathPrefix = "{self.root_directory}";\n'
+        if operation.is_ddl or not operation.parameters_types:
+            return pragma + operation.yql_text
+
+        return ydb.DataQuery(pragma + operation.yql_text, operation.parameters_types)
+
     @_handle_ydb_errors
     def _execute_dml(
         self, query: Union[ydb.DataQuery, str], parameters: Optional[Mapping[str, Any]] = None
@@ -163,8 +177,8 @@ def _describe_path(session: ydb.Session, table_path: str) -> ydb.SchemeEntry:
         return session._driver.scheme_client.describe_path(table_path)
 
     @staticmethod
-    def _list_directory(session: ydb.Session) -> ydb.Directory:
-        return session._driver.scheme_client.list_directory(session._driver._driver_config.database)
+    def _list_directory(session: ydb.Session, abs_dir_path: str) -> ydb.Directory:
+        return session._driver.scheme_client.list_directory(abs_dir_path)
 
     @staticmethod
     def _prepare(session: ydb.Session, query: str) -> ydb.DataQuery:
@@ -264,12 +278,12 @@ async def _describe_table(session: ydb.aio.table.Session, abs_table_path: str) -
         return await session.describe_table(abs_table_path)
 
     @staticmethod
-    async def _describe_path(session: ydb.aio.table.Session, table_path: str) -> ydb.SchemeEntry:
-        return await session._driver.scheme_client.describe_path(table_path)
+    async def _describe_path(session: ydb.aio.table.Session, abs_table_path: str) -> ydb.SchemeEntry:
+        return await session._driver.scheme_client.describe_path(abs_table_path)
 
     @staticmethod
-    async def _list_directory(session: ydb.aio.table.Session) -> ydb.Directory:
-        return await session._driver.scheme_client.list_directory(session._driver._driver_config.database)
+    async def _list_directory(session: ydb.aio.table.Session, abs_dir_path: str) -> ydb.Directory:
+        return await session._driver.scheme_client.list_directory(abs_dir_path)
 
     @staticmethod
     async def _execute_scheme(session: ydb.aio.table.Session, query: str) -> ydb.convert.ResultSets: