Skip to content

Commit

Permalink
feat: Adding SSL options for Postgres (#2644)
Browse files Browse the repository at this point in the history
* [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 <nborle@atb.com>

* [docs] SSL for Postgres

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

Signed-off-by: Neil Borle <nborle@atb.com>

Co-authored-by: Neil Borle <nborle@atb.com>
  • Loading branch information
NBor and NeilBATB authored May 10, 2022
1 parent ffa33ad commit 0e809c2
Show file tree
Hide file tree
Showing 5 changed files with 32 additions and 0 deletions.
6 changes: 6 additions & 0 deletions docs/reference/offline-stores/postgres.md
Original file line number Diff line number Diff line change
Expand Up @@ -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" %}
Expand All @@ -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
```
Expand Down
6 changes: 6 additions & 0 deletions docs/reference/online-stores/postgres.md
Original file line number Diff line number Diff line change
Expand Up @@ -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" %}
Expand All @@ -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 %}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Optional

import psycopg2
from psycopg2 import sql

Expand All @@ -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):
Expand All @@ -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
Expand Down
4 changes: 4 additions & 0 deletions sdk/python/feast/infra/utils/postgres/connection_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions sdk/python/feast/infra/utils/postgres/postgres_config.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Optional

from pydantic import StrictStr

from feast.repo_config import FeastConfigBaseModel
Expand All @@ -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

0 comments on commit 0e809c2

Please sign in to comment.