Skip to content

Commit 21931d5

Browse files
galen-ftgalen-ft
authored
fix: Extend SQL registry config with a sqlalchemy_config_kwargs key (#3997)
* Add SQLAlchemy Engine settings to registry config Signed-off-by: Galen <[email protected]> * Allow for dict values of any-type Signed-off-by: Galen Terziysky <[email protected]> * Extend PostgreSQL and MySQL registry tests and add examples to the docs Signed-off-by: Galen Terziysky <[email protected]> * Add "echo=False" to registry docs Signed-off-by: Galen Terziysky <[email protected]> * Fix mypy type checker error Signed-off-by: Galen Terziysky <[email protected]> * Accept change by the Black code formatter Signed-off-by: Galen Terziysky <[email protected]> --------- Signed-off-by: Galen <[email protected]> Signed-off-by: Galen Terziysky <[email protected]>
1 parent 42a7b81 commit 21931d5

5 files changed

Lines changed: 18 additions & 2 deletions

File tree

‎docs/getting-started/concepts/registry.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ registry:
5757
registry_type: sql
5858
path: postgresql://postgres:[email protected]:55001/feast
5959
cache_ttl_seconds: 60
60+
sqlalchemy_config_kwargs:
61+
echo: false
62+
pool_pre_ping: true
6063
```
6164
6265
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)

‎docs/tutorials/using-scalable-registry.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ registry:
2929
registry_type: sql
3030
path: postgresql://postgres:[email protected]:55001/feast
3131
cache_ttl_seconds: 60
32+
sqlalchemy_config_kwargs:
33+
echo: false
34+
pool_pre_ping: true
3235
```
3336
3437
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.

‎sdk/python/feast/infra/registry/sql.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from enum import Enum
55
from pathlib import Path
66
from threading import Lock
7-
from typing import Any, Callable, List, Optional, Set, Union
7+
from typing import Any, Callable, Dict, List, Optional, Set, Union
88

99
from pydantic import StrictStr
1010
from sqlalchemy import ( # type: ignore
@@ -190,6 +190,9 @@ class SqlRegistryConfig(RegistryConfig):
190190
""" str: Path to metadata store.
191191
If registry_type is 'sql', then this is a database URL as expected by SQLAlchemy """
192192

193+
sqlalchemy_config_kwargs: Dict[str, Any] = {"echo": False}
194+
""" Dict[str, Any]: Extra arguments to pass to SQLAlchemy.create_engine. """
195+
193196

194197
class SqlRegistry(BaseRegistry):
195198
def __init__(
@@ -199,7 +202,9 @@ def __init__(
199202
repo_path: Optional[Path],
200203
):
201204
assert registry_config is not None, "SqlRegistry needs a valid registry_config"
202-
self.engine: Engine = create_engine(registry_config.path, echo=False)
205+
self.engine: Engine = create_engine(
206+
registry_config.path, **registry_config.sqlalchemy_config_kwargs
207+
)
203208
metadata.create_all(self.engine)
204209
self.cached_registry_proto = self.proto()
205210
proto_registry_utils.init_project_metadata(self.cached_registry_proto, project)

‎sdk/python/feast/repo_config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,9 @@ class RegistryConfig(FeastBaseModel):
126126
s3_additional_kwargs: Optional[Dict[str, str]] = None
127127
""" Dict[str, str]: Extra arguments to pass to boto3 when writing the registry file to S3. """
128128

129+
sqlalchemy_config_kwargs: Dict[str, Any] = {}
130+
""" Dict[str, Any]: Extra arguments to pass to SQLAlchemy.create_engine. """
131+
129132

130133
class RepoConfig(FeastBaseModel):
131134
"""Repo config. Typically loaded from `feature_store.yaml`"""

‎sdk/python/tests/unit/test_sql_registry.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ def pg_registry():
7171
registry_config = RegistryConfig(
7272
registry_type="sql",
7373
path=f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
74+
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
7475
)
7576

7677
yield SqlRegistry(registry_config, "project", None)
@@ -106,6 +107,7 @@ def mysql_registry():
106107
registry_config = RegistryConfig(
107108
registry_type="sql",
108109
path=f"mysql+pymysql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
110+
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
109111
)
110112

111113
yield SqlRegistry(registry_config, "project", None)

0 commit comments

Comments
 (0)