diff --git a/sdk/python/feast/infra/registry/sql.py b/sdk/python/feast/infra/registry/sql.py index 0a037e519c2..8d2c88fba57 100644 --- a/sdk/python/feast/infra/registry/sql.py +++ b/sdk/python/feast/infra/registry/sql.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union, cast -from pydantic import StrictInt, StrictStr +from pydantic import StrictInt, StrictStr, field_validator from sqlalchemy import ( # type: ignore BigInteger, Column, @@ -306,6 +306,15 @@ class SqlRegistryConfig(RegistryConfig): thread_pool_executor_worker_count: StrictInt = 0 """ int: Number of worker threads to use for asynchronous caching in SQL Registry. If set to 0, it doesn't use ThreadPoolExecutor. """ + @field_validator("read_path") + def validate_read_path(cls, read_path: Optional[str]) -> Optional[str]: + # Mirror `RegistryConfig.validate_path`: a bare `postgresql://` read_path + # must be rewritten to the psycopg3 driver too, otherwise it silently + # falls back to psycopg2 while `path` uses psycopg3. + if read_path is not None: + return cls._normalize_postgres_scheme(read_path, "read_path") + return read_path + class SqlRegistry(CachingRegistry): def __init__( diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index b6ee3885b3c..775a62aab57 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -205,19 +205,33 @@ class RegistryConfig(FeastBaseModel): mcp: Optional[McpRegistryConfig] = None """ McpRegistryConfig: MCP (Model Context Protocol) configuration for the registry REST server. """ + @staticmethod + def _normalize_postgres_scheme(value: str, field_name: str) -> str: + """Rewrite a bare ``postgresql://`` URL to the psycopg3 driver, with a warning. + + SQLAlchemy resolves a bare ``postgresql://`` to the psycopg2 driver, while + feast standardizes on psycopg3 (``postgresql+psycopg``). Pass an explicit + ``postgresql+psycopg2`` to keep psycopg2. Shared by the ``path`` and + ``read_path`` validators so both endpoints normalize identically. + """ + if value.startswith("postgresql://"): + _logger.warning( + f"The `{field_name}` of the `RegistryConfig` starts with a plain " + "`postgresql` string. We are updating this to `postgresql+psycopg` " + "to ensure that the `psycopg3` driver is used by `sqlalchemy`. If " + f"you want to use `psycopg2` pass `postgresql+psycopg2` explicitly " + f"to `{field_name}`. To silence this warning, pass `postgresql+psycopg` " + f"explicitly to `{field_name}`." + ) + # Rewrite only the leading scheme, not any later occurrence (e.g. + # inside credentials or a query string). + return "postgresql+psycopg://" + value[len("postgresql://") :] + return value + @field_validator("path") def validate_path(cls, path: str, values: ValidationInfo) -> str: if values.data.get("registry_type") == "sql": - if path.startswith("postgresql://"): - _logger.warning( - "The `path` of the `RegistryConfig` starts with a plain " - "`postgresql` string. We are updating this to `postgresql+psycopg` " - "to ensure that the `psycopg3` driver is used by `sqlalchemy`. If " - "you want to use `psycopg2` pass `postgresql+psycopg2` explicitely " - "to `path`. To silence this warning, pass `postgresql+psycopg` " - "explicitely to `path`." - ) - return path.replace("postgresql://", "postgresql+psycopg://") + return cls._normalize_postgres_scheme(path, "path") return path diff --git a/sdk/python/tests/unit/infra/registry/test_sql_registry.py b/sdk/python/tests/unit/infra/registry/test_sql_registry.py index 9c21cfc85d0..5c94ce744eb 100644 --- a/sdk/python/tests/unit/infra/registry/test_sql_registry.py +++ b/sdk/python/tests/unit/infra/registry/test_sql_registry.py @@ -65,6 +65,72 @@ def shared_sqlite_db_path(): yield path +def test_read_path_normalizes_bare_postgresql_scheme(caplog): + """A bare `postgresql://` read_path is rewritten to the psycopg3 driver, the + same way `path` is, and logs a migration warning. Without this, read_path + silently falls back to psycopg2 while path uses psycopg3.""" + with caplog.at_level(logging.WARNING): + config = SqlRegistryConfig( + registry_type="sql", + path="postgresql://localhost:5432/db", + read_path="postgresql://localhost:5432/replica", + ) + assert config.path == "postgresql+psycopg://localhost:5432/db" + assert config.read_path == "postgresql+psycopg://localhost:5432/replica" + # The migration warning names the read_path field specifically. + assert "`read_path` of the `RegistryConfig`" in caplog.text + + +def test_read_path_normalized_when_registry_type_defaulted(): + """read_path normalization must still fire when registry_type is left to its + default ('sql') rather than passed explicitly — the common real-world config.""" + config = SqlRegistryConfig( + path="postgresql://localhost/db", + read_path="postgresql://localhost/replica", + ) + assert config.registry_type == "sql" + assert config.read_path == "postgresql+psycopg://localhost/replica" + + +@pytest.mark.parametrize( + "read_path", + [ + "postgresql+psycopg2://localhost/replica", # explicit psycopg2 preserved + "postgresql+psycopg://localhost/replica", # already psycopg3 + "mysql://localhost/replica", # non-postgres left untouched + ], +) +def test_read_path_leaves_explicit_scheme_untouched(read_path): + config = SqlRegistryConfig( + registry_type="sql", + path="sqlite:///unused.db", + read_path=read_path, + ) + assert config.read_path == read_path + + +def test_read_path_none_stays_none(): + config = SqlRegistryConfig(registry_type="sql", path="sqlite:///unused.db") + assert config.read_path is None + + +@pytest.mark.parametrize("field_name", ["path", "read_path"]) +def test_postgres_normalization_only_rewrites_leading_scheme(field_name): + raw_url = "postgresql://localhost/db?target=postgresql://replica" + config_values = { + "registry_type": "sql", + "path": "sqlite:///unused.db", + } + config_values[field_name] = raw_url + + config = SqlRegistryConfig(**config_values) + + assert ( + getattr(config, field_name) + == "postgresql+psycopg://localhost/db?target=postgresql://replica" + ) + + def test_proto_columns_use_longblob_on_mysql(): """On MySQL and MariaDB, serialized-proto columns must compile to LONGBLOB.