|
3 | 3 | from contextlib import contextmanager |
4 | 4 | from typing import Any, Optional |
5 | 5 |
|
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 |
7 | 9 |
|
8 | 10 | from orb.infrastructure.logging.logger import get_logger |
9 | 11 | from orb.infrastructure.storage.base.strategy import BaseStorageStrategy |
@@ -411,24 +413,26 @@ def count_by_column(self, column: str) -> dict[str, int]: |
411 | 413 | gracefully to the list-and-count slow path if needed. |
412 | 414 | """ |
413 | 415 | # 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. |
417 | 420 | if column not in self.columns: |
418 | 421 | raise StorageError( |
419 | 422 | f"count_by_column: column {column!r} is not in the registered " |
420 | 423 | f"schema for table {self.table_name!r}" |
421 | 424 | ) |
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) |
428 | 432 | with self.lock_manager.read_lock(): |
429 | 433 | try: |
430 | 434 | with self.connection_manager.get_session() as session: |
431 | | - result = session.execute(sql_statement) |
| 435 | + result = session.execute(stmt) |
432 | 436 | rows = result.fetchall() |
433 | 437 | counts: dict[str, int] = {} |
434 | 438 | for row in rows: |
|
0 commit comments