From 74528cfc3321d28317983f39b856136eedf64f80 Mon Sep 17 00:00:00 2001
From: Job Almekinders <55230856+job-almekinders@users.noreply.github.com>
Date: Thu, 8 Aug 2024 16:51:49 +0200
Subject: [PATCH] fix: Escape special characters in the Postgres password
 (#4394)

* Apply fix

Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>

* Add special characters to postgres online store test

Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>

* Fix linting error

Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>

---------

Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>
---
 .../infra/utils/postgres/connection_utils.py     | 16 +++++++++-------
 .../universal/online_store/postgres.py           |  8 ++++----
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/sdk/python/feast/infra/utils/postgres/connection_utils.py b/sdk/python/feast/infra/utils/postgres/connection_utils.py
index 7b37ea981f..3749fc2fc1 100644
--- 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
index e409862641..622ee99e14 100644
--- 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,