security(sql): validate SQL identifiers and queries to prevent injection

modin/db_conn.py

@@ -23,6 +23,7 @@
 driver or a worker wants one.
 """
 
+import re
 from typing import Any, Dict, Optional, Sequence
 
 _PSYCOPG_LIB_NAME = "psycopg2"
@@ -35,6 +36,34 @@ class UnsupportedDatabaseException(Exception):
     pass
 
 
+class InvalidQueryException(Exception):
+    """Raised when a SQL query fails security validation."""
+
+    pass
+
+
+# Regex for a simple SELECT query without comment markers or statement terminators
+_SELECT_QUERY_RE = re.compile(r"^\s*SELECT\s", re.IGNORECASE)
+_QUERY_FORBIDDEN_RE = re.compile(r"[;\-]|/\*|\*/|#", re.IGNORECASE)
+
+
+def _validate_select_query(query: str) -> None:
+    """
+    Validate that ``query`` looks like a single SELECT statement.
+
+    Raises InvalidQueryException if the query contains comment markers,
+    statement terminators, or does not start with SELECT.
+    """
+    if not _SELECT_QUERY_RE.match(query):
+        raise InvalidQueryException(
+            "Query must start with SELECT."
+        )
+    if _QUERY_FORBIDDEN_RE.search(query):
+        raise InvalidQueryException(
+            "Query contains forbidden characters (comment markers or statement terminators)."
+        )
+
+
 class ModinDatabaseConnection:
     """
     Creates a SQL database connection.
@@ -134,6 +163,7 @@ def column_names_query(self, query: str) -> str:
         -------
         str
         """
+        _validate_select_query(query)
         # This query looks odd, but it works in both PostgreSQL and Microsoft
         # SQL, which doesn't let you use a "limit" clause to select 0 rows.
         return f"SELECT * FROM ({query}) AS _MODIN_COUNT_QUERY WHERE 1 = 0"
@@ -151,6 +181,7 @@ def row_count_query(self, query: str) -> str:
         -------
         str
         """
+        _validate_select_query(query)
         return f"SELECT COUNT(*) FROM ({query}) AS _MODIN_COUNT_QUERY"
 
     def partition_query(self, query: str, limit: int, offset: int) -> str:
@@ -170,6 +201,7 @@ def partition_query(self, query: str, limit: int, offset: int) -> str:
         -------
         str
         """
+        _validate_select_query(query)
         return (
             (
                 f"SELECT * FROM ({query}) AS _MODIN_COUNT_QUERY ORDER BY(SELECT NULL)"

modin/experimental/core/io/sql/utils.py

@@ -13,10 +13,30 @@
 
 """Utilities for experimental SQL format type IO functions implementations."""
 
+import re
+
 import pandas
 import pandas._libs.lib as lib
 from sqlalchemy import MetaData, Table, create_engine, inspect, text
 
+# Regex for a valid SQL identifier (alphanumeric + underscore, no leading digit)
+_SQL_IDENTIFIER_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
+
+
+def _validate_sql_identifier(name: str, field_name: str = "identifier") -> None:
+    """
+    Validate that ``name`` is a safe SQL identifier.
+
+    Raises ValueError if the name contains characters outside [a-zA-Z0-9_]
+    or starts with a digit.
+    """
+    if not _SQL_IDENTIFIER_RE.match(name):
+        raise ValueError(
+            f"Invalid SQL {field_name}: {name!r}. "
+            f"Only alphanumeric characters and underscores are allowed, "
+            f"and the first character must be a letter or underscore."
+        )
+
 from modin.core.storage_formats.pandas.parsers import _split_result_for_readers
 
 
@@ -131,6 +151,7 @@ def build_query_from_table(name):
     str
         Query string.
     """
+    _validate_sql_identifier(name, field_name="table name")
     return "SELECT * FROM {0}".format(name)
 
 
@@ -252,6 +273,7 @@ def query_put_bounders(query, partition_column, start, end):  # pragma: no cover
     str
         Query string with boundaries.
     """
+    _validate_sql_identifier(partition_column, field_name="partition_column")
     where = " WHERE TMP_TABLE.{0} >= {1} AND TMP_TABLE.{0} <= {2}".format(
         partition_column, start, end
     )