Utils: Use psycopg3

Use new ConnectionPool Pass kwargs as named argument Use executemany over execute_values Remove not-required open argument in psycopg.connect Improve Use SpooledTemporaryFile Use max_size and add docstring Properly write with StringIO Utils: Use SpooledTemporaryFile over StringIO object Add replace Fix df_to_postgres_table Remove import Utils Signed-off-by: Job Almekinders <job.almekinders@teampicnic.com>
feast-dev · HaoXuAI · Jul 1, 2024 · Jun 19, 2024 · Jun 19, 2024 · Jun 19, 2024
commit 0a9bcedcd4d4b82ab9bebcb45b949b28a6fc9d75
@@ -1,50 +1,59 @@
-from typing import Dict
+from typing import Any, Dict
 
 import numpy as np
 import pandas as pd
-import psycopg2
-import psycopg2.extras
+import psycopg
 import pyarrow as pa
-from psycopg2.pool import SimpleConnectionPool
+from psycopg.connection import Connection
+from psycopg_pool import ConnectionPool
 
 from feast.infra.utils.postgres.postgres_config import PostgreSQLConfig
 from feast.type_map import arrow_to_pg_type
 
 
-def _get_conn(config: PostgreSQLConfig):
-    conn = psycopg2.connect(
-        dbname=config.database,
-        host=config.host,
-        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),
+def _get_conn(config: PostgreSQLConfig) -> Connection:
+    """Get a psycopg `Connection`."""
+    conn = psycopg.connect(
+        conninfo=_get_conninfo(config),
         keepalives_idle=config.keepalives_idle,
+        **_get_conn_kwargs(config),
     )
     return conn
 
 
-def _get_connection_pool(config: PostgreSQLConfig):
-    return SimpleConnectionPool(
-        config.min_conn,
-        config.max_conn,
-        dbname=config.database,
-        host=config.host,
-        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),
+def _get_connection_pool(config: PostgreSQLConfig) -> ConnectionPool:
+    """Get a psycopg `ConnectionPool`."""
+    return ConnectionPool(
+        conninfo=_get_conninfo(config),
+        min_size=config.min_conn,
+        max_size=config.max_conn,
+        open=False,
+        kwargs=_get_conn_kwargs(config),
     )
 
 
+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}"
+    )
+
+
+def _get_conn_kwargs(config: PostgreSQLConfig) -> Dict[str, Any]:
+    """Get the additional `kwargs` required for connection objects."""
+    return {
+        "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),
+    }
+
+
 def _df_to_create_table_sql(entity_df, table_name) -> str:
     pa_table = pa.Table.from_pandas(entity_df)
     columns = [
@@ -63,16 +72,14 @@ def df_to_postgres_table(
     """
     Create a table for the data frame, insert all the values, and return the table schema
     """
+    nr_columns = df.shape[1]
+    placeholders = ", ".join(["%s"] * nr_columns)
+    query = f"INSERT INTO {table_name} VALUES ({placeholders})"
+    values = df.replace({np.NaN: None}).to_numpy().tolist()
+
     with _get_conn(config) as conn, conn.cursor() as cur:
         cur.execute(_df_to_create_table_sql(df, table_name))
-        psycopg2.extras.execute_values(
-            cur,
-            f"""
-            INSERT INTO {table_name}
-            VALUES %s
-            """,
-            df.replace({np.NaN: None}).to_numpy(),
-        )
+        cur.executemany(query, values)
         return dict(zip(df.columns, df.dtypes))