Merge pull request #39 from kabulov/add_credentials_2 add credentials. set in connect_args for engine.

Author: rekby

test/test_core.py

@@ -597,6 +597,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/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: