Harden module installation and imports

src/govoplan_core/db/migrations.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass, replace
 import json
 import os
 from pathlib import Path
+import re
 from typing import Any
 
 from alembic import command
@@ -44,6 +45,7 @@ MIGRATION_TASK_PHASES = (*PRE_MIGRATION_TASK_PHASES, *POST_MIGRATION_TASK_PHASES
 MIGRATION_TRACK_RELEASE = "release"
 MIGRATION_TRACK_DEV = "dev"
 MIGRATION_TRACKS = (MIGRATION_TRACK_RELEASE, MIGRATION_TRACK_DEV)
+_SQL_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
 
 _NAMESPACE_TABLE_RENAMES = (
     ("tenants", "tenancy_tenants"),
@@ -521,19 +523,25 @@ def _backfill_user_lock_state_for_create_all_schema(database_url: str) -> None:
 
 
 def _row_count(connection, table_name: str) -> int:
-    quoted = connection.dialect.identifier_preparer.quote(table_name)
-    return int(connection.execute(text(f"SELECT COUNT(*) FROM {quoted}")).scalar_one())
+    quoted = _quoted_table_name(connection, table_name)
+    return int(connection.execute(text(f"SELECT COUNT(*) FROM {quoted}")).scalar_one())  # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
 
 
 def _drop_table(connection, table_name: str) -> None:
-    quoted = connection.dialect.identifier_preparer.quote(table_name)
-    connection.execute(text(f"DROP TABLE {quoted}"))
+    quoted = _quoted_table_name(connection, table_name)
+    connection.execute(text(f"DROP TABLE {quoted}"))  # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
 
 
 def _rename_table(connection, old_name: str, new_name: str) -> None:
-    quoted_old = connection.dialect.identifier_preparer.quote(old_name)
-    quoted_new = connection.dialect.identifier_preparer.quote(new_name)
-    connection.execute(text(f"ALTER TABLE {quoted_old} RENAME TO {quoted_new}"))
+    quoted_old = _quoted_table_name(connection, old_name)
+    quoted_new = _quoted_table_name(connection, new_name)
+    connection.execute(text(f"ALTER TABLE {quoted_old} RENAME TO {quoted_new}"))  # nosemgrep: python.sqlalchemy.security.audit.avoid-sqlalchemy-text.avoid-sqlalchemy-text
+
+
+def _quoted_table_name(connection, table_name: str) -> str:
+    if not _SQL_IDENTIFIER_RE.fullmatch(table_name):
+        raise ValueError(f"Unsafe table identifier: {table_name!r}")
+    return connection.dialect.identifier_preparer.quote(table_name)
 
 
 def _reconcile_scope_table_names(connection, tables: set[str]) -> bool: