Skip to content

Commit a401d37

Browse files
committed
refactor(storage/sql): use SQLAlchemy Core for count_by_column
Replace the f-string + text() construction with select/func.count over a reflected Table and named column. Eliminates the raw SQL build entirely so static analysis no longer has a string-interpolated SQL statement to flag, and the allowlist check on `column` remains in place as defence-in-depth against future callers passing arbitrary identifiers.
1 parent 5492894 commit a401d37

1 file changed

Lines changed: 15 additions & 11 deletions

File tree

src/orb/infrastructure/storage/sql/strategy.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@
33
from contextlib import contextmanager
44
from typing import Any, Optional
55

6-
from sqlalchemy import text
6+
from sqlalchemy import MetaData, Table
7+
from sqlalchemy import column as sa_column
8+
from sqlalchemy import func, select, text
79

810
from orb.infrastructure.logging.logger import get_logger
911
from orb.infrastructure.storage.base.strategy import BaseStorageStrategy
@@ -411,24 +413,26 @@ def count_by_column(self, column: str) -> dict[str, int]:
411413
gracefully to the list-and-count slow path if needed.
412414
"""
413415
# Validate the column against the strategy's registered columns dict
414-
# before interpolation. SQL identifiers cannot be bound parameters, so
415-
# the only safe path is allowlist-then-interpolate. Anything else is
416-
# rejected.
416+
# before building the query. This is the only place untrusted strings
417+
# could ever enter the SQL build; rejecting unknown columns keeps the
418+
# query construction below restricted to identifiers we registered at
419+
# construction time.
417420
if column not in self.columns:
418421
raise StorageError(
419422
f"count_by_column: column {column!r} is not in the registered "
420423
f"schema for table {self.table_name!r}"
421424
)
422-
# table_name and columns keys are constructor-time constants set by the
423-
# repository layer; both have just been validated against the in-memory
424-
# schema map above. No request-time string crosses this boundary.
425-
sql_statement = text(
426-
f"SELECT {column} AS bucket, COUNT(*) AS cnt FROM {self.table_name} GROUP BY {column}"
427-
)
425+
# Build the SELECT via SQLAlchemy Core constructs — no raw SQL string
426+
# interpolation. The column object is created by name (validated
427+
# above) and the Table is reflected from MetaData by name (also
428+
# validated since self.table_name is a constructor-time constant).
429+
bucket = sa_column(column)
430+
table = Table(self.table_name, MetaData())
431+
stmt = select(bucket.label("bucket"), func.count().label("cnt")).select_from(table).group_by(bucket)
428432
with self.lock_manager.read_lock():
429433
try:
430434
with self.connection_manager.get_session() as session:
431-
result = session.execute(sql_statement)
435+
result = session.execute(stmt)
432436
rows = result.fetchall()
433437
counts: dict[str, int] = {}
434438
for row in rows:

0 commit comments

Comments
 (0)