Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 6 commits into from
Mar 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions docs/getting-started/concepts/registry.md
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,9 @@ registry:
registry_type: sql
path: postgresql://postgres:[email protected]: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)
Expand Down
3 changes: 3 additions & 0 deletions docs/tutorials/using-scalable-registry.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ registry:
registry_type: sql
path: postgresql://postgres:[email protected]: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.
Expand Down
9 changes: 7 additions & 2 deletions sdk/python/feast/infra/registry/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__(
Expand All @@ -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)
Expand Down
3 changes: 3 additions & 0 deletions sdk/python/feast/repo_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,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`"""
Expand Down
2 changes: 2 additions & 0 deletions sdk/python/tests/unit/test_sql_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
Loading