You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "amoghrajesh (via GitHub)" <gi...@apache.org> on 2026/04/24 09:12:52 UTC

[PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

amoghrajesh opened a new pull request, #65759:
URL: https://github.com/apache/airflow/pull/65759

    <!-- SPDX-License-Identifier: Apache-2.0
         https://www.apache.org/licenses/LICENSE-2.0 -->
   
   <!--
   Thank you for contributing!
   
   Please provide above a brief description of the changes made in this pull request.
   Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
   
   Please make sure that your code changes are covered with tests.
   And in case of new features or big changes remember to adjust the documentation.
   
   Feel free to ping (in general) for the review if you do not see reaction for a few days
   (72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
   
   In case of an existing issue, reference it using one of the following:
   
   * closes: #ISSUE
   * related: #ISSUE
   -->
   
   ---
   
   ##### Was generative AI tooling used to co-author this PR?
   
   <!--
   If generative AI tooling has been used in the process of authoring this PR, please
   change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
   -->
   
   - [x] Yes (please specify the tool below)
   
   <!--
   Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
   -->
   
   # AIP-103 PR 1: DB Schema + Backend Abstraction
   
   This PR lays the foundation for AIP-103 (Task State Management) — a key-value state store scoped to task instances and assets. It implements the DB layer, the backend abstraction, and config — everything needed to persist and retrieve state before any API or SDK work begins.
   
   ---
   
   ## What's in this PR
   
   ### 1. ORM Models
   
   Two new models: `TaskStateModel` and `AssetStateModel` in `airflow/models/`.
   
   `task_state` PK is `(dag_run_id, task_id, map_index, key)`:
   - `dag_run_id` is an integer FK to `dag_run.id` with CASCADE DELETE — task state is automatically wiped when a run is deleted, no orphan cleanup job needed.
   - `dag_id` and `run_id` are denormalized strings alongside `dag_run_id`. The integer FK is for the JOIN and CASCADE; the strings are for fast Execution API lookups without a join and for surfacing `updated_by_run` in the Core API response.
   - `map_index` (default `-1`) gives each mapped task instance its own namespace. Without it, all instances of the same mapped task share one row per key — instance 0 writes `remote_job_id=app_001`, instance 1 overwrites it with `app_002`, instance 0 retries and reconnects to the wrong job.
   - `updated_by_try` is nullable — it has no meaning for asset state and is absent there.
   
   `asset_state` PK is `(asset_id, key)`:
   - `asset_id` FK to `asset.id` with CASCADE DELETE.
   - No `dag_run_id` — asset state is intentionally run-independent. A watermark written in run 1 is readable by run 2. This is what distinguishes it from task state.
   
   Two tables instead of one: the GC lifecycles, index shapes, and FK targets are entirely different for the two scopes. A single discriminated table would mean permanently-null columns on every row and would conflate two very different cleanup models.
   
   What the schema implies:
   - Retries within the same run share state — same `(dag_run_id, task_id, map_index)` will have same rows. Try 3 reads `remote_job_id` written by previous try and reconnects. This is the core AIP-103 use case.
   - Each new DAG run starts with a clean task state namespace — no cross-run leakage via task state.
   - Null values are not representable — absence is row non-existence. "Not set" and "set to null" are not ambiguous.
   
   ### 2. Alembic Migration
   
   Migration `0112_3_3_0_add_task_state_and_asset_state` creates both tables.
   
   ### 3. BaseStateBackend in `shared/state`
   
   `TaskScope`, `AssetScope`, `StateScope`, and `BaseStateBackend` live in a new shared distribution (`apache-airflow-shared-state`) rather than `airflow-core`. This is the only placement that avoids circular dependencies across all three consumers: provider packages (e.g. S3/GCS custom backends) cannot depend on `airflow-core`; `task-sdk` needs the scope dataclasses for context accessors in PR 3; `airflow-core` needs them for `MetastoreStateBackend`.
   
   The interface takes a single `scope` parameter — the backend routes internally on `isinstance(scope, TaskScope)`. Async variants (`aget`, `aset`, `adelete`, `aclear`) are provided as defaults that delegate to the sync methods, so sync backends get async support for free. These are required for AIP-98 compatibility — async tasks must be able to call `await task_state.aset(...)` without blocking the event loop.
   
   The `[state]` config section (`backend` key) and `resolve_state_backend()` follow the `resolve_xcom_backend()` pattern — validate the configured class is a `BaseStateBackend` subclass before returning it.
   
   ### 4. MetastoreStateBackend
   
   Default implementation in `airflow/state/metastore.py`. Routes on scope type, resolves `dag_run_id` from `(dag_id, run_id)` before writing task state, and follows the standard `@provide_session` / `NEW_SESSION` pattern used throughout `airflow-core`.
   
   ---
   
   ## Next steps
   
   - PR 2: Execution API + Core API endpoints
   - PR 3: Task SDK — comms, supervisor, context accessors
   - PR 4: GC and retention config
   
   
   ---
   
   * Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
   * For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
   * When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
   * For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3152703316


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   It's either:
   
   ```ini
   [state]
   backend = someclass
   other = option
   ```
   or
   
   ```ini
   [core]
   state_backend = someclass
   state_other = option
   ```
   right?
   
   I think I prefer the dedicated state section (though I do find "state" a little bit generic in both cases as so much in Airflow is about/has state.)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145867282


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)

Review Comment:
   I think that would break the use case AIP-103 is trying to handle. Each retry gets a new `task_instance.id` UUID — verified this empirically too. 11 tries on a task with retries=10, 11 unique UUIDs across task_instance + task_instance_history.
   
   ```
   [Breeze:3.10.19] root@eaabf698afd2:/opt/airflow$  python /opt/airflow/dev/check_ti_uuid_across_retries.py
   2026-04-27T08:23:06.814648Z [info     ] setup plugin alembic.autogenerate.schemas [alembic.runtime.plugins] loc=plugins.py:37
   2026-04-27T08:23:06.814749Z [info     ] setup plugin alembic.autogenerate.tables [alembic.runtime.plugins] loc=plugins.py:37
   2026-04-27T08:23:06.814794Z [info     ] setup plugin alembic.autogenerate.types [alembic.runtime.plugins] loc=plugins.py:37
   2026-04-27T08:23:06.814838Z [info     ] setup plugin alembic.autogenerate.constraints [alembic.runtime.plugins] loc=plugins.py:37
   2026-04-27T08:23:06.814889Z [info     ] setup plugin alembic.autogenerate.defaults [alembic.runtime.plugins] loc=plugins.py:37
   2026-04-27T08:23:06.814968Z [info     ] setup plugin alembic.autogenerate.comments [alembic.runtime.plugins] loc=plugins.py:37
   === task_instance (current row) ===
     id=019dce01-ef14-74ff-9cf9-f028420d8a84  try_number=11  state=failed
   
   === task_instance_history (one row per past try) ===
     id=019dce01-1899-750b-b37e-c98048cc5340  try_number=1  state=failed
     id=019dce01-240f-7fc5-9385-65124bab87bd  try_number=2  state=failed
     id=019dce01-3a7f-771c-9865-3ae6ec04aa90  try_number=3  state=failed
     id=019dce01-51ad-774a-b84e-514d7b6b1f32  try_number=4  state=failed
     id=019dce01-66c4-7daf-afa0-85737a5c34b6  try_number=5  state=failed
     id=019dce01-7d62-78bf-a1b2-c441b09cf032  try_number=6  state=failed
     id=019dce01-922b-7790-8b4b-79f608209506  try_number=7  state=failed
     id=019dce01-aad4-7e5e-8895-eb3353ab81af  try_number=8  state=failed
     id=019dce01-c0d1-7135-b7dc-67ee806c386f  try_number=9  state=failed
     id=019dce01-d809-7f47-8f47-302231b38f2b  try_number=10  state=failed
   
   Total tries: 11  Unique UUIDs: 1
   ```
   
   Take this as example:
   
   - Try 1 runs, writes remote_job_id = "app_001" — stored under UUID `aaa`
   - Try 1 fails
   - Try 2 starts — it has UUID `bbb`
   - Try 2 looks up remote_job_id under UUID `bbb` and finds nothing
   - Try 2 submits a new job instead of reconnecting to app_001
   - With `(dag_run_id, task_id, map_index)` as the key, all retries share the same rows because that triple combination never changes between tries of the same task within a run
   
   This is how that looks:
   
   - Try 1 writes remote_job_id = "app_001" under `(dag_run_id=42, task_id="submit_spark", map_index=-1)`
   - Try 1 fails
   - Try 2 starts — same DAG run, same task, same map index -> same key `(42, "submit_spark", -1)`
   - Try 2 looks up `remote_job_id` → finds "app_001"
   - Try 2 reconnects to the existing job instead of submitting a new one



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145893924


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Hmmm interesting take. Anything under `core` gives somewhat a sense of it being a "core" configuration tied to airflow core. I personally do not like `xcom_backend` placement in `core` and I think maintaining `state` as a separate config can also support the future tasks where we will have more configs come in like retention configurations. I think of state as a distinct subsystem and prefer separate section but its arguable.
   
   Any thoughts @ashb @kaxil @vikramkoka @jroachgolf84 @XD-DENG ?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3148965130


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)

Review Comment:
   OUch, I thought the UUID is stable for the task instance and a new UUID might be generated for history only. Yeah then it is not suitable.... and we need to use the composite key. I thought the UUID is being used also to link XCom, Rendered fields etc.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3139603735


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+

Review Comment:
   `asset_id` is an integer PK in the DB, the execution API will receive it as an integer in the URL path when needed and I dont think there is a serialization boundary needing a string. int is correct. Made that change



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3146183947


##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -0,0 +1,42 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow._shared.state import (
+    AssetScope as AssetScope,
+    BaseStateBackend as BaseStateBackend,
+    StateScope as StateScope,
+    TaskScope as TaskScope,
+)
+
+_DEFAULT_BACKEND = "airflow.state.metastore.MetastoreStateBackend"

Review Comment:
   Handled in [handle silent failure on resolve_state_backend](https://github.com/apache/airflow/pull/65759/commits/1d428fee7b1eac2f919999be3f1abf48aae9774c)



##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)

Review Comment:
   Wow nice catch. It would be a bummer for Mysql with this limit, I attempted to fix it for both task_state.value and asset_state.value by using `Text().with_variant(MEDIUMTEXT, "mysql")`, same pattern as `Variable._val` instead of documenting it and making it harder to use. WDYT?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3146282601


##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state.py:
##########
@@ -0,0 +1,88 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: 8a7b3c2d1e4f
+Revises: 9fabad868fdb
+Create Date: 2026-04-24 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+from sqlalchemy import Column, ForeignKeyConstraint, Integer, String, Text
+
+from airflow._shared.timezones import timezone
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "8a7b3c2d1e4f"

Review Comment:
   Thanks for the catch, looked it up and figured that we have a breeze utility to do it, and I deleted the hand-written migration and regenerated with `breeze generate-migration-file` — new revision is `fde9ed84d07b`.
   
   
   Has been handled in b93c6e51a1



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3162336009


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,351 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.typing_compat import assert_never
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+            case _:
+                assert_never(scope)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+                case _:
+                    assert_never(scope)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:
+        conditions = [
+            TaskStateModel.dag_id == scope.dag_id,
+            TaskStateModel.run_id == scope.run_id,
+            TaskStateModel.task_id == scope.task_id,
+        ]
+        if scope.map_index != -1:

Review Comment:
   Handled in d1316837b2



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141520879


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)
+    updated_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+    updated_by_try: Mapped[int | None] = mapped_column(Integer, nullable=True)

Review Comment:
   Oh nice catch.
   
   `updated_by_try` is never populated in this PR, do dropping it from this PR entirely. I will add a new field
   `updated_by_ti_id` as an UUID in the execution API PR which I will raise next to avoid adding dead weight here.
   
   I have made an internal note regarding this.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141816779


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)

Review Comment:
   Open question: Should we not directly link this table to a task via the Task UUID? That would reduce the number of columns drastically and this is why we added the UUID



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145980225


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   It just falls through, so we might need handling:
   ```
   value = 10
   match value:
       case 1:
           print("one")
       case 2:
           print("two")
   print("done")
   done
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154020748


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Yeah certainly. I prefer going with `state_store` and yes we will have per backend keys coming in next PRs.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3146185578


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)

Review Comment:
   Handled this in [1d24107e93](https://github.com/apache/airflow/pull/65759#discussion_r3146185578)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154098298


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)

Review Comment:
   Thanks, done.



##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: fde9ed84d07b
+Revises: 9fabad868fdb
+Create Date: 2026-04-27 09:29:20.165461
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "fde9ed84d07b"
+down_revision = "9fabad868fdb"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+    """Apply add task_state and asset_state tables."""
+    op.create_table(
+        "asset_state",
+        sa.Column("asset_id", sa.Integer(), nullable=False),
+        sa.Column("key", sa.String(length=512), nullable=False),
+        sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
+        sa.Column("updated_at", UtcDateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["asset_id"], ["asset.id"], name="asset_state_asset_fkey", ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"),
+    )
+    op.create_table(
+        "task_state",
+        sa.Column("dag_run_id", sa.Integer(), nullable=False),
+        sa.Column("task_id", sa.String(length=250), nullable=False),

Review Comment:
   Good catch. Switched both the migration `(task_id, dag_id, run_id columns in task_state)` and the ORM model to use `StringID()`, same pattern as migration number: 0108. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3162294170


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,351 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.typing_compat import assert_never
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+            case _:
+                assert_never(scope)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+                case _:
+                    assert_never(scope)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:
+        conditions = [
+            TaskStateModel.dag_id == scope.dag_id,
+            TaskStateModel.run_id == scope.run_id,
+            TaskStateModel.task_id == scope.task_id,
+        ]
+        if scope.map_index != -1:

Review Comment:
   Going with option 2. Adding an explicit `clear_all_map_indices: bool = False` parameter on clear / aclear, dropping the -1 sentinel instead. Default is now exact match on `scope.map_index` — same as delete, no asymmetry / confusion. If you want a WIDE scoped wipe you have to spell it out: `clear(scope, all_map_indices=True)`. Felt safer than option 1 -> accidental wipes from a forgotten map_index just no-op now instead of nuking every mapped instance.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh closed pull request #65759: AIP-103: Setting up foundation with ORM and Backend Abstraction 
URL: https://github.com/apache/airflow/pull/65759


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145848683


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   Yup. I wonder what those match statements do at runtime when nothing matches. Do we need a default that raises, or does python handle that for us already? 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3139612080


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2987,7 +2987,19 @@ dag_processor:
       type: string
       example: ~
       default: "warning"
-
+state:

Review Comment:
   Haha, handled.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145913335


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:

Review Comment:
   Handled in 3f03dc7492



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:

Review Comment:
   Handled in 3f03dc7492



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145749951


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   Here you go: a14d28e6b7 :)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145966115


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task(self, scope: TaskScope, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+            )
+        )
+
+    def _get_asset(self, scope: AssetScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+                AssetStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_asset(self, scope: AssetScope, key: str, value: str, *, session: Session) -> None:
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["asset_id", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["asset_id", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_asset(self, scope: AssetScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+                AssetStateModel.key == key,
+            )
+        )
+
+    def _clear_asset(self, scope: AssetScope, *, session: Session) -> None:
+        session.execute(
+            delete(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+            )
+        )
+
+    async def _aget_task(self, scope: TaskScope, key: str, *, session: AsyncSession) -> str | None:
+        row = await session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    async def _aset_task(self, scope: TaskScope, key: str, value: str, *, session: AsyncSession) -> None:
+        dag_run_id = await session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        # get_dialect_name expects a sync Session; sync_session is the underlying Session the async wrapper delegates to
+        dialect = get_dialect_name(session.sync_session)

Review Comment:
   Good catch. Added `TestMetastoreStateBackendAsync` covering all async methods for both task and asset scopes against the real metastore by using  `@pytest.mark.asyncio + pytest-asyncio` that handles the event loop, and create_session_async() writes to the same real DB so it is a genuine integration test. 
   
   > session.sync_session.get_bind() for an async-engine-backed session is the kind of thing worth verifying empirically rather than by inspection.
   
   
   The newly written async tests should cover this — `aset` calls `get_dialect_name(session.sync_session)` on a real async-engine-backed session, so if that path misbehaves the tests will catch it empirically



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(

Review Comment:
   Handled in session.sync_session.get_bind() for an async-engine-backed session is the kind of thing worth verifying empirically rather than by inspection.
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145969082


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task(self, scope: TaskScope, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+            )
+        )
+
+    def _get_asset(self, scope: AssetScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+                AssetStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_asset(self, scope: AssetScope, key: str, value: str, *, session: Session) -> None:
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["asset_id", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["asset_id", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_asset(self, scope: AssetScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+                AssetStateModel.key == key,
+            )
+        )
+
+    def _clear_asset(self, scope: AssetScope, *, session: Session) -> None:
+        session.execute(
+            delete(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+            )
+        )
+
+    async def _aget_task(self, scope: TaskScope, key: str, *, session: AsyncSession) -> str | None:
+        row = await session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    async def _aset_task(self, scope: TaskScope, key: str, value: str, *, session: AsyncSession) -> None:
+        dag_run_id = await session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        # get_dialect_name expects a sync Session; sync_session is the underlying Session the async wrapper delegates to
+        dialect = get_dialect_name(session.sync_session)

Review Comment:
   Handled in 2b29b8098c



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145997406


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":

Review Comment:
   Sure, that can be done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #65759:
URL: https://github.com/apache/airflow/pull/65759#issuecomment-4345334515

   Rerunning with full tests


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154266869


##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -0,0 +1,363 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import select
+
+from airflow._shared.timezones import timezone
+from airflow.models.asset import AssetModel
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+from airflow.state import AssetScope, TaskScope, resolve_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+from airflow.utils.session import create_session
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+
+
+@pytest.fixture
+def backend() -> MetastoreStateBackend:
+    return MetastoreStateBackend()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+@pytest.fixture
+def dag_run_committed() -> DagRun:
+    """DagRun committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        run = DagRun(
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            run_type=DagRunType.SCHEDULED,
+            logical_date=timezone.datetime(2026, 4, 24),
+            run_after=timezone.datetime(2026, 4, 24),
+        )
+        session.add(run)
+        session.flush()
+        session.expunge(run)
+    return run
+
+
+@pytest.fixture
+def asset(session: Session) -> AssetModel:
+    a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+    session.add(a)
+    session.flush()
+    return a
+
+
+@pytest.fixture
+def asset_committed() -> AssetModel:
+    """AssetModel committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+        session.add(a)
+        session.flush()
+        session.expunge(a)
+    return a
+
+
+class TestMetastoreStateBackendTaskScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_1234"
+
+    def test_set_twice_overrides_existing_value(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """Calling set twice on the same key updates the value in place."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_try1", session=session)
+        session.flush()
+        backend.set(scope, "job_id", "app_try2", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_try2"
+
+    def test_set_stores_dag_run_id_fk(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """set resolves dag_run_id from (dag_id, run_id) and persists it as the FK."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == DAG_ID,
+                TaskStateModel.task_id == TASK_ID,
+                TaskStateModel.key == "job_id",
+            )
+        )
+        assert row is not None
+        assert row.dag_run_id == dag_run.id
+
+    def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        backend.set(scope, "checkpoint", "step_3", session=session)
+        session.flush()
+
+        backend.delete(scope, "job_id", session=session)
+        session.flush()
+
+        assert backend.get(scope, "job_id", session=session) is None
+        assert backend.get(scope, "checkpoint", session=session) == "step_3"
+
+    def test_delete_is_noop_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.delete(scope, "nonexistent", session=session)
+
+    def test_clear_removes_all_keys(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        """clear removes every key under the given task scope."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        for key in ("job_id", "checkpoint", "watermark"):
+            backend.set(scope, key, f"val_{key}", session=session)
+        session.flush()
+
+        backend.clear(scope, session=session)
+        session.flush()
+
+        for key in ("job_id", "checkpoint", "watermark"):
+            assert backend.get(scope, key, session=session) is None
+
+    def test_map_index_isolation(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        """Mapped task instances with different map_index values have isolated namespaces."""
+        scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0)
+        scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1)
+
+        backend.set(scope0, "job_id", "app_0", session=session)
+        backend.set(scope1, "job_id", "app_1", session=session)
+        session.flush()
+
+        assert backend.get(scope0, "job_id", session=session) == "app_0"
+        assert backend.get(scope1, "job_id", session=session) == "app_1"
+
+    def test_set_raises_for_missing_dag_run(self, session: Session, backend: MetastoreStateBackend):
+        scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID)
+        with pytest.raises(ValueError, match="No DagRun found"):
+            backend.set(scope, "job_id", "app_1234", session=session)
+
+    def test_clear_scoped_to_map_index(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """clear for map_index=0 does not remove state belonging to map_index=1."""
+        scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0)
+        scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1)
+
+        backend.set(scope0, "job_id", "app_0", session=session)
+        backend.set(scope1, "job_id", "app_1", session=session)
+        session.flush()
+
+        backend.clear(scope0, session=session)
+        session.flush()
+
+        assert backend.get(scope0, "job_id", session=session) is None
+        assert backend.get(scope1, "job_id", session=session) == "app_1"
+
+
+class TestMetastoreStateBackendAssetScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        session.flush()
+        assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z"
+
+    def test_set_twice_overwrites_existing_value(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-01T00:00:00Z", session=session)
+        session.flush()
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        session.flush()
+        assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z"
+
+    def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        backend.set(scope, "file_count", "42", session=session)
+        session.flush()
+
+        backend.delete(scope, "watermark", session=session)
+        session.flush()
+
+        assert backend.get(scope, "watermark", session=session) is None
+        assert backend.get(scope, "file_count", session=session) == "42"
+
+    def test_delete_is_noop_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        backend.delete(scope, "nonexistent", session=session)
+
+    def test_clear_removes_all_keys(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        for key in ("watermark", "file_count", "last_error"):
+            backend.set(scope, key, f"val_{key}", session=session)
+        session.flush()
+
+        backend.clear(scope, session=session)
+        session.flush()
+
+        for key in ("watermark", "file_count", "last_error"):
+            assert backend.get(scope, key, session=session) is None
+
+    def test_different_assets_are_isolated(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        asset2 = AssetModel(uri="s3://bucket/other", name="other_asset", group="test")
+        session.add(asset2)
+        session.flush()
+
+        scope1 = AssetScope(asset_id=asset.id)
+        scope2 = AssetScope(asset_id=asset2.id)
+
+        backend.set(scope1, "watermark", "asset1_value", session=session)
+        session.flush()
+
+        assert backend.get(scope2, "watermark", session=session) is None
+
+
+@pytest.mark.asyncio(loop_scope="class")
+class TestMetastoreStateBackendAsync:
+    async def test_aset_and_aget_task_roundtrip(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        result = await backend.aget(scope, "job_id")
+        assert result == "app_async"
+
+    async def test_adelete_task_removes_key(self, backend: MetastoreStateBackend, dag_run_committed: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        await backend.adelete(scope, "job_id")
+        assert await backend.aget(scope, "job_id") is None
+
+    async def test_aclear_task_removes_all_keys(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        await backend.aset(scope, "checkpoint", "step_1")
+        await backend.aclear(scope)
+        assert await backend.aget(scope, "job_id") is None
+        assert await backend.aget(scope, "checkpoint") is None
+
+    async def test_aset_and_aget_asset_roundtrip(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        result = await backend.aget(scope, "watermark")
+        assert result == "2026-04-27T00:00:00Z"
+
+    async def test_adelete_asset_removes_key(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        await backend.adelete(scope, "watermark")
+        assert await backend.aget(scope, "watermark") is None
+
+    async def test_aclear_asset_removes_all_keys(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        await backend.aset(scope, "file_count", "42")
+        await backend.aclear(scope)
+        assert await backend.aget(scope, "watermark") is None
+        assert await backend.aget(scope, "file_count") is None
+
+    async def test_aset_task_raises_for_missing_dag_run(self, backend: MetastoreStateBackend):
+        scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID)
+        with pytest.raises(ValueError, match="No DagRun found"):
+            await backend.aset(scope, "job_id", "app_async")
+
+
+class TestResolveStateBackend:
+    def test_default_resolves_to_metastore_backend(self):

Review Comment:
   Handled in dbcda100fb



##########
shared/state/tests/state/test_state.py:
##########
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from airflow_shared.state import BaseStateBackend, StateScope, TaskScope
+
+
+class TestBaseStateBackend:
+    @pytest.fixture
+    def backend(self) -> BaseStateBackend:
+        class ConcreteBackend(BaseStateBackend):
+            def __init__(self):
+                self.calls: list[str] = []
+
+            def get(self, scope: StateScope, key: str) -> str | None:
+                self.calls.append("get")
+                return "value"
+
+            def set(self, scope: StateScope, key: str, value: str) -> None:
+                self.calls.append("set")
+
+            def delete(self, scope: StateScope, key: str) -> None:
+                self.calls.append("delete")
+
+            def clear(self, scope: StateScope) -> None:
+                self.calls.append("clear")
+
+            async def aget(self, scope: StateScope, key: str) -> str | None:
+                self.calls.append("aget")
+                return "value"
+
+            async def aset(self, scope: StateScope, key: str, value: str) -> None:
+                self.calls.append("aset")
+
+            async def adelete(self, scope: StateScope, key: str) -> None:
+                self.calls.append("adelete")
+
+            async def aclear(self, scope: StateScope) -> None:
+                self.calls.append("aclear")
+
+        return ConcreteBackend()
+
+    def test_aget_is_callable(self, backend):

Review Comment:
   Handled in dbcda100fb



##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -0,0 +1,363 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import select
+
+from airflow._shared.timezones import timezone
+from airflow.models.asset import AssetModel
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+from airflow.state import AssetScope, TaskScope, resolve_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+from airflow.utils.session import create_session
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+
+
+@pytest.fixture
+def backend() -> MetastoreStateBackend:
+    return MetastoreStateBackend()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+@pytest.fixture
+def dag_run_committed() -> DagRun:
+    """DagRun committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        run = DagRun(
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            run_type=DagRunType.SCHEDULED,
+            logical_date=timezone.datetime(2026, 4, 24),
+            run_after=timezone.datetime(2026, 4, 24),
+        )
+        session.add(run)
+        session.flush()
+        session.expunge(run)
+    return run
+
+
+@pytest.fixture
+def asset(session: Session) -> AssetModel:
+    a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+    session.add(a)
+    session.flush()
+    return a
+
+
+@pytest.fixture
+def asset_committed() -> AssetModel:
+    """AssetModel committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+        session.add(a)
+        session.flush()
+        session.expunge(a)
+    return a
+
+
+class TestMetastoreStateBackendTaskScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_1234"
+
+    def test_set_twice_overrides_existing_value(

Review Comment:
   Handled in dbcda100fb



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153690847


##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: fde9ed84d07b
+Revises: 9fabad868fdb
+Create Date: 2026-04-27 09:29:20.165461
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "fde9ed84d07b"
+down_revision = "9fabad868fdb"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+    """Apply add task_state and asset_state tables."""
+    op.create_table(
+        "asset_state",
+        sa.Column("asset_id", sa.Integer(), nullable=False),
+        sa.Column("key", sa.String(length=512), nullable=False),
+        sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
+        sa.Column("updated_at", UtcDateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["asset_id"], ["asset.id"], name="asset_state_asset_fkey", ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"),
+    )
+    op.create_table(
+        "task_state",
+        sa.Column("dag_run_id", sa.Integer(), nullable=False),
+        sa.Column("task_id", sa.String(length=250), nullable=False),

Review Comment:
   `task_id`, `dag_id`, `run_id` are declared as `sa.String(length=250)` here (and in the `asset_state` table), but the ORM uses `String(ID_LEN, **COLLATION_ARGS)` which on MySQL adds `utf8mb3_bin`. The migration ships with the server-default collation, so FK joins / WHERE comparisons across `task_state` and `dag_run` either error with `Illegal mix of collations` or skip the index on MySQL. Use `from airflow.migrations.db_types import StringID` and `StringID()` for these columns -- migration `0108_3_2_0_fix_migration_file_ORM_inconsistencies.py` fixed the same drift in bulk for existing tables.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153682161


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: int
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):

Review Comment:
   Thanks. I went with documentation over splitting. Splitting into `TaskStateBackend` + `AssetStateBackend` would double the config surface ([state] task_backend + [state] asset_backend), require two separate instances everywhere the backend is used, and make the StateScope union type redundant imo.
   
   Updated `BaseStateBackend` with a class-level docstring showing the standard dispatch pattern, and added "Must handle both TaskScope and AssetScope" not in each abstract method. Custom backend authors should have everything they need from the interface doc alone.
   
   
   OR we can have abstract methods like: `get_task_state`, `get_asset_state`, the cost of which will be that abstract method count doubles. Any preference? I am ok with both



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153628438


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   [state_store] sounds better yeah. I do not want to bind it to task ONLY cos there's asset state as well.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154114471


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:

Review Comment:
   Undecided what the best option here is but I think I concur most with `-1` means "all map indices for this task" for mapped tasks. ie: if not specified, we delete all. Making the change and documenting it too



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3146185578


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)

Review Comment:
   Handled this in 1d24107e93



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141819721


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1

Review Comment:
   Also as above, should we not use the task UUID?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154253985


##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -0,0 +1,363 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import select
+
+from airflow._shared.timezones import timezone
+from airflow.models.asset import AssetModel
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+from airflow.state import AssetScope, TaskScope, resolve_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+from airflow.utils.session import create_session
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+
+
+@pytest.fixture
+def backend() -> MetastoreStateBackend:
+    return MetastoreStateBackend()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+@pytest.fixture
+def dag_run_committed() -> DagRun:
+    """DagRun committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        run = DagRun(
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            run_type=DagRunType.SCHEDULED,
+            logical_date=timezone.datetime(2026, 4, 24),
+            run_after=timezone.datetime(2026, 4, 24),
+        )
+        session.add(run)
+        session.flush()
+        session.expunge(run)
+    return run
+
+
+@pytest.fixture
+def asset(session: Session) -> AssetModel:
+    a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+    session.add(a)
+    session.flush()
+    return a
+
+
+@pytest.fixture
+def asset_committed() -> AssetModel:
+    """AssetModel committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+        session.add(a)
+        session.flush()
+        session.expunge(a)
+    return a
+
+
+class TestMetastoreStateBackendTaskScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_1234"
+
+    def test_set_twice_overrides_existing_value(

Review Comment:
   Very good catch. I was under the impression that if `pytestmark = pytest.mark.db_test` was added, it would run on all backends. The conftest skips tests without a matching backend marker when `--backend postgres/mysql` is passed, so the upsert tests were only ever running on sqlite, and I was mistaken. Added `@pytest.mark.backend("postgres", "mysql", "sqlite")` to both
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154270379


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Handled it below, lmk what you think.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153683622


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: int
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):

Review Comment:
   Handled in e20daac1c8 btw
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #65759:
URL: https://github.com/apache/airflow/pull/65759#issuecomment-4350511743

   Thanks for reviews folks, merging this one in, I handled all comments and resolved it, but if you think you need some changes to be done, let me know, I can take it as a follow up


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3149221770


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Ooh, interesting... Since State is going to be a core part of Airflow (part of the base installation), I think it might make sense to have it in core. But I can see both sides of it.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3138882327


##########
airflow-core/tests/unit/models/test_task_state.py:
##########
@@ -0,0 +1,234 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+
+from tests_common.test_utils.db import clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+class TestTaskStateModel:
+    def test_insert_and_read(self, session: Session, dag_run: DagRun):
+        row = TaskStateModel(
+            dag_run_id=dag_run.id,
+            task_id=TASK_ID,
+            map_index=-1,
+            key="remote_job_id",
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            value="application_1234",
+            updated_by_try=1,
+        )
+        session.add(row)
+        session.flush()
+
+        result = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_run_id == dag_run.id,
+                TaskStateModel.task_id == TASK_ID,
+                TaskStateModel.map_index == -1,
+                TaskStateModel.key == "remote_job_id",
+            )
+        )
+        assert result is not None
+        assert result.value == "application_1234"
+        assert result.updated_by_try == 1
+        assert result.dag_id == DAG_ID
+        assert result.run_id == RUN_ID
+
+    def test_duplicate_pk_raises(self, session: Session, dag_run: DagRun):
+        """Inserting a second row with the same (dag_run_id, task_id, map_index, key) raises IntegrityError."""
+        session.add(
+            TaskStateModel(
+                dag_run_id=dag_run.id,
+                task_id=TASK_ID,
+                map_index=-1,
+                key="remote_job_id",
+                dag_id=DAG_ID,
+                run_id=RUN_ID,
+                value="first",
+            )
+        )
+        session.flush()
+
+        session.add(
+            TaskStateModel(
+                dag_run_id=dag_run.id,
+                task_id=TASK_ID,
+                map_index=-1,
+                key="remote_job_id",
+                dag_id=DAG_ID,
+                run_id=RUN_ID,
+                value="duplicate",
+            )
+        )
+        with pytest.raises(IntegrityError):
+            session.flush()
+
+    def test_retries_share_state(self, session: Session, dag_run: DagRun):
+        """Validates that retries share state."""
+        row = TaskStateModel(
+            dag_run_id=dag_run.id,
+            task_id=TASK_ID,
+            map_index=-1,
+            key="remote_job_id",
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            value="application_try1",
+            updated_by_try=1,
+        )
+        session.add(row)
+        session.flush()
+
+        row.value = "application_try1"

Review Comment:
   `row.value = "application_try1"` is the same value as the initial write on line 132. The test only asserts `updated_by_try` changed but doesn't verify the value was actually updated by the retry.
   
   Should be something like `row.value = "application_try2"` with a corresponding `assert result.value == "application_try2"`.
   



##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: str
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):
+    """
+    Abstract backend for reading and writing task and asset state.
+
+    - Sync methods are abstract.
+    - Async methods default to calling their sync counterparts so that sync-only backends (e.g. the default metastore backend)
+    get async support for free. Backends with native async I/O (e.g. S3) should
+    override the async methods.
+    """
+
+    @abstractmethod
+    def get(self, scope: StateScope, key: str) -> str | None:
+        """Return the stored value, or None if the key does not exist."""
+
+    @abstractmethod
+    def set(self, scope: StateScope, key: str, value: str) -> None:
+        """Write or overwrite the value for the given key."""
+
+    @abstractmethod
+    def delete(self, scope: StateScope, key: str) -> None:
+        """Delete a single key. No-op if the key does not exist."""
+
+    @abstractmethod
+    def clear(self, scope: StateScope) -> None:
+        """Delete all keys under the given scope."""
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        return self.get(scope, key)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        self.set(scope, key, value)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        self.delete(scope, key)
+
+    async def aclear(self, scope: StateScope) -> None:
+        self.clear(scope)

Review Comment:
   The async impls call the sync methods directly, which blocks the event loop. This will be a problem for triggers, which run in an async event loop and would call the backend directly (not via the Execution API).
   
   These need native async implementations — at minimum the MetastoreStateBackend needs async SQLAlchemy session support for the async variants. The base class should probably make the async methods abstract too (alongside the sync ones), so backend authors are forced to provide real async implementations rather than silently inheriting blocking defaults.
   
   The PR description says these are "required for AIP-98 compatibility — async tasks must be able to call `await task_state.aset(...)` without blocking the event loop" — but as written they do block.
   



##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+

Review Comment:
   `asset_id` is an integer FK in the DB but typed as `str` here. This forces `int(scope.asset_id)` casts in every asset method in MetastoreStateBackend.
   
   If the field is always numeric, typing it as `int` removes the scattered casts. If it needs to be `str` for serialization across the Execution API (PR 2), a comment explaining why would help.
   



##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2987,7 +2987,19 @@ dag_processor:
       type: string
       example: ~
       default: "warning"
-
+state:

Review Comment:
   nit: missing blank line before `state:` — other top-level sections in this file are separated by blank lines. Same between `state:` and `profiling:` below.
   



##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)
+    updated_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+    updated_by_try: Mapped[int | None] = mapped_column(Integer, nullable=True)

Review Comment:
   `updated_by_try` is never set by `MetastoreStateBackend._set_task`, and `TaskScope` doesn't carry a try number — so this column will always be NULL for rows written through the backend.
   
   More broadly: should this be the `TaskInstance.id` UUID instead of an integer try number? The worker's JWT is scoped to `task_instance_id` (a UUID), so that's what the caller naturally has. A UUID is globally unique and directly joinable to `task_instance` / `task_instance_history`, whereas an integer `try_number` is only meaningful combined with the rest of the PK scope you already have. Something like `updated_by_ti_id` (UUID FK to `task_instance.id`) would be more useful for auditing and wouldn't require an extra lookup on the write path.
   



##########
shared/state/.pre-commit-config.yaml:
##########
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+---
+default_stages: [pre-commit, pre-push]
+minimum_prek_version: '0.3.4'
+default_language_version:
+  python: python3
+repos:
+  - repo: local
+    hooks:
+      - id: mypy-shared-state
+        name: Run mypy for shared-statse

Review Comment:
   Typo:
   
   ```suggestion
           name: Run mypy for shared-state
   ```
   



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,162 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, provide_session
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        if row is None:
+            dag_run_id = session.scalar(
+                select(DagRun.id).where(
+                    DagRun.dag_id == scope.dag_id,
+                    DagRun.run_id == scope.run_id,
+                )
+            )
+            row = TaskStateModel(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+            )
+            session.add(row)
+        row.value = value
+        row.updated_at = timezone.utcnow()

Review Comment:
   SELECT-then-INSERT without locking. Two concurrent calls for the same `(scope, key)` can both see `None`, both INSERT, and one gets an IntegrityError. Same pattern in `_set_asset`.
   
   A database-level UPSERT (`INSERT...ON CONFLICT UPDATE`) would be atomic and also more efficient (single round-trip instead of two). In practice concurrent writes to the same key should be rare, but this is easy to get right from the start.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153600156


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   `[task_state]` or `[state_store]`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #65759:
URL: https://github.com/apache/airflow/pull/65759#issuecomment-4345332492

   @jscheffl @jroachgolf84 @ashb do you wanna take another look? Planning to merge it once CI is green otherwise


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #65759:
URL: https://github.com/apache/airflow/pull/65759#issuecomment-4349997280

   The failure: https://github.com/apache/airflow/actions/runs/25119239062/job/73621449539?pr=65759 was intermittent, retriggered that job


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154027093


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   This one needs handling.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3139558794


##########
airflow-core/tests/unit/models/test_task_state.py:
##########
@@ -0,0 +1,234 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+
+from tests_common.test_utils.db import clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+class TestTaskStateModel:
+    def test_insert_and_read(self, session: Session, dag_run: DagRun):
+        row = TaskStateModel(
+            dag_run_id=dag_run.id,
+            task_id=TASK_ID,
+            map_index=-1,
+            key="remote_job_id",
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            value="application_1234",
+            updated_by_try=1,
+        )
+        session.add(row)
+        session.flush()
+
+        result = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_run_id == dag_run.id,
+                TaskStateModel.task_id == TASK_ID,
+                TaskStateModel.map_index == -1,
+                TaskStateModel.key == "remote_job_id",
+            )
+        )
+        assert result is not None
+        assert result.value == "application_1234"
+        assert result.updated_by_try == 1
+        assert result.dag_id == DAG_ID
+        assert result.run_id == RUN_ID
+
+    def test_duplicate_pk_raises(self, session: Session, dag_run: DagRun):
+        """Inserting a second row with the same (dag_run_id, task_id, map_index, key) raises IntegrityError."""
+        session.add(
+            TaskStateModel(
+                dag_run_id=dag_run.id,
+                task_id=TASK_ID,
+                map_index=-1,
+                key="remote_job_id",
+                dag_id=DAG_ID,
+                run_id=RUN_ID,
+                value="first",
+            )
+        )
+        session.flush()
+
+        session.add(
+            TaskStateModel(
+                dag_run_id=dag_run.id,
+                task_id=TASK_ID,
+                map_index=-1,
+                key="remote_job_id",
+                dag_id=DAG_ID,
+                run_id=RUN_ID,
+                value="duplicate",
+            )
+        )
+        with pytest.raises(IntegrityError):
+            session.flush()
+
+    def test_retries_share_state(self, session: Session, dag_run: DagRun):
+        """Validates that retries share state."""
+        row = TaskStateModel(
+            dag_run_id=dag_run.id,
+            task_id=TASK_ID,
+            map_index=-1,
+            key="remote_job_id",
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            value="application_try1",
+            updated_by_try=1,
+        )
+        session.add(row)
+        session.flush()
+
+        row.value = "application_try1"

Review Comment:
   Good catch, fixed it — second write now uses "application_try2"



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3139611008


##########
shared/state/.pre-commit-config.yaml:
##########
@@ -0,0 +1,31 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+---
+default_stages: [pre-commit, pre-push]
+minimum_prek_version: '0.3.4'
+default_language_version:
+  python: python3
+repos:
+  - repo: local
+    hooks:
+      - id: mypy-shared-state
+        name: Run mypy for shared-statse

Review Comment:
   Thanks, handled.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145869136


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1

Review Comment:
   Replied here: https://github.com/apache/airflow/pull/65759#discussion_r3145867282



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jroachgolf84 (via GitHub)" <gi...@apache.org>.
jroachgolf84 commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141859388


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:

Review Comment:
   I don't want to get too pedantic, but `_delete_task` (and other things like `_get_asset` or `_set_asset`) are a bit misleading. It sounds like the Task itself is being deleted; not just the state for that Task.
   
   Might it make sense for it to be something like `_delete_task_state? If not, no worries, just throwing thoughts out there.
   
   I guess as I say it, I'm realizing that this interface wouldn't be exposed to users...



##########
airflow-core/src/airflow/models/asset_state.py:
##########


Review Comment:
   This is exactly what I had in mind for this; awesome!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3161053289


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,351 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.typing_compat import assert_never
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+            case _:
+                assert_never(scope)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+            case _:
+                assert_never(scope)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+                case _:
+                    assert_never(scope)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+                case _:
+                    assert_never(scope)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:
+        conditions = [
+            TaskStateModel.dag_id == scope.dag_id,
+            TaskStateModel.run_id == scope.run_id,
+            TaskStateModel.task_id == scope.task_id,
+        ]
+        if scope.map_index != -1:

Review Comment:
   `_clear_task_state` drops the `map_index` filter when `map_index == -1`, but `_delete_task_state` above keeps the exact-match filter. Same `TaskScope(map_index=-1)` value, opposite semantics: `delete` only touches the non-mapped row, `clear` wipes every map index. That asymmetry is a foot-gun even when documented. Either drop the `map_index` filter in delete too when `map_index == -1`, or replace the sentinel with an explicit `clear_all_map_indices: bool = False` flag. Same applies to `_aclear_task_state` below.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141637614


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   These could be ripe for using match statements 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3139566589


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,162 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, provide_session
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        if row is None:
+            dag_run_id = session.scalar(
+                select(DagRun.id).where(
+                    DagRun.dag_id == scope.dag_id,
+                    DagRun.run_id == scope.run_id,
+                )
+            )
+            row = TaskStateModel(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+            )
+            session.add(row)
+        row.value = value
+        row.updated_at = timezone.utcnow()

Review Comment:
   Fixed with dialect-specific UPSERT — same pattern as `Variable.set`: https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/models/variable.py#L260-L319



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141820990


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Should we not add this to the core section like we do with `xcom_backend`. Sounds a bit strange to open a section for a single entry (will there be more?)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141521870


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: str
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):
+    """
+    Abstract backend for reading and writing task and asset state.
+
+    - Sync methods are abstract.
+    - Async methods default to calling their sync counterparts so that sync-only backends (e.g. the default metastore backend)
+    get async support for free. Backends with native async I/O (e.g. S3) should
+    override the async methods.
+    """
+
+    @abstractmethod
+    def get(self, scope: StateScope, key: str) -> str | None:
+        """Return the stored value, or None if the key does not exist."""
+
+    @abstractmethod
+    def set(self, scope: StateScope, key: str, value: str) -> None:
+        """Write or overwrite the value for the given key."""
+
+    @abstractmethod
+    def delete(self, scope: StateScope, key: str) -> None:
+        """Delete a single key. No-op if the key does not exist."""
+
+    @abstractmethod
+    def clear(self, scope: StateScope) -> None:
+        """Delete all keys under the given scope."""
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        return self.get(scope, key)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        self.set(scope, key, value)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        self.delete(scope, key)
+
+    async def aclear(self, scope: StateScope) -> None:
+        self.clear(scope)

Review Comment:
   ACK, made the change for that and updated PR desc too.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145823012


##########
airflow-core/src/airflow/models/asset_state.py:
##########


Review Comment:
   Cool!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145922853


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(

Review Comment:
   Fixed. Both sync and async paths now raise `ValueError` before attempting the INSERT.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3152960588


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   That's right, I prefer dedicated `state` section too. Hmm I wonder if I can find a better name for config actually. Naming is hard!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154268856


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   Same as https://github.com/apache/airflow/pull/65759#discussion_r3153690839, handled



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141521870


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: str
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):
+    """
+    Abstract backend for reading and writing task and asset state.
+
+    - Sync methods are abstract.
+    - Async methods default to calling their sync counterparts so that sync-only backends (e.g. the default metastore backend)
+    get async support for free. Backends with native async I/O (e.g. S3) should
+    override the async methods.
+    """
+
+    @abstractmethod
+    def get(self, scope: StateScope, key: str) -> str | None:
+        """Return the stored value, or None if the key does not exist."""
+
+    @abstractmethod
+    def set(self, scope: StateScope, key: str, value: str) -> None:
+        """Write or overwrite the value for the given key."""
+
+    @abstractmethod
+    def delete(self, scope: StateScope, key: str) -> None:
+        """Delete a single key. No-op if the key does not exist."""
+
+    @abstractmethod
+    def clear(self, scope: StateScope) -> None:
+        """Delete all keys under the given scope."""
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        return self.get(scope, key)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        self.set(scope, key, value)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        self.delete(scope, key)
+
+    async def aclear(self, scope: StateScope) -> None:
+        self.clear(scope)

Review Comment:
   ACK, made the change for that and updated PR desc too.
   
   
   There were some refs in how we do this in the API, handled accordingly



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141669278


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)

Review Comment:
   Just to avoid confusion, you mean https://docs.python.org/3/tutorial/controlflow.html#match-statements yes?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145905489


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:

Review Comment:
   Yeah its private but since you brought it up, let's fix it



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145913335


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:

Review Comment:
   Handled in 3f03dc7492



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153690860


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:

Review Comment:
   `_clear_task_state` filters `WHERE map_index = scope.map_index`, and `TaskScope.map_index` defaults to `-1`. So `backend.clear(TaskScope(dag_id, run_id, task_id))` only deletes non-mapped rows -- for a mapped task with state at map_index 0/1/2, this is a silent no-op. The new `test_clear_scoped_to_map_index` actually demonstrates the surprise (clear at index 0 leaves index 1 alone), but there's no path to clear all map indices for a task. Either drop the default on `map_index`, or special-case `-1` inside `clear`/`aclear` to mean "all map indices for this task", or document the limit on both `TaskScope` and `clear()`.



##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -0,0 +1,363 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import select
+
+from airflow._shared.timezones import timezone
+from airflow.models.asset import AssetModel
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+from airflow.state import AssetScope, TaskScope, resolve_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+from airflow.utils.session import create_session
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+
+
+@pytest.fixture
+def backend() -> MetastoreStateBackend:
+    return MetastoreStateBackend()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+@pytest.fixture
+def dag_run_committed() -> DagRun:
+    """DagRun committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        run = DagRun(
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            run_type=DagRunType.SCHEDULED,
+            logical_date=timezone.datetime(2026, 4, 24),
+            run_after=timezone.datetime(2026, 4, 24),
+        )
+        session.add(run)
+        session.flush()
+        session.expunge(run)
+    return run
+
+
+@pytest.fixture
+def asset(session: Session) -> AssetModel:
+    a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+    session.add(a)
+    session.flush()
+    return a
+
+
+@pytest.fixture
+def asset_committed() -> AssetModel:
+    """AssetModel committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+        session.add(a)
+        session.flush()
+        session.expunge(a)
+    return a
+
+
+class TestMetastoreStateBackendTaskScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_1234"
+
+    def test_set_twice_overrides_existing_value(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """Calling set twice on the same key updates the value in place."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_try1", session=session)
+        session.flush()
+        backend.set(scope, "job_id", "app_try2", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_try2"
+
+    def test_set_stores_dag_run_id_fk(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """set resolves dag_run_id from (dag_id, run_id) and persists it as the FK."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == DAG_ID,
+                TaskStateModel.task_id == TASK_ID,
+                TaskStateModel.key == "job_id",
+            )
+        )
+        assert row is not None
+        assert row.dag_run_id == dag_run.id
+
+    def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        backend.set(scope, "checkpoint", "step_3", session=session)
+        session.flush()
+
+        backend.delete(scope, "job_id", session=session)
+        session.flush()
+
+        assert backend.get(scope, "job_id", session=session) is None
+        assert backend.get(scope, "checkpoint", session=session) == "step_3"
+
+    def test_delete_is_noop_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.delete(scope, "nonexistent", session=session)
+
+    def test_clear_removes_all_keys(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        """clear removes every key under the given task scope."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        for key in ("job_id", "checkpoint", "watermark"):
+            backend.set(scope, key, f"val_{key}", session=session)
+        session.flush()
+
+        backend.clear(scope, session=session)
+        session.flush()
+
+        for key in ("job_id", "checkpoint", "watermark"):
+            assert backend.get(scope, key, session=session) is None
+
+    def test_map_index_isolation(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        """Mapped task instances with different map_index values have isolated namespaces."""
+        scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0)
+        scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1)
+
+        backend.set(scope0, "job_id", "app_0", session=session)
+        backend.set(scope1, "job_id", "app_1", session=session)
+        session.flush()
+
+        assert backend.get(scope0, "job_id", session=session) == "app_0"
+        assert backend.get(scope1, "job_id", session=session) == "app_1"
+
+    def test_set_raises_for_missing_dag_run(self, session: Session, backend: MetastoreStateBackend):
+        scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID)
+        with pytest.raises(ValueError, match="No DagRun found"):
+            backend.set(scope, "job_id", "app_1234", session=session)
+
+    def test_clear_scoped_to_map_index(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """clear for map_index=0 does not remove state belonging to map_index=1."""
+        scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0)
+        scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1)
+
+        backend.set(scope0, "job_id", "app_0", session=session)
+        backend.set(scope1, "job_id", "app_1", session=session)
+        session.flush()
+
+        backend.clear(scope0, session=session)
+        session.flush()
+
+        assert backend.get(scope0, "job_id", session=session) is None
+        assert backend.get(scope1, "job_id", session=session) == "app_1"
+
+
+class TestMetastoreStateBackendAssetScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        session.flush()
+        assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z"
+
+    def test_set_twice_overwrites_existing_value(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-01T00:00:00Z", session=session)
+        session.flush()
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        session.flush()
+        assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z"
+
+    def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        backend.set(scope, "file_count", "42", session=session)
+        session.flush()
+
+        backend.delete(scope, "watermark", session=session)
+        session.flush()
+
+        assert backend.get(scope, "watermark", session=session) is None
+        assert backend.get(scope, "file_count", session=session) == "42"
+
+    def test_delete_is_noop_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        backend.delete(scope, "nonexistent", session=session)
+
+    def test_clear_removes_all_keys(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        for key in ("watermark", "file_count", "last_error"):
+            backend.set(scope, key, f"val_{key}", session=session)
+        session.flush()
+
+        backend.clear(scope, session=session)
+        session.flush()
+
+        for key in ("watermark", "file_count", "last_error"):
+            assert backend.get(scope, key, session=session) is None
+
+    def test_different_assets_are_isolated(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        asset2 = AssetModel(uri="s3://bucket/other", name="other_asset", group="test")
+        session.add(asset2)
+        session.flush()
+
+        scope1 = AssetScope(asset_id=asset.id)
+        scope2 = AssetScope(asset_id=asset2.id)
+
+        backend.set(scope1, "watermark", "asset1_value", session=session)
+        session.flush()
+
+        assert backend.get(scope2, "watermark", session=session) is None
+
+
+@pytest.mark.asyncio(loop_scope="class")
+class TestMetastoreStateBackendAsync:
+    async def test_aset_and_aget_task_roundtrip(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        result = await backend.aget(scope, "job_id")
+        assert result == "app_async"
+
+    async def test_adelete_task_removes_key(self, backend: MetastoreStateBackend, dag_run_committed: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        await backend.adelete(scope, "job_id")
+        assert await backend.aget(scope, "job_id") is None
+
+    async def test_aclear_task_removes_all_keys(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        await backend.aset(scope, "checkpoint", "step_1")
+        await backend.aclear(scope)
+        assert await backend.aget(scope, "job_id") is None
+        assert await backend.aget(scope, "checkpoint") is None
+
+    async def test_aset_and_aget_asset_roundtrip(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        result = await backend.aget(scope, "watermark")
+        assert result == "2026-04-27T00:00:00Z"
+
+    async def test_adelete_asset_removes_key(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        await backend.adelete(scope, "watermark")
+        assert await backend.aget(scope, "watermark") is None
+
+    async def test_aclear_asset_removes_all_keys(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        await backend.aset(scope, "file_count", "42")
+        await backend.aclear(scope)
+        assert await backend.aget(scope, "watermark") is None
+        assert await backend.aget(scope, "file_count") is None
+
+    async def test_aset_task_raises_for_missing_dag_run(self, backend: MetastoreStateBackend):
+        scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID)
+        with pytest.raises(ValueError, match="No DagRun found"):
+            await backend.aset(scope, "job_id", "app_async")
+
+
+class TestResolveStateBackend:
+    def test_default_resolves_to_metastore_backend(self):

Review Comment:
   `test_default_resolves_to_metastore_backend` calls `resolve_state_backend()` with no `conf_vars` override, so it depends on global config state -- if a previous test sets `state.backend = ""` and forgets cleanup, this fails non-deterministically. It also passes by accident: it's exercising the `config.yml` default rather than the `if not clazz` branch in `resolve_state_backend`. Add a separate test wrapped in `conf_vars({("state", "backend"): ""})` that asserts the explicit-empty path (which today returns `MetastoreStateBackend`, but the previous review suggested raising instead -- pick one and lock it in), and rename this one to `test_resolve_returns_default_from_config_yml`.



##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -0,0 +1,363 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import select
+
+from airflow._shared.timezones import timezone
+from airflow.models.asset import AssetModel
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+from airflow.state import AssetScope, TaskScope, resolve_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+from airflow.utils.session import create_session
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+
+
+@pytest.fixture
+def backend() -> MetastoreStateBackend:
+    return MetastoreStateBackend()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+@pytest.fixture
+def dag_run_committed() -> DagRun:
+    """DagRun committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        run = DagRun(
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            run_type=DagRunType.SCHEDULED,
+            logical_date=timezone.datetime(2026, 4, 24),
+            run_after=timezone.datetime(2026, 4, 24),
+        )
+        session.add(run)
+        session.flush()
+        session.expunge(run)
+    return run
+
+
+@pytest.fixture
+def asset(session: Session) -> AssetModel:
+    a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+    session.add(a)
+    session.flush()
+    return a
+
+
+@pytest.fixture
+def asset_committed() -> AssetModel:
+    """AssetModel committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+        session.add(a)
+        session.flush()
+        session.expunge(a)
+    return a
+
+
+class TestMetastoreStateBackendTaskScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_1234"
+
+    def test_set_twice_overrides_existing_value(

Review Comment:
   These set/get/upsert tests only run against the default backend in CI (sqlite). A typo in the `pg_insert.on_conflict_do_update(...)` or `mysql_insert.on_duplicate_key_update(...)` branches in `_build_upsert_stmt` wouldn't be caught -- and given the dialect routing is the load-bearing piece for this PR, that's a real gap. Add a `pytest.mark.backend("postgres", "mysql", "sqlite")` parametrisation on at least `test_set_twice_overrides_existing_value` so each upsert branch gets exercised under its real dialect.



##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2997,6 +2997,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Re-flagging the section-name concern from the prior round. `state` is heavily overloaded in Airflow (`TI.state`, `DR.state`, `TaskInstanceState`, `state_color_*`, etc.), so users will read `[state]` and reasonably expect it to be about TI/DR state behaviour, not a key/value blob backend. `[state_store]` parallels `[secrets]` and reads cleanly; `[task_state]` also works. Worth picking now since AIP-103 will likely add per-backend keys (retention, GC) in the next PRs and renaming a config section after release is painful.



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:

Review Comment:
   The `match scope:` blocks here and in the seven sibling methods (lines 80, 88, 96, 104, 112, 120, 128) have no `case _:` arm. `StateScope = TaskScope | AssetScope` today, but adding a future scope subtype makes `get`/`aget` silently return `None` and `set`/`delete`/`clear` silently no-op (writes vanish, no log, no error). Add `case _: assert_never(scope)` (from `typing`) to all eight match blocks so mypy enforces exhaustiveness on every new subtype and runtime fails loudly. This also doubles as the dispatch documentation flagged in the previous review.



##########
shared/state/tests/state/test_state.py:
##########
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from airflow_shared.state import BaseStateBackend, StateScope, TaskScope
+
+
+class TestBaseStateBackend:
+    @pytest.fixture
+    def backend(self) -> BaseStateBackend:
+        class ConcreteBackend(BaseStateBackend):
+            def __init__(self):
+                self.calls: list[str] = []
+
+            def get(self, scope: StateScope, key: str) -> str | None:
+                self.calls.append("get")
+                return "value"
+
+            def set(self, scope: StateScope, key: str, value: str) -> None:
+                self.calls.append("set")
+
+            def delete(self, scope: StateScope, key: str) -> None:
+                self.calls.append("delete")
+
+            def clear(self, scope: StateScope) -> None:
+                self.calls.append("clear")
+
+            async def aget(self, scope: StateScope, key: str) -> str | None:
+                self.calls.append("aget")
+                return "value"
+
+            async def aset(self, scope: StateScope, key: str, value: str) -> None:
+                self.calls.append("aset")
+
+            async def adelete(self, scope: StateScope, key: str) -> None:
+                self.calls.append("adelete")
+
+            async def aclear(self, scope: StateScope) -> None:
+                self.calls.append("aclear")
+
+        return ConcreteBackend()
+
+    def test_aget_is_callable(self, backend):

Review Comment:
   `test_aget_is_callable` and the three siblings (`test_aset_is_callable`, `test_adelete_is_callable`, `test_aclear_is_callable`) test that `asyncio.run` runs an `async def` and that `list.append` works -- none of them exercise anything specific to `BaseStateBackend`. They'll keep passing if you delete every async method from the ABC, or replace them with sync ones. Replace with: (a) a test that instantiating an incomplete `BaseStateBackend` subclass raises `TypeError` ("abstract"), and (b) one that asserts `BaseStateBackend.__abstractmethods__` contains every `a*` method paired with its sync sibling -- that's the actual contract this file should be defending.



##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)

Review Comment:
   While here, switch the ORM to `StringID()` from `airflow.models.base` instead of the raw `String(ID_LEN, **COLLATION_ARGS)` form for `task_id`, `dag_id`, `run_id` (and the equivalents in `asset_state.py`). Every recent model (backfill, deadline, taskinstance) uses the helper -- keeps the collation in one place and matches what the migration should be doing per the previous comment.



##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: fde9ed84d07b
+Revises: 9fabad868fdb
+Create Date: 2026-04-27 09:29:20.165461
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "fde9ed84d07b"
+down_revision = "9fabad868fdb"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+    """Apply add task_state and asset_state tables."""
+    op.create_table(
+        "asset_state",
+        sa.Column("asset_id", sa.Integer(), nullable=False),
+        sa.Column("key", sa.String(length=512), nullable=False),
+        sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
+        sa.Column("updated_at", UtcDateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["asset_id"], ["asset.id"], name="asset_state_asset_fkey", ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"),
+    )
+    op.create_table(
+        "task_state",
+        sa.Column("dag_run_id", sa.Integer(), nullable=False),
+        sa.Column("task_id", sa.String(length=250), nullable=False),

Review Comment:
   `task_id`, `dag_id`, `run_id` are declared as `sa.String(length=250)` here (and in the `asset_state` table), but the ORM uses `String(ID_LEN, **COLLATION_ARGS)` which on MySQL adds `utf8mb3_bin`. The migration ships with the server-default collation, so FK joins / WHERE comparisons across `task_state` and `dag_run` either error with `Illegal mix of collations` or skip the index on MySQL. Use `from airflow.migrations.db_types import StringID` and `StringID()` for these columns -- migration #0108 fixed the same drift in bulk for existing tables.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154214256


##########
shared/state/tests/state/test_state.py:
##########
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from airflow_shared.state import BaseStateBackend, StateScope, TaskScope
+
+
+class TestBaseStateBackend:
+    @pytest.fixture
+    def backend(self) -> BaseStateBackend:
+        class ConcreteBackend(BaseStateBackend):
+            def __init__(self):
+                self.calls: list[str] = []
+
+            def get(self, scope: StateScope, key: str) -> str | None:
+                self.calls.append("get")
+                return "value"
+
+            def set(self, scope: StateScope, key: str, value: str) -> None:
+                self.calls.append("set")
+
+            def delete(self, scope: StateScope, key: str) -> None:
+                self.calls.append("delete")
+
+            def clear(self, scope: StateScope) -> None:
+                self.calls.append("clear")
+
+            async def aget(self, scope: StateScope, key: str) -> str | None:
+                self.calls.append("aget")
+                return "value"
+
+            async def aset(self, scope: StateScope, key: str, value: str) -> None:
+                self.calls.append("aset")
+
+            async def adelete(self, scope: StateScope, key: str) -> None:
+                self.calls.append("adelete")
+
+            async def aclear(self, scope: StateScope) -> None:
+                self.calls.append("aclear")
+
+        return ConcreteBackend()
+
+    def test_aget_is_callable(self, backend):

Review Comment:
   Thank you, good catch. I replaced all four of those async callable tests with two tests that actually defend the contract: `test_incomplete_subclass_raises_type_error` (instantiating a partial implementation raises TypeError) and `test_abstract_methods_cover_full_interface` (asserts __abstractmethods__ is exactly the 8 expected methods).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3146012065


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":

Review Comment:
   Handled in 171f77061b



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3153690832


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2997,6 +2997,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   Re-flagging the section-name concern from the prior round. `state` is heavily overloaded in Airflow (`TI.state`, `DR.state`, `TaskInstanceState`, `state_color_*`, etc.), so users will read `[state]` and reasonably expect it to be about TI/DR state behaviour, not a key/value blob backend. `[state_store]` parallels `[secrets]` and reads cleanly; `[task_state]` also works. Worth picking now since AIP-103 will likely add per-backend keys (retention, GC) in the next PRs and renaming a config section after release is painful.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3142869619


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)

Review Comment:
   `Text` maps to a 64KB-capped `TEXT` on MySQL but is unlimited on PostgreSQL/SQLite. If a user stores a JSON blob over 64KB on MySQL they hit a silent truncation (or a strict-mode error) where the same value succeeds on Postgres. Either document the cap, switch to `LONGTEXT` via `Text().with_variant(...)`, or hard-cap with `String(...)` to make the limit explicit and portable. The same applies to `asset_state.value` and the migration's `Column("value", Text, ...)`.



##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: int
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):

Review Comment:
   The interface accepts `StateScope` for every method but each backend has to dispatch on `isinstance(scope, TaskScope)` internally (see `MetastoreStateBackend`). Worth either documenting in the docstring which scope types each method must accept, or splitting into `TaskStateBackend` and `AssetStateBackend` so custom backends don't need to repeat the dispatch boilerplate.



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(

Review Comment:
   If the DagRun does not exist, this returns `None` and the INSERT then fails on the `dag_run_id NOT NULL` constraint with a confusing IntegrityError. Worth raising a clear error here (e.g. `ValueError(f"No DagRun found for dag_id={...} run_id={...}")`) so callers see what actually went wrong.



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":

Review Comment:
   The dialect-routing block is duplicated four times (sync task, sync asset, async task, async asset) with the same pg/mysql/sqlite branches. Could this be pulled into a small helper that takes the model class, the conflict columns and the values dict, and returns the dialect-specific upsert statement? It would make the four call sites two lines each.



##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -0,0 +1,42 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow._shared.state import (
+    AssetScope as AssetScope,
+    BaseStateBackend as BaseStateBackend,
+    StateScope as StateScope,
+    TaskScope as TaskScope,
+)
+
+_DEFAULT_BACKEND = "airflow.state.metastore.MetastoreStateBackend"

Review Comment:
   This duplicates the default in `config_templates/config.yml`. If one changes the other goes stale. You can drop this constant and let `conf.getimport("state", "backend")` (no fallback) read the schema default directly.



##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -0,0 +1,42 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow._shared.state import (
+    AssetScope as AssetScope,
+    BaseStateBackend as BaseStateBackend,
+    StateScope as StateScope,
+    TaskScope as TaskScope,
+)
+
+_DEFAULT_BACKEND = "airflow.state.metastore.MetastoreStateBackend"
+
+
+def resolve_state_backend() -> type[BaseStateBackend]:
+    from airflow.configuration import conf
+
+    clazz = conf.getimport("state", "backend", fallback=_DEFAULT_BACKEND)
+    if not clazz:

Review Comment:
   `getimport` only returns falsy if the user explicitly sets `state.backend = ` to an empty string, since the fallback already points at the metastore. Silently swapping in `MetastoreStateBackend` for an explicitly-cleared value seems like the wrong behaviour, either drop the branch or raise on empty. Either way the inline `from airflow.state.metastore import MetastoreStateBackend` is odd here, since the same class is already what `_DEFAULT_BACKEND` resolves to.



##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state.py:
##########
@@ -0,0 +1,88 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: 8a7b3c2d1e4f
+Revises: 9fabad868fdb
+Create Date: 2026-04-24 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+from sqlalchemy import Column, ForeignKeyConstraint, Integer, String, Text
+
+from airflow._shared.timezones import timezone
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "8a7b3c2d1e4f"

Review Comment:
   This revision id looks hand-typed (`8a7b3c2d1e4f` is sequential nibbles). Other migrations in this folder use random hex generated by `alembic revision`. Worth regenerating with the standard tool so it can't accidentally collide with a future auto-generated id.



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(TaskStateModel).values(
+                dag_run_id=dag_run_id,
+                dag_id=scope.dag_id,
+                run_id=scope.run_id,
+                task_id=scope.task_id,
+                map_index=scope.map_index,
+                key=key,
+                value=value,
+                updated_at=now,
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["dag_run_id", "task_id", "map_index", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_task(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task(self, scope: TaskScope, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+            )
+        )
+
+    def _get_asset(self, scope: AssetScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+                AssetStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_asset(self, scope: AssetScope, key: str, value: str, *, session: Session) -> None:
+        now = timezone.utcnow()
+        dialect = get_dialect_name(session)
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        if dialect == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            stmt = pg_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["asset_id", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        elif dialect == "mysql":
+            from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+            stmt = mysql_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_duplicate_key_update(value=value, updated_at=now)
+        else:
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            stmt = sqlite_insert(AssetStateModel).values(
+                asset_id=scope.asset_id, key=key, value=value, updated_at=now
+            )
+            stmt = stmt.on_conflict_do_update(
+                index_elements=["asset_id", "key"],
+                set_=dict(value=value, updated_at=now),
+            )
+        session.execute(stmt)
+
+    def _delete_asset(self, scope: AssetScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+                AssetStateModel.key == key,
+            )
+        )
+
+    def _clear_asset(self, scope: AssetScope, *, session: Session) -> None:
+        session.execute(
+            delete(AssetStateModel).where(
+                AssetStateModel.asset_id == scope.asset_id,
+            )
+        )
+
+    async def _aget_task(self, scope: TaskScope, key: str, *, session: AsyncSession) -> str | None:
+        row = await session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    async def _aset_task(self, scope: TaskScope, key: str, value: str, *, session: AsyncSession) -> None:
+        dag_run_id = await session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        now = timezone.utcnow()
+        stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+        # get_dialect_name expects a sync Session; sync_session is the underlying Session the async wrapper delegates to
+        dialect = get_dialect_name(session.sync_session)

Review Comment:
   The async `_aset_task` and `_aset_asset` paths aren't covered by tests in this PR. `test_metastore.py` only exercises the sync methods; the shared `test_state.py` uses a fake backend. Could you add at least one test per async method against the real metastore backend? `session.sync_session.get_bind()` for an async-engine-backed session is the kind of thing worth verifying empirically rather than by inspection.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3149801681


##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -2988,6 +2988,20 @@ dag_processor:
       example: ~
       default: "warning"
 
+state:

Review Comment:
   then even if potentially mis-placed in core - what would be a better key? Then it should be consistent for state as like with xcom :-D



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3146115184


##########
airflow-core/src/airflow/state/__init__.py:
##########
@@ -0,0 +1,42 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow._shared.state import (
+    AssetScope as AssetScope,
+    BaseStateBackend as BaseStateBackend,
+    StateScope as StateScope,
+    TaskScope as TaskScope,
+)
+
+_DEFAULT_BACKEND = "airflow.state.metastore.MetastoreStateBackend"
+
+
+def resolve_state_backend() -> type[BaseStateBackend]:
+    from airflow.configuration import conf
+
+    clazz = conf.getimport("state", "backend", fallback=_DEFAULT_BACKEND)
+    if not clazz:

Review Comment:
   Good catch, handled it here: 1d428fee7b



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141533314


##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+    """Identifies the state namespace for a single task instance (or retry thereof)."""
+
+    dag_id: str
+    run_id: str
+    task_id: str
+    map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+    """Identifies the state namespace for an asset."""
+
+    asset_id: str
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):
+    """
+    Abstract backend for reading and writing task and asset state.
+
+    - Sync methods are abstract.
+    - Async methods default to calling their sync counterparts so that sync-only backends (e.g. the default metastore backend)
+    get async support for free. Backends with native async I/O (e.g. S3) should
+    override the async methods.
+    """
+
+    @abstractmethod
+    def get(self, scope: StateScope, key: str) -> str | None:
+        """Return the stored value, or None if the key does not exist."""
+
+    @abstractmethod
+    def set(self, scope: StateScope, key: str, value: str) -> None:
+        """Write or overwrite the value for the given key."""
+
+    @abstractmethod
+    def delete(self, scope: StateScope, key: str) -> None:
+        """Delete a single key. No-op if the key does not exist."""
+
+    @abstractmethod
+    def clear(self, scope: StateScope) -> None:
+        """Delete all keys under the given scope."""
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        return self.get(scope, key)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        self.set(scope, key, value)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        self.delete(scope, key)
+
+    async def aclear(self, scope: StateScope) -> None:
+        self.clear(scope)

Review Comment:
   Handled in feff6d553b



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3141533610


##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)
+    map_index: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False, server_default="-1")
+    key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS), nullable=False, primary_key=True)
+
+    dag_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+    run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
+
+    value: Mapped[str] = mapped_column(Text, nullable=False)
+    updated_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
+    updated_by_try: Mapped[int | None] = mapped_column(Integer, nullable=True)

Review Comment:
   Removed in feff6d553b



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154196803


##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -0,0 +1,363 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from sqlalchemy import select
+
+from airflow._shared.timezones import timezone
+from airflow.models.asset import AssetModel
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+from airflow.state import AssetScope, TaskScope, resolve_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+from airflow.utils.session import create_session
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+DAG_ID = "test_dag"
+TASK_ID = "test_task"
+RUN_ID = "scheduled__2026-04-24"
+
+
+@pytest.fixture(autouse=True)
+def clean_tables():
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+    yield
+    clear_db_dags()
+    clear_db_runs()
+    clear_db_assets()
+
+
+@pytest.fixture
+def backend() -> MetastoreStateBackend:
+    return MetastoreStateBackend()
+
+
+@pytest.fixture
+def dag_run(session: Session) -> DagRun:
+    run = DagRun(
+        dag_id=DAG_ID,
+        run_id=RUN_ID,
+        run_type=DagRunType.SCHEDULED,
+        logical_date=timezone.datetime(2026, 4, 24),
+        run_after=timezone.datetime(2026, 4, 24),
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+@pytest.fixture
+def dag_run_committed() -> DagRun:
+    """DagRun committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        run = DagRun(
+            dag_id=DAG_ID,
+            run_id=RUN_ID,
+            run_type=DagRunType.SCHEDULED,
+            logical_date=timezone.datetime(2026, 4, 24),
+            run_after=timezone.datetime(2026, 4, 24),
+        )
+        session.add(run)
+        session.flush()
+        session.expunge(run)
+    return run
+
+
+@pytest.fixture
+def asset(session: Session) -> AssetModel:
+    a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+    session.add(a)
+    session.flush()
+    return a
+
+
+@pytest.fixture
+def asset_committed() -> AssetModel:
+    """AssetModel committed to DB so async sessions (which open their own connections) can see it."""
+    with create_session() as session:
+        a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test")
+        session.add(a)
+        session.flush()
+        session.expunge(a)
+    return a
+
+
+class TestMetastoreStateBackendTaskScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_1234"
+
+    def test_set_twice_overrides_existing_value(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """Calling set twice on the same key updates the value in place."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_try1", session=session)
+        session.flush()
+        backend.set(scope, "job_id", "app_try2", session=session)
+        session.flush()
+        assert backend.get(scope, "job_id", session=session) == "app_try2"
+
+    def test_set_stores_dag_run_id_fk(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """set resolves dag_run_id from (dag_id, run_id) and persists it as the FK."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        session.flush()
+
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == DAG_ID,
+                TaskStateModel.task_id == TASK_ID,
+                TaskStateModel.key == "job_id",
+            )
+        )
+        assert row is not None
+        assert row.dag_run_id == dag_run.id
+
+    def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.set(scope, "job_id", "app_1234", session=session)
+        backend.set(scope, "checkpoint", "step_3", session=session)
+        session.flush()
+
+        backend.delete(scope, "job_id", session=session)
+        session.flush()
+
+        assert backend.get(scope, "job_id", session=session) is None
+        assert backend.get(scope, "checkpoint", session=session) == "step_3"
+
+    def test_delete_is_noop_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        backend.delete(scope, "nonexistent", session=session)
+
+    def test_clear_removes_all_keys(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        """clear removes every key under the given task scope."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        for key in ("job_id", "checkpoint", "watermark"):
+            backend.set(scope, key, f"val_{key}", session=session)
+        session.flush()
+
+        backend.clear(scope, session=session)
+        session.flush()
+
+        for key in ("job_id", "checkpoint", "watermark"):
+            assert backend.get(scope, key, session=session) is None
+
+    def test_map_index_isolation(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun):
+        """Mapped task instances with different map_index values have isolated namespaces."""
+        scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0)
+        scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1)
+
+        backend.set(scope0, "job_id", "app_0", session=session)
+        backend.set(scope1, "job_id", "app_1", session=session)
+        session.flush()
+
+        assert backend.get(scope0, "job_id", session=session) == "app_0"
+        assert backend.get(scope1, "job_id", session=session) == "app_1"
+
+    def test_set_raises_for_missing_dag_run(self, session: Session, backend: MetastoreStateBackend):
+        scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID)
+        with pytest.raises(ValueError, match="No DagRun found"):
+            backend.set(scope, "job_id", "app_1234", session=session)
+
+    def test_clear_scoped_to_map_index(
+        self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+    ):
+        """clear for map_index=0 does not remove state belonging to map_index=1."""
+        scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0)
+        scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1)
+
+        backend.set(scope0, "job_id", "app_0", session=session)
+        backend.set(scope1, "job_id", "app_1", session=session)
+        session.flush()
+
+        backend.clear(scope0, session=session)
+        session.flush()
+
+        assert backend.get(scope0, "job_id", session=session) is None
+        assert backend.get(scope1, "job_id", session=session) == "app_1"
+
+
+class TestMetastoreStateBackendAssetScope:
+    def test_get_returns_none_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        assert backend.get(scope, "missing", session=session) is None
+
+    def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        session.flush()
+        assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z"
+
+    def test_set_twice_overwrites_existing_value(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-01T00:00:00Z", session=session)
+        session.flush()
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        session.flush()
+        assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z"
+
+    def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel):
+        scope = AssetScope(asset_id=asset.id)
+        backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session)
+        backend.set(scope, "file_count", "42", session=session)
+        session.flush()
+
+        backend.delete(scope, "watermark", session=session)
+        session.flush()
+
+        assert backend.get(scope, "watermark", session=session) is None
+        assert backend.get(scope, "file_count", session=session) == "42"
+
+    def test_delete_is_noop_for_missing_key(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        backend.delete(scope, "nonexistent", session=session)
+
+    def test_clear_removes_all_keys(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset.id)
+        for key in ("watermark", "file_count", "last_error"):
+            backend.set(scope, key, f"val_{key}", session=session)
+        session.flush()
+
+        backend.clear(scope, session=session)
+        session.flush()
+
+        for key in ("watermark", "file_count", "last_error"):
+            assert backend.get(scope, key, session=session) is None
+
+    def test_different_assets_are_isolated(
+        self, session: Session, backend: MetastoreStateBackend, asset: AssetModel
+    ):
+        asset2 = AssetModel(uri="s3://bucket/other", name="other_asset", group="test")
+        session.add(asset2)
+        session.flush()
+
+        scope1 = AssetScope(asset_id=asset.id)
+        scope2 = AssetScope(asset_id=asset2.id)
+
+        backend.set(scope1, "watermark", "asset1_value", session=session)
+        session.flush()
+
+        assert backend.get(scope2, "watermark", session=session) is None
+
+
+@pytest.mark.asyncio(loop_scope="class")
+class TestMetastoreStateBackendAsync:
+    async def test_aset_and_aget_task_roundtrip(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        result = await backend.aget(scope, "job_id")
+        assert result == "app_async"
+
+    async def test_adelete_task_removes_key(self, backend: MetastoreStateBackend, dag_run_committed: DagRun):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        await backend.adelete(scope, "job_id")
+        assert await backend.aget(scope, "job_id") is None
+
+    async def test_aclear_task_removes_all_keys(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        await backend.aset(scope, "job_id", "app_async")
+        await backend.aset(scope, "checkpoint", "step_1")
+        await backend.aclear(scope)
+        assert await backend.aget(scope, "job_id") is None
+        assert await backend.aget(scope, "checkpoint") is None
+
+    async def test_aset_and_aget_asset_roundtrip(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        result = await backend.aget(scope, "watermark")
+        assert result == "2026-04-27T00:00:00Z"
+
+    async def test_adelete_asset_removes_key(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        await backend.adelete(scope, "watermark")
+        assert await backend.aget(scope, "watermark") is None
+
+    async def test_aclear_asset_removes_all_keys(
+        self, backend: MetastoreStateBackend, asset_committed: AssetModel
+    ):
+        scope = AssetScope(asset_id=asset_committed.id)
+        await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z")
+        await backend.aset(scope, "file_count", "42")
+        await backend.aclear(scope)
+        assert await backend.aget(scope, "watermark") is None
+        assert await backend.aget(scope, "file_count") is None
+
+    async def test_aset_task_raises_for_missing_dag_run(self, backend: MetastoreStateBackend):
+        scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID)
+        with pytest.raises(ValueError, match="No DagRun found"):
+            await backend.aset(scope, "job_id", "app_async")
+
+
+class TestResolveStateBackend:
+    def test_default_resolves_to_metastore_backend(self):

Review Comment:
   Thanks, fixed both. Renamed `test_default_resolves_to_metastore_backend` to `test_resolve_returns_configured_backend` and wrapped it in `conf_vars` so it doesn't depend on any ambient config value. I have also added a `test_empty_backend_raises_value_error` and I am going with raise (not silent fallback) since an explicit `backend = ""` is a misconfiguration and `config.yml` sets the real default so you will only land here if something actively blanked the config



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154265389


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:

Review Comment:
   Handled in dbcda100fb



##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: fde9ed84d07b
+Revises: 9fabad868fdb
+Create Date: 2026-04-27 09:29:20.165461
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "fde9ed84d07b"
+down_revision = "9fabad868fdb"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+    """Apply add task_state and asset_state tables."""
+    op.create_table(
+        "asset_state",
+        sa.Column("asset_id", sa.Integer(), nullable=False),
+        sa.Column("key", sa.String(length=512), nullable=False),
+        sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
+        sa.Column("updated_at", UtcDateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["asset_id"], ["asset.id"], name="asset_state_asset_fkey", ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"),
+    )
+    op.create_table(
+        "task_state",
+        sa.Column("dag_run_id", sa.Integer(), nullable=False),
+        sa.Column("task_id", sa.String(length=250), nullable=False),

Review Comment:
   Handled in dbcda100fb



##########
airflow-core/src/airflow/models/task_state.py:
##########
@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text
+from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow._shared.timezones import timezone
+from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class TaskStateModel(Base):
+    """
+    Persists key/value state for a task within a single DAG run.
+
+    Scoped to (dag_run_id, task_id, map_index). Retries of the same task share
+    the same rows — that is the point. Different DAG runs have different dag_run_id
+    values so they get independent namespaces automatically.
+    """
+
+    __tablename__ = "task_state"
+
+    dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True)

Review Comment:
   Handled in dbcda100fb



##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:

Review Comment:
   Handled in dbcda100fb



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154114471


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:
+            case TaskScope():
+                return self._get_task_state(scope, key, session=session)
+            case AssetScope():
+                return self._get_asset_state(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._set_task_state(scope, key, value, session=session)
+            case AssetScope():
+                self._set_asset_state(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._delete_task_state(scope, key, session=session)
+            case AssetScope():
+                self._delete_asset_state(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        match scope:
+            case TaskScope():
+                self._clear_task_state(scope, session=session)
+            case AssetScope():
+                self._clear_asset_state(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    return await self._aget_task_state(scope, key, session=session)
+                case AssetScope():
+                    return await self._aget_asset_state(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aset_task_state(scope, key, value, session=session)
+                case AssetScope():
+                    await self._aset_asset_state(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._adelete_task_state(scope, key, session=session)
+                case AssetScope():
+                    await self._adelete_asset_state(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            match scope:
+                case TaskScope():
+                    await self._aclear_task_state(scope, session=session)
+                case AssetScope():
+                    await self._aclear_asset_state(scope, session=session)
+
+    def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(
+            select(DagRun.id).where(
+                DagRun.dag_id == scope.dag_id,
+                DagRun.run_id == scope.run_id,
+            )
+        )
+        if dag_run_id is None:
+            raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}")
+        now = timezone.utcnow()
+        values = dict(
+            dag_run_id=dag_run_id,
+            dag_id=scope.dag_id,
+            run_id=scope.run_id,
+            task_id=scope.task_id,
+            map_index=scope.map_index,
+            key=key,
+            value=value,
+            updated_at=now,
+        )
+        stmt = _build_upsert_stmt(
+            get_dialect_name(session),
+            TaskStateModel,
+            ["dag_run_id", "task_id", "map_index", "key"],
+            values,
+            dict(value=value, updated_at=now),
+        )
+        session.execute(stmt)
+
+    def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None:
+        session.execute(
+            delete(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+
+    def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None:

Review Comment:
   Undecided what the best option here is but I think I concur most with `-1` means "all map indices for this task" for mapped tasks. ie: if not specified, we delete all.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154063774


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,334 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+def _build_upsert_stmt(
+    dialect: str | None,
+    model: type,
+    conflict_cols: list[str],
+    values: dict,
+    update_fields: dict,
+):
+    """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement."""
+    stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert
+    if dialect == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+        stmt = pg_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    elif dialect == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        stmt = mysql_insert(model).values(**values)
+        stmt = stmt.on_duplicate_key_update(**update_fields)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(model).values(**values)
+        stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields)
+    return stmt
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        match scope:

Review Comment:
   Ooh nice, didn't know about `assert_never`, handling it with that:
   ```python
   value = 10
   match value:
       case 1:
           print("one")
       case 2:
           print("two")
       case _:
           assert_never(value)
   Traceback (most recent call last):
     File "/Users/amoghdesai/Documents/OSS/repos/airflow/.venv/lib/python3.13/site-packages/IPython/core/interactiveshell.py", line 3747, in run_code
       exec(code_obj, self.user_global_ns, self.user_ns)
       ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     File "<ipython-input-3-21b042fd93f7>", line 8, in <module>
       assert_never(value)
       ~~~~~~~~~~~~^^^^^^^
     File "/opt/homebrew/Cellar/python@3.13/3.13.3_1/Frameworks/Python.framework/Versions/3.13/lib/python3.13/typing.py", line 2627, in assert_never
       raise AssertionError(f"Expected code to be unreachable, but got: {value}")
   AssertionError: Expected code to be unreachable, but got: 10
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3154080748


##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add task_state and asset_state tables.
+
+Revision ID: fde9ed84d07b
+Revises: 9fabad868fdb
+Create Date: 2026-04-27 09:29:20.165461
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import mysql
+
+from airflow.utils.sqlalchemy import UtcDateTime
+
+# revision identifiers, used by Alembic.
+revision = "fde9ed84d07b"
+down_revision = "9fabad868fdb"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+    """Apply add task_state and asset_state tables."""
+    op.create_table(
+        "asset_state",
+        sa.Column("asset_id", sa.Integer(), nullable=False),
+        sa.Column("key", sa.String(length=512), nullable=False),
+        sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
+        sa.Column("updated_at", UtcDateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["asset_id"], ["asset.id"], name="asset_state_asset_fkey", ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"),
+    )
+    op.create_table(
+        "task_state",
+        sa.Column("dag_run_id", sa.Integer(), nullable=False),
+        sa.Column("task_id", sa.String(length=250), nullable=False),

Review Comment:
   Good catch. Switched both the migration `(task_id, dag_id, run_id columns in task_state)` and the ORM model to use `StringID(), same pattern as migration number: 0108. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #65759:
URL: https://github.com/apache/airflow/pull/65759#discussion_r3145967145


##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -0,0 +1,397 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import delete, select
+
+from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope
+from airflow._shared.timezones import timezone
+from airflow.models.asset_state import AssetStateModel
+from airflow.models.dagrun import DagRun
+from airflow.models.task_state import TaskStateModel
+from airflow.utils.session import NEW_SESSION, create_session_async, provide_session
+from airflow.utils.sqlalchemy import get_dialect_name
+
+if TYPE_CHECKING:
+    from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
+    from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
+    from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
+
+
+class MetastoreStateBackend(BaseStateBackend):
+    """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
+
+    @provide_session
+    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+        if isinstance(scope, TaskScope):
+            return self._get_task(scope, key, session=session)
+        return self._get_asset(scope, key, session=session)
+
+    @provide_session
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._set_task(scope, key, value, session=session)
+        else:
+            self._set_asset(scope, key, value, session=session)
+
+    @provide_session
+    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._delete_task(scope, key, session=session)
+        else:
+            self._delete_asset(scope, key, session=session)
+
+    @provide_session
+    def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None:
+        if isinstance(scope, TaskScope):
+            self._clear_task(scope, session=session)
+        else:
+            self._clear_asset(scope, session=session)
+
+    async def aget(self, scope: StateScope, key: str) -> str | None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                return await self._aget_task(scope, key, session=session)
+            return await self._aget_asset(scope, key, session=session)
+
+    async def aset(self, scope: StateScope, key: str, value: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aset_task(scope, key, value, session=session)
+            else:
+                await self._aset_asset(scope, key, value, session=session)
+
+    async def adelete(self, scope: StateScope, key: str) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._adelete_task(scope, key, session=session)
+            else:
+                await self._adelete_asset(scope, key, session=session)
+
+    async def aclear(self, scope: StateScope) -> None:
+        async with create_session_async() as session:
+            if isinstance(scope, TaskScope):
+                await self._aclear_task(scope, session=session)
+            else:
+                await self._aclear_asset(scope, session=session)
+
+    def _get_task(self, scope: TaskScope, key: str, *, session: Session) -> str | None:
+        row = session.scalar(
+            select(TaskStateModel).where(
+                TaskStateModel.dag_id == scope.dag_id,
+                TaskStateModel.run_id == scope.run_id,
+                TaskStateModel.task_id == scope.task_id,
+                TaskStateModel.map_index == scope.map_index,
+                TaskStateModel.key == key,
+            )
+        )
+        return row.value if row is not None else None
+
+    def _set_task(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None:
+        dag_run_id = session.scalar(

Review Comment:
   Handled in 2b29b8098c
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] AIP-103: Setting up foundation with ORM and Backend Abstraction [airflow]

Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #65759:
URL: https://github.com/apache/airflow/pull/65759


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org