diff --git a/src/amp/loaders/implementations/postgresql_loader.py b/src/amp/loaders/implementations/postgresql_loader.py index 2591b52..adcb2a0 100644 --- a/src/amp/loaders/implementations/postgresql_loader.py +++ b/src/amp/loaders/implementations/postgresql_loader.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +import re from typing import Any, Dict, List, Optional, Union import pyarrow as pa @@ -485,6 +486,10 @@ def _handle_reorg(self, invalidation_ranges: List[BlockRange], table_name: str, conn = self.pool.getconn() try: with conn.cursor() as cur: + # Validate table_name to prevent SQL injection + if not re.match(r'^[a-zA-Z0-9_]+$', str(table_name)): + raise ValueError("Invalid input") + # Build list of unique IDs to delete unique_batch_ids = list(set(bid.unique_id for bid in all_affected_batch_ids))