diff --git a/docs/getting-started/concepts/registry.md b/docs/getting-started/concepts/registry.md index f7d4a5b3e1..8ac32ce87b 100644 --- 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 index a87aedd9b9..30b8e01ed5 100644 --- 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 index d57bcc7c0a..1e5c2a8725 100644 --- 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 index 263ba81e39..5708754622 100644 --- 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 index 722e318b0c..4ca41423c1 100644 --- 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)