Address code review feedback: fix SQL injection, QueueMetrics bug, reduce boilerplate

pgmq_sqlalchemy/operation.py

@@ -1,12 +1,12 @@
 from typing import List, Optional, Tuple, Dict, Any, Union
 import re
+import json
 
 from sqlalchemy import text
 from sqlalchemy.orm import Session
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from .schema import Message, QueueMetrics
-from ._utils import encode_dict_to_psql, encode_list_to_psql
 
 
 class PGMQOperation:
@@ -106,19 +106,31 @@ def _get_list_queues_statement() -> Tuple[str, Dict[str, Any]]:
         return "select queue_name from pgmq.list_queues();", {}
 
     @staticmethod
-    def _get_send_statement(queue_name: str, message: dict, delay: int) -> str:
-        """Get statement for send (no params, using f-string)."""
-        encoded_message = encode_dict_to_psql(message)
-        return f"select * from pgmq.send('{queue_name}',{encoded_message},{delay});"
+    def _get_send_statement(
+        queue_name: str, message: dict, delay: int
+    ) -> Tuple[str, Dict[str, Any]]:
+        """Get statement and params for send."""
+        return (
+            "select * from pgmq.send(:queue_name, :message::jsonb, :delay);",
+            {
+                "queue_name": queue_name,
+                "message": json.dumps(message),
+                "delay": delay,
+            },
+        )
 
     @staticmethod
     def _get_send_batch_statement(
         queue_name: str, messages: List[dict], delay: int
-    ) -> str:
-        """Get statement for send_batch (no params, using f-string)."""
-        encoded_messages = encode_list_to_psql(messages)
+    ) -> Tuple[str, Dict[str, Any]]:
+        """Get statement and params for send_batch."""
         return (
-            f"select * from pgmq.send_batch('{queue_name}',{encoded_messages},{delay});"
+            "select * from pgmq.send_batch(:queue_name, :messages::jsonb, :delay);",
+            {
+                "queue_name": queue_name,
+                "messages": json.dumps(messages),
+                "delay": delay,
+            },
         )
 
     @staticmethod
@@ -546,8 +558,8 @@ def send(
         Returns:
             The message ID.
         """
-        stmt = PGMQOperation._get_send_statement(queue_name, message, delay)
-        row = session.execute(text(stmt)).fetchone()
+        stmt, params = PGMQOperation._get_send_statement(queue_name, message, delay)
+        row = session.execute(text(stmt), params).fetchone()
         if commit:
             session.commit()
         return row[0]
@@ -573,8 +585,8 @@ async def send_async(
         Returns:
             The message ID.
         """
-        stmt = PGMQOperation._get_send_statement(queue_name, message, delay)
-        row = (await session.execute(text(stmt))).fetchone()
+        stmt, params = PGMQOperation._get_send_statement(queue_name, message, delay)
+        row = (await session.execute(text(stmt), params)).fetchone()
         if commit:
             await session.commit()
         return row[0]
@@ -600,8 +612,8 @@ def send_batch(
         Returns:
             List of message IDs.
         """
-        stmt = PGMQOperation._get_send_batch_statement(queue_name, messages, delay)
-        rows = session.execute(text(stmt)).fetchall()
+        stmt, params = PGMQOperation._get_send_batch_statement(queue_name, messages, delay)
+        rows = session.execute(text(stmt), params).fetchall()
         if commit:
             session.commit()
         return [row[0] for row in rows]
@@ -627,8 +639,8 @@ async def send_batch_async(
         Returns:
             List of message IDs.
         """
-        stmt = PGMQOperation._get_send_batch_statement(queue_name, messages, delay)
-        rows = (await session.execute(text(stmt))).fetchall()
+        stmt, params = PGMQOperation._get_send_batch_statement(queue_name, messages, delay)
+        rows = (await session.execute(text(stmt), params)).fetchall()
         if commit:
             await session.commit()
         return [row[0] for row in rows]
@@ -1250,7 +1262,6 @@ def metrics(
             newest_msg_age_sec=row[2],
             oldest_msg_age_sec=row[3],
             total_messages=row[4],
-            scrape_time=row[5],
         )
 
     @staticmethod
@@ -1282,7 +1293,6 @@ async def metrics_async(
             newest_msg_age_sec=row[2],
             oldest_msg_age_sec=row[3],
             total_messages=row[4],
-            scrape_time=row[5],
         )
 
     @staticmethod
@@ -1313,7 +1323,6 @@ def metrics_all(
                 newest_msg_age_sec=row[2],
                 oldest_msg_age_sec=row[3],
                 total_messages=row[4],
-                scrape_time=row[5],
             )
             for row in rows
         ]
@@ -1346,7 +1355,6 @@ async def metrics_all_async(
                 newest_msg_age_sec=row[2],
                 oldest_msg_age_sec=row[3],
                 total_messages=row[4],
-                scrape_time=row[5],
             )
             for row in rows
         ]