Ability to pass custom DriverConfig kwargs

Author: vgvoleg

tests/test_connections.py

@@ -372,6 +372,16 @@ def connection(
         finally:
             conn.close()
 
+    def test_connect_with_custom_driver_config_kwargs(
+        self, connection_kwargs: dict
+    ) -> None:
+        connection_kwargs["driver_config_kwargs"] = {
+            "grpc_keep_alive_timeout": 777,
+        }
+        connection = dbapi.connect(**connection_kwargs)
+        assert connection._driver._driver_config.grpc_keep_alive_timeout == 777
+        connection.close()
+
     @pytest.mark.parametrize(
         ("isolation_level", "read_only"),
         [
@@ -461,6 +471,17 @@ def close() -> None:
 
             await greenlet_spawn(close)
 
+    @pytest.mark.asyncio
+    async def test_connect_with_custom_driver_config_kwargs(
+        self, connection_kwargs: dict
+    ) -> None:
+        connection_kwargs["driver_config_kwargs"] = {
+            "grpc_keep_alive_timeout": 777,
+        }
+        connection = await dbapi.async_connect(**connection_kwargs)
+        assert connection._driver._driver_config.grpc_keep_alive_timeout == 777
+        await connection.close()
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         ("isolation_level", "read_only"),

ydb_dbapi/connections.py

@@ -79,6 +79,7 @@ def __init__(
         ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
         root_certificates_path: str | None = None,
         root_certificates: str | None = None,
+        driver_config_kwargs: dict | None = None,
         **kwargs: dict,
     ) -> None:
         protocol = protocol if protocol else "grpc"
@@ -89,6 +90,8 @@ def __init__(
 
         self.connection_kwargs: dict = kwargs
 
+        driver_config_kwargs = driver_config_kwargs or {}
+
         self._shared_session_pool: bool = False
 
         self._tx_context: TxContext | AsyncTxContext | None = None
@@ -113,6 +116,7 @@ def __init__(
                 credentials=self.credentials,
                 query_client_settings=self._get_client_settings(),
                 root_certificates=root_certificates,
+                **driver_config_kwargs,
             )
             self._driver = self._driver_cls(driver_config)
             self._session_pool = self._pool_cls(self._driver, size=5)
@@ -197,6 +201,7 @@ def __init__(
         ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
         root_certificates_path: str | None = None,
         root_certificates: str | None = None,
+        driver_config_kwargs: dict | None = None,
         **kwargs: dict,
     ) -> None:
         super().__init__(
@@ -209,6 +214,7 @@ def __init__(
             ydb_session_pool=ydb_session_pool,
             root_certificates_path=root_certificates_path,
             root_certificates=root_certificates,
+            driver_config_kwargs=driver_config_kwargs,
             **kwargs,
         )
         self._current_cursor: Cursor | None = None
@@ -390,6 +396,7 @@ def __init__(
         ydb_session_pool: SessionPool | AsyncSessionPool | None = None,
         root_certificates_path: str | None = None,
         root_certificates: str | None = None,
+        driver_config_kwargs: dict | None = None,
         **kwargs: dict,
     ) -> None:
         super().__init__(
@@ -402,6 +409,7 @@ def __init__(
             ydb_session_pool=ydb_session_pool,
             root_certificates_path=root_certificates_path,
             root_certificates=root_certificates,
+            driver_config_kwargs=driver_config_kwargs,
             **kwargs,
         )
         self._current_cursor: AsyncCursor | None = None