fix: Extend SQL registry config with a sqlalchemy_config_kwargs key (#3997)

galen-ft · web-flow · commit 21931d59f8a2 · 2024-03-08T21:35:34.000-08:00
* Add SQLAlchemy Engine settings to registry config

Signed-off-by: Galen &lt;galen.terziysky@ft.com&gt;

* Allow for dict values of any-type

Signed-off-by: Galen Terziysky &lt;105213101+galen-ft@users.noreply.github.com&gt;

* Extend PostgreSQL and MySQL registry tests and add examples to the docs

Signed-off-by: Galen Terziysky &lt;105213101+galen-ft@users.noreply.github.com&gt;

* Add "echo=False" to registry docs

Signed-off-by: Galen Terziysky &lt;105213101+galen-ft@users.noreply.github.com&gt;

* Fix mypy type checker error

Signed-off-by: Galen Terziysky &lt;105213101+galen-ft@users.noreply.github.com&gt;

* Accept change by the Black code formatter

Signed-off-by: Galen Terziysky &lt;105213101+galen-ft@users.noreply.github.com&gt;

---------

Signed-off-by: Galen &lt;galen.terziysky@ft.com&gt;
Signed-off-by: Galen Terziysky &lt;105213101+galen-ft@users.noreply.github.com&gt;
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)

Original file line number	Diff line number	Diff line change
`@@ -71,6 +71,7 @@ def pg_registry():`
`71`	`71`	`registry_config = RegistryConfig(`
`72`	`72`	`registry_type="sql",`
`73`	`73`	`path=f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",`
	`74`	`+ sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},`
`74`	`75`	`)`
`75`	`76`
`76`	`77`	`yield SqlRegistry(registry_config, "project", None)`
`@@ -106,6 +107,7 @@ def mysql_registry():`
`106`	`107`	`registry_config = RegistryConfig(`
`107`	`108`	`registry_type="sql",`
`108`	`109`	`path=f"mysql+pymysql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",`
	`110`	`+ sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},`
`109`	`111`	`)`
`110`	`112`
`111`	`113`	`yield SqlRegistry(registry_config, "project", None)`