fix(postgres): use psycopg2.sql.Identifier/SQL for table/sequence names to prevent SQL injection and satisfy PR feedback

src/vector_db/postgres_vector_service.py

@@ -10,6 +10,7 @@ from typing import Any, Dict, List, Optional
 
 import psycopg2
 import psycopg2.extras
+from psycopg2 import sql
 
 logger = logging.getLogger(__name__)
 
@@ -68,8 +69,9 @@ class PostgresVectorService:
 
                 # Create table with initial structure (dimension will be added later)
                 cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS {table_name} (
+                    sql.SQL(
+                        """
+                    CREATE TABLE IF NOT EXISTS {} (
                         id SERIAL PRIMARY KEY,
                         content TEXT NOT NULL,
                         embedding vector,
@@ -77,18 +79,18 @@ class PostgresVectorService:
                         created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                         updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                     );
-                """.format(
-                        table_name=self.table_name
-                    )
+                """
+                    ).format(sql.Identifier(self.table_name))
                 )
 
                 # Create index for text search
                 cur.execute(
-                    """
-                    CREATE INDEX IF NOT EXISTS idx_{table}_content
-                    ON {table} USING gin(to_tsvector('english', content));
-                """.format(
-                        table=self.table_name
+                    sql.SQL(
+                        "CREATE INDEX IF NOT EXISTS {} "
+                        "ON {} USING gin(to_tsvector('english', content));"
+                    ).format(
+                        sql.Identifier(f"idx_{self.table_name}_content"),
+                        sql.Identifier(self.table_name),
                     )
                 )
 
@@ -123,28 +125,29 @@ class PostgresVectorService:
                 if result and ("vector(%s)" % dimension) not in str(result):
                     # Drop existing index if it exists
                     cur.execute(
-                        "DROP INDEX IF EXISTS idx_{table}_embedding_cosine".format(
-                            table=self.table_name
+                        sql.SQL("DROP INDEX IF EXISTS {}; ").format(
+                            sql.Identifier(f"idx_{self.table_name}_embedding_cosine")
                         )
                     )
 
                     # Alter column to correct dimension
                     cur.execute(
-                        "ALTER TABLE {table} "
-                        "ALTER COLUMN embedding TYPE vector({dim})".format(
-                            table=self.table_name, dim=dimension
+                        sql.SQL(
+                            "ALTER TABLE {} ALTER COLUMN embedding TYPE vector({});"
+                        ).format(
+                            sql.Identifier(self.table_name), sql.Literal(dimension)
                         )
                     )
 
                     # Create optimized index for similarity search
                     cur.execute(
-                        """
-                        CREATE INDEX IF NOT EXISTS idx_{table}_embedding_cosine
-                        ON {table}
-                        USING ivfflat (embedding vector_cosine_ops)
-                        WITH (lists = 100);
-                    """.format(
-                            table=self.table_name
+                        sql.SQL(
+                            "CREATE INDEX IF NOT EXISTS {} ON {} "
+                            "USING ivfflat (embedding vector_cosine_ops) "
+                            "WITH (lists = 100);"
+                        ).format(
+                            sql.Identifier(f"idx_{self.table_name}_embedding_cosine"),
+                            sql.Identifier(self.table_name),
                         )
                     )
 
@@ -192,19 +195,13 @@ class PostgresVectorService:
         with self._get_connection() as conn:
             with conn.cursor() as cur:
                 for text, embedding, metadata in zip(texts, embeddings, metadatas):
-                    # Insert document and get ID
+                    # Insert document and get ID (table name composed safely)
                     cur.execute(
-                        """
-                        INSERT INTO %s (content, embedding, metadata)
-                        VALUES (%s, %s, %s)
-                        RETURNING id;
-                    """,
-                        (
-                            self.table_name,
-                            text,
-                            embedding,
-                            psycopg2.extras.Json(metadata),
-                        ),
+                        sql.SQL(
+                            "INSERT INTO {} (content, embedding, metadata) "
+                            "VALUES (%s, %s, %s) RETURNING id;"
+                        ).format(sql.Identifier(self.table_name)),
+                        (text, embedding, psycopg2.extras.Json(metadata)),
                     )
 
                     doc_id = cur.fetchone()[0]
@@ -254,14 +251,18 @@ class PostgresVectorService:
             if conditions:
                 where_clause = "WHERE " + " AND ".join(conditions)
 
-        query = f"""
+        # Compose query safely with identifier for table name. where_clause
+        # contains only parameter placeholders (%s) and logical operators.
+        query = sql.SQL(
+            """
             SELECT id, content, metadata,
                    1 - (embedding <=> %s) as similarity_score
-            FROM {self.table_name}
-            {where_clause}
+            FROM {}
+            {}
             ORDER BY embedding <=> %s
             LIMIT %s;
         """
+        ).format(sql.Identifier(self.table_name), sql.SQL(where_clause))
 
         with self._get_connection() as conn:
             with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
@@ -283,15 +284,18 @@ class PostgresVectorService:
         with self._get_connection() as conn:
             with conn.cursor() as cur:
                 # Get document count
-                cur.execute("SELECT COUNT(*) FROM %s", (self.table_name,))
+                cur.execute(
+                    sql.SQL("SELECT COUNT(*) FROM {};").format(
+                        sql.Identifier(self.table_name)
+                    )
+                )
                 doc_count = cur.fetchone()[0]
 
                 # Get table size
                 cur.execute(
-                    """
-                    SELECT pg_size_pretty(pg_total_relation_size(%s)) as size;
-                """,
-                    (self.table_name,),
+                    sql.SQL(
+                        "SELECT pg_size_pretty(pg_total_relation_size({})) as size;"
+                    ).format(sql.Identifier(self.table_name))
                 )
                 table_size = cur.fetchone()[0]
 
@@ -335,10 +339,9 @@ class PostgresVectorService:
                 int_ids = [int(doc_id) for doc_id in document_ids]
 
                 cur.execute(
-                    f"""
-                    DELETE FROM {self.table_name}
-                    WHERE id = ANY(%s)
-                """,
+                    sql.SQL("DELETE FROM {} WHERE id = ANY(%s);").format(
+                        sql.Identifier(self.table_name)
+                    ),
                     (int_ids,),
                 )
 
@@ -357,14 +360,22 @@ class PostgresVectorService:
         """
         with self._get_connection() as conn:
             with conn.cursor() as cur:
-                cur.execute("SELECT COUNT(*) FROM %s", (self.table_name,))
+                cur.execute(
+                    sql.SQL("SELECT COUNT(*) FROM {};").format(
+                        sql.Identifier(self.table_name)
+                    )
+                )
                 count_before = cur.fetchone()[0]
 
-                cur.execute("DELETE FROM %s", (self.table_name,))
+                cur.execute(
+                    sql.SQL("DELETE FROM {};").format(sql.Identifier(self.table_name))
+                )
 
                 # Reset the sequence
                 cur.execute(
-                    "ALTER SEQUENCE %s_id_seq RESTART WITH 1", (self.table_name,)
+                    sql.SQL("ALTER SEQUENCE {} RESTART WITH 1;").format(
+                        sql.Identifier(f"{self.table_name}_id_seq")
+                    )
                 )
 
                 conn.commit()
@@ -411,11 +422,10 @@ class PostgresVectorService:
         updates.append("updated_at = CURRENT_TIMESTAMP")
         params.append(int(document_id))
 
-        query = f"""
-            UPDATE {self.table_name}
-            SET {', '.join(updates)}
-            WHERE id = %s
-        """
+        # Compose update query with safe identifier for the table name.
+        query = sql.SQL(
+            "UPDATE {} SET " + ", ".join(updates) + " WHERE id = %s"
+        ).format(sql.Identifier(self.table_name))
 
         with self._get_connection() as conn:
             with conn.cursor() as cur:
@@ -443,12 +453,11 @@ class PostgresVectorService:
         with self._get_connection() as conn:
             with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                 cur.execute(
-                    """
-                    SELECT id, content, metadata, created_at, updated_at
-                    FROM %s
-                    WHERE id = %s
-                """,
-                    (self.table_name, int(document_id)),
+                    sql.SQL(
+                        "SELECT id, content, metadata, created_at, "
+                        "updated_at FROM {} WHERE id = %s;"
+                    ).format(sql.Identifier(self.table_name)),
+                    (int(document_id),),
                 )
 
                 row = cur.fetchone()