Commit f320b87
Migrate SQA declarative classes to SA 2.0 Mapped[T] + adopt SQLAlchemy 2.0 in bento_kernel (#5205)
Summary:
Pull Request resolved: #5205
Migrates the Ax SA declarative classes from SA 1.x `Column[T]` class-body annotations to SA 2.0 native `Mapped[T]` annotations, keeping `Column(...)` as the runtime constructor (NOT `mapped_column(...)`). This is the SA 1.4-compatible bridging form that gives us the type-narrowing benefits of `Mapped[T]` at downstream callsites while keeping the OSS Ax dual-version contract — runtime works on both SA 1.4 and SA 2.0.
Why not `mapped_column(...)`: `mapped_column` is SA 2.0-only. Several Meta consumers (e.g. ad ranking AutoML/AutoParamFinder targets) still pin SA 1.4 in their PACKAGE files, and the OSS Ax repo also supports SA 1.4. A pure SA 2.0 migration would `ImportError` at module load in those contexts. `Mapped[T] = Column(...)` runs identically under both versions:
- SA 2.0: `Mapped` is a typed descriptor; `__get__` resolves to `T` at instance access; declarative scanner evaluates the annotation via `typing.get_type_hints` and maps the `Column` correctly. `__allow_unmapped__ = True` on `SQABase` enables this hybrid form alongside the strict `mapped_column` form.
- SA 1.4: `Mapped` is importable as a class (yes — SA 1.4.17 exports it from `sqlalchemy.orm`); annotations stay as strings due to `from __future__ import annotations`; SA 1.4's declarative scanner never introspects the annotation, only the `Column(...)` RHS.
Trade-off: under SA 2.0 stubs pyre sees `Column[T]` on the RHS as incompatible with `Mapped[T]` on the LHS, so each declarative class file carries a file-level `# pyre-ignore-all-errors[8]`. The benefit is paid back at every downstream callsite: `experiment_sqa.name` resolves to `str` (not `Column[str]`) for the type-checker, eliminating the cascade of `# pyre-ignore[6]: Column[T] vs plain T param` suppressions that pure-Column class declarations would require. Concretely, the cleanup diff D106016101 removes ~22 inline pyre-ignores that were previously needed because the SA 2.0 typed stubs flagged every callsite that passed `experiment_sqa.X` to a function expecting plain `X`.
The migration also corrects several pre-existing annotation lies: places where the old `Column[T]` annotation was narrower than the runtime `Column(..., nullable=True)` default have been widened to `Mapped[T | None]` to match the actual schema. Nullability rule applied uniformly: source of truth is the runtime `Column(..., nullable=)` flag (with `primary_key=True` implying not-null), NOT the prior annotation. See the contract comment at the top of `ax/storage/sqa_store/sqa_classes.py` for what future edits must respect — `mapped_column`'s automatic nullable inference from `Mapped[T]` does NOT apply here because we're using `Column(...)`, so each new column MUST explicitly pass `nullable=False` (or `primary_key=True`) for `Mapped[T]` non-Optional, and either omit `nullable=` or pass `nullable=True` for `Mapped[T | None]`.
Mapped import uses a `try/except ImportError` guard at module load. SA 2.0 needs `Mapped` in the module namespace at class-definition time (the declarative scanner evaluates the annotation strings); SA 1.4 doesn't introspect annotations, so silently catching the unlikely ImportError keeps the file loadable even in stripped-down SA installations.
Files touched:
- `fbcode/ax/storage/sqa_store/sqa_classes.py`: all 13 declarative classes migrated to `Mapped[T] = Column(...)` form. Removed prior `mapped_column` import. Added file-level `# pyre-ignore-all-errors[8]` with explanation and the future-edit contract.
- `fbcode/ax/fb/storage/sqa_store/sqa_classes.py`: `SQAExperimentFB`'s 2 relationships migrated to `Mapped[list[T]] = relationship(...)`. No file-level [8] needed because `relationship()` returns Mapped-compatible under SA 2.0 stubs. `association_proxy` lines keep their existing `# pyre-ignore[8]` (association_proxy isn't a Mapped attr).
- `fbcode/ax/fb/storage/sqa_store/metadata.py`: 4 classes (ExperimentOwner, ExperimentTag, ExperimentTask, ExperimentOncallRotation) migrated. `hybrid_property.expression` `Column(...)` returns intentionally NOT migrated — they're SQL expressions, not Mapped attrs.
Also bundled (per the prior diff title): bumps the `bento_kernel_pts` package PACKAGE pin to SQLAlchemy 2.0 so the PTS Bento kernel adopts SA 2.0 alongside the rest of `fbcode/ax/`.
Differential Revision: D105247335
fbshipit-source-id: 207d8018fbf5115694f0c9c294512be3e585f96b1 parent 1bf02d0 commit f320b87
2 files changed
Lines changed: 201 additions & 190 deletions
0 commit comments