fix: Escape special characters in the Postgres password (#4394) (#4396)

* Apply fix * Add special characters to postgres online store test * Fix linting error --------- Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>
feast-dev · Aug 9, 2024 · b456984 · b456984
1 parent 44f0e84
commit b456984
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/sdk/python/feast/infra/utils/postgres/connection_utils.py b/sdk/python/feast/infra/utils/postgres/connection_utils.py
@@ -5,6 +5,7 @@
 import psycopg
 import pyarrow as pa
 from psycopg import AsyncConnection, Connection
+from psycopg.conninfo import make_conninfo
 from psycopg_pool import AsyncConnectionPool, ConnectionPool
 
 from feast.infra.utils.postgres.postgres_config import PostgreSQLConfig
@@ -55,13 +56,14 @@ async def _get_connection_pool_async(config: PostgreSQLConfig) -> AsyncConnectio
 
 def _get_conninfo(config: PostgreSQLConfig) -> str:
     """Get the `conninfo` argument required for connection objects."""
-    return (
-        f"postgresql://{config.user}"
-        f":{config.password}"
-        f"@{config.host}"
-        f":{int(config.port)}"
-        f"/{config.database}"
-    )
+    psycopg_config = {
+        "user": config.user,
+        "password": config.password,
+        "host": config.host,
+        "port": int(config.port),
+        "dbname": config.database,
+    }
+    return make_conninfo(conninfo="", **psycopg_config)
 
 
 def _get_conn_kwargs(config: PostgreSQLConfig) -> Dict[str, Any]:

diff --git a/sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py b/sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py
@@ -16,7 +16,7 @@ def __init__(self, project_name: str, **kwargs):
         self.container = PostgresContainer(
             "postgres:16",
             username="root",
-            password="test",
+            password="test!@#$%",
             dbname="test",
         ).with_exposed_ports(5432)
 
@@ -26,7 +26,7 @@ def create_online_store(self) -> Dict[str, str]:
             "host": "localhost",
             "type": "postgres",
             "user": "root",
-            "password": "test",
+            "password": "test!@#$%",
             "database": "test",
             "port": self.container.get_exposed_port(5432),
         }
@@ -42,7 +42,7 @@ def __init__(self, project_name: str, **kwargs):
         self.container = (
             DockerContainer("pgvector/pgvector:pg16")
             .with_env("POSTGRES_USER", "root")
-            .with_env("POSTGRES_PASSWORD", "test")
+            .with_env("POSTGRES_PASSWORD", "test!@#$%")
             .with_env("POSTGRES_DB", "test")
             .with_exposed_ports(5432)
             .with_volume_mapping(
@@ -65,7 +65,7 @@ def create_online_store(self) -> Dict[str, Any]:
             "host": "localhost",
             "type": "postgres",
             "user": "root",
-            "password": "test",
+            "password": "test!@#$%",
             "database": "test",
             "pgvector_enabled": True,
             "vector_len": 2,