feat: Adding SSL options for Postgres (feast-dev#2644)

NBor · NeilBATB · web-flow · commit 0e809c2c86a8 · 2022-05-10T14:42:31.000-07:00
* [feat] Allow SSL for Postgres

Allow user to set SSL options in the feature_story.yml file for the
offline, online, and registry sections

Signed-off-by: Neil Borle &lt;nborle@atb.com&gt;

* [docs] SSL for Postgres

Updated the reference docs for offline/online stores to describe the SSL
options available

Signed-off-by: Neil Borle &lt;nborle@atb.com&gt;

Co-authored-by: Neil Borle &lt;nborle@atb.com&gt;
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
