Merge branch 'main' into feat/ydb-secondary-indicies

Author: LuckySting

CHANGELOG.md

@@ -1,3 +1,8 @@
+* Added declare for yql statement variables (opt in) - temporary flag
+
+## 0.0.1b12 ##
+* supported ydb connection credentials
+
 ## 0.0.1b11 ##
 * test release
 

setup.py

@@ -13,7 +13,7 @@
 
 setuptools.setup(
     name="ydb-sqlalchemy",
-    version="0.0.1b11",  # AUTOVERSION
+    version="0.0.1b12",  # AUTOVERSION
     description="YDB Dialect for SQLAlchemy",
     author="Yandex LLC",
     author_email="ydb@yandex-team.ru",

test/test_core.py

@@ -83,6 +83,11 @@ def test_sa_crud(self, connection):
             (5, "c"),
         ]
 
+    def test_sa_crud_with_add_declare(self):
+        engine = sa.create_engine(config.db_url, _add_declare_for_yql_stmt_vars=True)
+        with engine.connect() as connection:
+            self.test_sa_crud(connection)
+
 
 class TestSimpleSelect(TablesTest):
     __backend__ = True
@@ -597,6 +602,57 @@ def ydb_pool(self, ydb_driver):
             loop.run_until_complete(session_pool.stop())
 
 
+class TestCredentials(TestBase):
+    __backend__ = True
+    __only_on__ = "yql+ydb"
+
+    @pytest.fixture(scope="class")
+    def table_client_settings(self):
+        yield (
+            ydb.TableClientSettings()
+            .with_native_date_in_result_sets(True)
+            .with_native_datetime_in_result_sets(True)
+            .with_native_timestamp_in_result_sets(True)
+            .with_native_interval_in_result_sets(True)
+            .with_native_json_in_result_sets(False)
+        )
+
+    @pytest.fixture(scope="class")
+    def driver_config_for_credentials(self, table_client_settings):
+        url = config.db_url
+        endpoint = f"grpc://{url.host}:{url.port}"
+        database = url.database
+
+        yield ydb.DriverConfig(
+            endpoint=endpoint,
+            database=database,
+            table_client_settings=table_client_settings,
+        )
+
+    def test_ydb_credentials_good(self, table_client_settings, driver_config_for_credentials):
+        credentials_good = ydb.StaticCredentials(
+            driver_config=driver_config_for_credentials,
+            user="root",
+            password="1234",
+        )
+        engine = sa.create_engine(config.db_url, connect_args={"credentials": credentials_good})
+        with engine.connect() as conn:
+            result = conn.execute(sa.text("SELECT 1 as value"))
+            assert result.fetchone()
+
+    def test_ydb_credentials_bad(self, table_client_settings, driver_config_for_credentials):
+        credentials_bad = ydb.StaticCredentials(
+            driver_config=driver_config_for_credentials,
+            user="root",
+            password="56",
+        )
+        engine = sa.create_engine(config.db_url, connect_args={"credentials": credentials_bad})
+        with pytest.raises(Exception) as excinfo:
+            with engine.connect() as conn:
+                conn.execute(sa.text("SELECT 1 as value"))
+        assert "Invalid password" in str(excinfo.value)
+
+
 class TestUpsert(TablesTest):
     __backend__ = True
 

ydb_sqlalchemy/_version.py

@@ -1 +1 @@
-VERSION = "0.0.1b11"
+VERSION = "0.0.1b12"

ydb_sqlalchemy/dbapi/connection.py

@@ -37,6 +37,7 @@ def __init__(
         self.endpoint = f"grpc://{host}:{port}"
         self.database = database
         self.conn_kwargs = conn_kwargs
+        self.credentials = self.conn_kwargs.pop("credentials", None)
 
         if "ydb_session_pool" in self.conn_kwargs:  # Use session pool managed manually
             self._shared_session_pool = True
@@ -157,6 +158,7 @@ def _create_driver(self):
             endpoint=self.endpoint,
             database=self.database,
             table_client_settings=self._get_table_client_settings(),
+            credentials=self.credentials,
         )
         driver = self._ydb_driver_class(driver_config)
         try:

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -625,11 +625,14 @@ class YqlDialect(StrCompileDialect):
     def import_dbapi(cls: Any):
         return dbapi.YdbDBApi()
 
-    def __init__(self, json_serializer=None, json_deserializer=None, **kwargs):
+    def __init__(self, json_serializer=None, json_deserializer=None, _add_declare_for_yql_stmt_vars=False, **kwargs):
         super().__init__(**kwargs)
 
         self._json_deserializer = json_deserializer
         self._json_serializer = json_serializer
+        # NOTE: _add_declare_for_yql_stmt_vars is temporary and is soon to be removed.
+        # no need in declare in yql statement here since ydb 24-1
+        self._add_declare_for_yql_stmt_vars = _add_declare_for_yql_stmt_vars
 
     def _describe_table(self, connection, table_name, schema=None) -> ydb.TableDescription:
         if schema is not None:
@@ -752,6 +755,12 @@ def _format_variables(
         formatted_statement = formatted_statement.replace("%%", "%")
         return formatted_statement, formatted_parameters
 
+    def _add_declare_for_yql_stmt_vars_impl(self, statement, parameters_types):
+        declarations = "\n".join(
+            [f"DECLARE {param_name} as {str(param_type)};" for param_name, param_type in parameters_types.items()]
+        )
+        return f"{declarations}\n{statement}"
+
     def _make_ydb_operation(
         self,
         statement: str,
@@ -765,6 +774,8 @@ def _make_ydb_operation(
             parameters_types = context.compiled.get_bind_types(parameters)
             parameters_types = {f"${k}": v for k, v in parameters_types.items()}
             statement, parameters = self._format_variables(statement, parameters, execute_many)
+            if self._add_declare_for_yql_stmt_vars:
+                statement = self._add_declare_for_yql_stmt_vars_impl(statement, parameters_types)
             return dbapi.YdbQuery(yql_text=statement, parameters_types=parameters_types, is_ddl=is_ddl), parameters
 
         statement, parameters = self._format_variables(statement, parameters, execute_many)