Remove redundant psycopg2 dependency

local_database/create_database.py

@@ -1,11 +1,9 @@
-import argparse
 import os
-import subprocess
 
-import psycopg2
-from psycopg2 import sql
+import psycopg
+from psycopg import sql
 
-from local_database.constants import LOCAL_SOURCE_COLLECTOR_DB_NAME, RESTORE_SH_DOCKER_PATH
+from local_database.constants import LOCAL_SOURCE_COLLECTOR_DB_NAME
 
 # Defaults (can be overridden via environment variables)
 POSTGRES_HOST = os.getenv("POSTGRES_HOST", "host.docker.internal")
@@ -16,15 +14,13 @@
 
 # Connect to the default 'postgres' database to create other databases
 def connect(database="postgres", autocommit=True):
-    conn = psycopg2.connect(
+    conn = psycopg.connect(
         dbname=database,
         user=POSTGRES_USER,
         password=POSTGRES_PASSWORD,
         host=POSTGRES_HOST,
         port=POSTGRES_PORT
     )
-    if autocommit:
-        conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
     return conn
 
 def create_database(db_name):
@@ -43,7 +39,7 @@ def create_database(db_name):
         try:
             cur.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(db_name)))
             print(f"✅ Created database: {db_name}")
-        except psycopg2.errors.DuplicateDatabase:
+        except psycopg.errors.DuplicateDatabase:
             print(f"⚠️  Database {db_name} already exists")
         except Exception as e:
             print(f"❌ Failed to create {db_name}: {e}")

pyproject.toml

@@ -23,7 +23,6 @@ dependencies = [
     "openai~=1.60.1",
     "pdap-access-manager==0.3.5",
     "playwright~=1.49.1",
-    "psycopg2-binary~=2.9.6",
     "psycopg[binary]~=3.1.20",
     "pydantic~=2.11.3",
     "pyjwt~=2.10.1",

util/db_manager.py

@@ -1,13 +1,13 @@
 import os
 
-import psycopg2
+import psycopg
 from dotenv import load_dotenv
 
 
 class DBManager:
 
     def __init__(self, db_name, user, password, host, port):
-        self.conn = psycopg2.connect(
+        self.conn = psycopg.connect(
             dbname=db_name,
             user=user,
             password=password,
@@ -41,6 +41,6 @@ def close(self):
     # Note: This is test code to evaluate whether the connection url works. Will be removed in final version.
     load_dotenv()
     conn_url = os.getenv("DIGITAL_OCEAN_DB_CONNECTION_URL")
-    conn = psycopg2.connect(conn_url)
+    conn = psycopg.connect(conn_url)
 
     pass