Migrate SQA declarative classes to SA 2.0 Mapped[T] + adopt SQLAlchemy 2.0 in bento_kernel (#5205)#5205
Open
LeoMoonStar wants to merge 2 commits into
Open
Migrate SQA declarative classes to SA 2.0 Mapped[T] + adopt SQLAlchemy 2.0 in bento_kernel (#5205)#5205LeoMoonStar wants to merge 2 commits into
LeoMoonStar wants to merge 2 commits into
Conversation
|
@LeoMoonStar has exported this pull request. If you are a Meta employee, you can view the originating Diff in D105247335. |
added 2 commits
May 22, 2026 00:12
Summary: The Meta-internal Ax storage extensions in ax/fb/storage/ have two SA 2.0 incompatibilities not present in the OSS surface: a raw SQL string passed to session.execute in fb sqa_store db.py (SA 2.0 requires text() wrapping), and external_store.py uses Connection.execute() for writes without explicit transaction (SA 2.0 removed implicit autocommit, so writes were silently rolling back), uses string-keyed Row indexing (SA 2.0 requires row._mapping[key]), and consumes a Result generator outside the connection context (SA 2.0 closes the Result on connection close). This diff wraps SHOW DATABASES with text(), switches _write to engine.begin() for transactional commit, migrates _decode_row to row._mapping access, and materializes the read_raw_data result list inside the with conn block. Adds tests_sa2 dual-version Buck targets for fb sqa_store, fb external_store, and fb prod_tests, plus a SQLAlchemy2CompatTest smoke test that exercises the libfb.py.db_locator -> creator -> engine -> session -> SELECT 1 path and asserts EXPECTED_SA_MAJOR. Reviewed By: mgarrard, yangjoanna Differential Revision: D104875016
…y 2.0 in bento_kernel (facebook#5205) Summary: 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
ed6cf7d to
b4a5d53
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #5205 +/- ##
==========================================
- Coverage 96.38% 96.38% -0.01%
==========================================
Files 617 617
Lines 69615 69619 +4
==========================================
+ Hits 67100 67102 +2
- Misses 2515 2517 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Migrates the Ax SA declarative classes from SA 1.x
Column[T]class-body annotations to SA 2.0 nativeMapped[T]annotations, keepingColumn(...)as the runtime constructor (NOTmapped_column(...)). This is the SA 1.4-compatible bridging form that gives us the type-narrowing benefits ofMapped[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_columnis 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 wouldImportErrorat module load in those contexts.Mapped[T] = Column(...)runs identically under both versions:Mappedis a typed descriptor;__get__resolves toTat instance access; declarative scanner evaluates the annotation viatyping.get_type_hintsand maps theColumncorrectly.__allow_unmapped__ = TrueonSQABaseenables this hybrid form alongside the strictmapped_columnform.Mappedis importable as a class (yes — SA 1.4.17 exports it fromsqlalchemy.orm); annotations stay as strings due tofrom __future__ import annotations; SA 1.4's declarative scanner never introspects the annotation, only theColumn(...)RHS.Trade-off: under SA 2.0 stubs pyre sees
Column[T]on the RHS as incompatible withMapped[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.nameresolves tostr(notColumn[str]) for the type-checker, eliminating the cascade of# pyre-ignore[6]: Column[T] vs plain T paramsuppressions 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 passedexperiment_sqa.Xto a function expecting plainX.The migration also corrects several pre-existing annotation lies: places where the old
Column[T]annotation was narrower than the runtimeColumn(..., nullable=True)default have been widened toMapped[T | None]to match the actual schema. Nullability rule applied uniformly: source of truth is the runtimeColumn(..., nullable=)flag (withprimary_key=Trueimplying not-null), NOT the prior annotation. See the contract comment at the top ofax/storage/sqa_store/sqa_classes.pyfor what future edits must respect —mapped_column's automatic nullable inference fromMapped[T]does NOT apply here because we're usingColumn(...), so each new column MUST explicitly passnullable=False(orprimary_key=True) forMapped[T]non-Optional, and either omitnullable=or passnullable=TrueforMapped[T | None].Mapped import uses a
try/except ImportErrorguard at module load. SA 2.0 needsMappedin 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 toMapped[T] = Column(...)form. Removed priormapped_columnimport. 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 toMapped[list[T]] = relationship(...). No file-level [8] needed becauserelationship()returns Mapped-compatible under SA 2.0 stubs.association_proxylines 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.expressionColumn(...)returns intentionally NOT migrated — they're SQL expressions, not Mapped attrs.Also bundled (per the prior diff title): bumps the
bento_kernel_ptspackage PACKAGE pin to SQLAlchemy 2.0 so the PTS Bento kernel adopts SA 2.0 alongside the rest offbcode/ax/.Differential Revision: D105247335