diff --git a/docs/reference/offline-stores/postgres.md b/docs/reference/offline-stores/postgres.md index b64b2350caa..7d04ac5843a 100644 --- a/docs/reference/offline-stores/postgres.md +++ b/docs/reference/offline-stores/postgres.md @@ -13,6 +13,8 @@ The PostgreSQL offline store is an offline store that provides support for readi * `to_df` to retrieve the pandas dataframe. * `to_arrow` to retrieve the dataframe as a PyArrow table. +* sslmode, sslkey_path, sslcert_path, and sslrootcert_path are optional + ## Example {% code title="feature_store.yaml" %} @@ -28,6 +30,10 @@ offline_store: db_schema: DB_SCHEMA user: DB_USERNAME password: DB_PASSWORD + sslmode: verify-ca + sslkey_path: /path/to/client-key.pem + sslcert_path: /path/to/client-cert.pem + sslrootcert_path: /path/to/server-ca.pem online_store: path: data/online_store.db ``` diff --git a/docs/reference/online-stores/postgres.md b/docs/reference/online-stores/postgres.md index 917673f278c..3e507c70fc7 100644 --- a/docs/reference/online-stores/postgres.md +++ b/docs/reference/online-stores/postgres.md @@ -6,6 +6,8 @@ The PostgreSQL online store provides support for materializing feature values in * Only the latest feature values are persisted +* sslmode, sslkey_path, sslcert_path, and sslrootcert_path are optional + ## Example {% code title="feature_store.yaml" %} @@ -21,6 +23,10 @@ online_store: db_schema: DB_SCHEMA user: DB_USERNAME password: DB_PASSWORD + sslmode: verify-ca + sslkey_path: /path/to/client-key.pem + sslcert_path: /path/to/client-cert.pem + sslrootcert_path: /path/to/server-ca.pem ``` {% endcode %} diff --git a/sdk/python/feast/infra/registry_stores/contrib/postgres/registry_store.py b/sdk/python/feast/infra/registry_stores/contrib/postgres/registry_store.py index 6b0ae1910dd..b3c0c6bd365 100644 --- a/sdk/python/feast/infra/registry_stores/contrib/postgres/registry_store.py +++ b/sdk/python/feast/infra/registry_stores/contrib/postgres/registry_store.py @@ -1,3 +1,5 @@ +from typing import Optional + import psycopg2 from psycopg2 import sql @@ -15,6 +17,10 @@ class PostgresRegistryConfig(RegistryConfig): db_schema: str user: str password: str + sslmode: Optional[str] + sslkey_path: Optional[str] + sslcert_path: Optional[str] + sslrootcert_path: Optional[str] class PostgreSQLRegistryStore(RegistryStore): @@ -26,6 +32,10 @@ def __init__(self, config: PostgresRegistryConfig, registry_path: str): db_schema=config.db_schema, user=config.user, password=config.password, + sslmode=getattr(config, "sslmode", None), + sslkey_path=getattr(config, "sslkey_path", None), + sslcert_path=getattr(config, "sslcert_path", None), + sslrootcert_path=getattr(config, "sslrootcert_path", None), ) self.table_name = config.path self.cache_ttl_seconds = config.cache_ttl_seconds diff --git a/sdk/python/feast/infra/utils/postgres/connection_utils.py b/sdk/python/feast/infra/utils/postgres/connection_utils.py index c3a7c913d4c..6dbb4a4bc01 100644 --- a/sdk/python/feast/infra/utils/postgres/connection_utils.py +++ b/sdk/python/feast/infra/utils/postgres/connection_utils.py @@ -17,6 +17,10 @@ def _get_conn(config: PostgreSQLConfig): port=int(config.port), user=config.user, password=config.password, + sslmode=config.sslmode, + sslkey=config.sslkey_path, + sslcert=config.sslcert_path, + sslrootcert=config.sslrootcert_path, options="-c search_path={}".format(config.db_schema or config.user), ) return conn diff --git a/sdk/python/feast/infra/utils/postgres/postgres_config.py b/sdk/python/feast/infra/utils/postgres/postgres_config.py index c7a581f0722..f22cc6c204e 100644 --- a/sdk/python/feast/infra/utils/postgres/postgres_config.py +++ b/sdk/python/feast/infra/utils/postgres/postgres_config.py @@ -1,3 +1,5 @@ +from typing import Optional + from pydantic import StrictStr from feast.repo_config import FeastConfigBaseModel @@ -10,3 +12,7 @@ class PostgreSQLConfig(FeastConfigBaseModel): db_schema: StrictStr = "public" user: StrictStr password: StrictStr + sslmode: Optional[StrictStr] = None + sslkey_path: Optional[StrictStr] = None + sslcert_path: Optional[StrictStr] = None + sslrootcert_path: Optional[StrictStr] = None