fix: Extend SQL registry config with a sqlalchemy_config_kwargs key (f…

…east-dev#3997) * Add SQLAlchemy Engine settings to registry config Signed-off-by: Galen <galen.terziysky@ft.com> * Allow for dict values of any-type Signed-off-by: Galen Terziysky <105213101+galen-ft@users.noreply.github.com> * Extend PostgreSQL and MySQL registry tests and add examples to the docs Signed-off-by: Galen Terziysky <105213101+galen-ft@users.noreply.github.com> * Add "echo=False" to registry docs Signed-off-by: Galen Terziysky <105213101+galen-ft@users.noreply.github.com> * Fix mypy type checker error Signed-off-by: Galen Terziysky <105213101+galen-ft@users.noreply.github.com> * Accept change by the Black code formatter Signed-off-by: Galen Terziysky <105213101+galen-ft@users.noreply.github.com> --------- Signed-off-by: Galen <galen.terziysky@ft.com> Signed-off-by: Galen Terziysky <105213101+galen-ft@users.noreply.github.com>
tqtensor · Mar 11, 2024 · 3c1c012 · 3c1c012
1 parent 32e3e19
commit 3c1c012
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 2 deletions.
diff --git a/docs/getting-started/concepts/registry.md b/docs/getting-started/concepts/registry.md
@@ -57,6 +57,9 @@ registry:
     registry_type: sql
     path: postgresql://postgres:mysecretpassword@127.0.0.1:55001/feast
     cache_ttl_seconds: 60
+    sqlalchemy_config_kwargs:
+        echo: false
+        pool_pre_ping: true
 ```
 
 This supports any SQLAlchemy compatible database as a backend. The exact schema can be seen in [sql.py](https://github.com/feast-dev/feast/blob/master/sdk/python/feast/infra/registry/sql.py)

diff --git a/docs/tutorials/using-scalable-registry.md b/docs/tutorials/using-scalable-registry.md
@@ -29,6 +29,9 @@ registry:
     registry_type: sql
     path: postgresql://postgres:mysecretpassword@127.0.0.1:55001/feast
     cache_ttl_seconds: 60
+    sqlalchemy_config_kwargs:
+        echo: false
+        pool_pre_ping: true
 ```
 
 Specifically, the registry_type needs to be set to sql in the registry config block. On doing so, the path should refer to the [Database URL](https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls) for the database to be used, as expected by SQLAlchemy. No other additional commands are currently needed to configure this registry.

diff --git a/sdk/python/feast/infra/registry/sql.py b/sdk/python/feast/infra/registry/sql.py
@@ -4,7 +4,7 @@
 from enum import Enum
 from pathlib import Path
 from threading import Lock
-from typing import Any, Callable, List, Optional, Set, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Union
 
 from pydantic import StrictStr
 from sqlalchemy import (  # type: ignore
@@ -190,6 +190,9 @@ class SqlRegistryConfig(RegistryConfig):
     """ str: Path to metadata store.
     If registry_type is 'sql', then this is a database URL as expected by SQLAlchemy """
 
+    sqlalchemy_config_kwargs: Dict[str, Any] = {"echo": False}
+    """ Dict[str, Any]: Extra arguments to pass to SQLAlchemy.create_engine. """
+
 
 class SqlRegistry(BaseRegistry):
     def __init__(
@@ -199,7 +202,9 @@ def __init__(
         repo_path: Optional[Path],
     ):
         assert registry_config is not None, "SqlRegistry needs a valid registry_config"
-        self.engine: Engine = create_engine(registry_config.path, echo=False)
+        self.engine: Engine = create_engine(
+            registry_config.path, **registry_config.sqlalchemy_config_kwargs
+        )
         metadata.create_all(self.engine)
         self.cached_registry_proto = self.proto()
         proto_registry_utils.init_project_metadata(self.cached_registry_proto, project)

diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -126,6 +126,9 @@ class RegistryConfig(FeastBaseModel):
     s3_additional_kwargs: Optional[Dict[str, str]] = None
     """ Dict[str, str]: Extra arguments to pass to boto3 when writing the registry file to S3. """
 
+    sqlalchemy_config_kwargs: Dict[str, Any] = {}
+    """ Dict[str, Any]: Extra arguments to pass to SQLAlchemy.create_engine. """
+
 
 class RepoConfig(FeastBaseModel):
     """Repo config. Typically loaded from `feature_store.yaml`"""

diff --git a/sdk/python/tests/unit/test_sql_registry.py b/sdk/python/tests/unit/test_sql_registry.py
@@ -71,6 +71,7 @@ def pg_registry():
     registry_config = RegistryConfig(
         registry_type="sql",
         path=f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
+        sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
     )
 
     yield SqlRegistry(registry_config, "project", None)
@@ -106,6 +107,7 @@ def mysql_registry():
     registry_config = RegistryConfig(
         registry_type="sql",
         path=f"mysql+pymysql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
+        sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
     )
 
     yield SqlRegistry(registry_config, "project", None)