feast-dev · HaoXuAI · Mar 9, 2024 · Feb 16, 2024 · Feb 16, 2024 · Mar 7, 2024
@@ -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)

@@ -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.

@@ -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)

@@ -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`"""

@@ -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)