#
# 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 logging
import warnings
from collections.abc import Generator
from datetime import timedelta
from typing import TYPE_CHECKING
from unittest import mock
from unittest.mock import patch

import pytest
from sqlalchemy import delete, func, inspect as sa_inspect, select
from sqlalchemy.exc import OperationalError, SAWarning

import airflow.dag_processing.collection
from airflow._shared.timezones import timezone as tz
from airflow.configuration import conf
from airflow.dag_processing.collection import (
    AssetModelOperation,
    DagModelOperation,
    _get_latest_runs_stmt,
    _get_latest_runs_stmt_partitioned,
    _update_dag_tags,
    update_dag_parsing_results_in_db,
)
from airflow.exceptions import SerializationError
from airflow.models import DagModel, DagRun
from airflow.models.asset import (
    AssetActive,
    AssetModel,
    DagScheduleAssetNameReference,
    DagScheduleAssetUriReference,
)
from airflow.models.dag import DagTag
from airflow.models.dagbundle import DagBundleModel
from airflow.models.errors import ParseImportError
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.trigger import Trigger
from airflow.partition_mappers.base import RollupMapper
from airflow.partition_mappers.temporal import StartOfDayMapper
from airflow.partition_mappers.window import DayWindow
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.triggers.file import FileDeleteTrigger
from airflow.sdk import DAG, Asset, AssetAlias, AssetAll, AssetWatcher
from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
from airflow.serialization.definitions.assets import SerializedAsset
from airflow.serialization.encoders import encode_trigger, ensure_serialized_asset
from airflow.serialization.serialized_objects import LazyDeserializedDAG
from airflow.timetables.simple import PartitionedAtRuntime
from airflow.triggers.base import BaseEventTrigger
from airflow.utils.types import DagRunType

from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
    clear_db_assets,
    clear_db_dags,
    clear_db_import_errors,
    clear_db_serialized_dags,
    clear_db_triggers,
)

if TYPE_CHECKING:
    from kgb import SpyAgency

mark_fab_auth_manager_test = pytest.mark.skipif(
    condition="FabAuthManager" not in conf.get("core", "auth_manager"),
    reason="This is only for FabAuthManager. Please set the environment variable `AIRFLOW__CORE__AUTH_MANAGER` to `airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager` in `files/airflow-breeze-config/environment_variables.env` before running breeze shell. To run the test, add the flag `--keep-env-variables` to the pytest command.",
)


def test_statement_latest_runs_one_dag():
    with warnings.catch_warnings():
        warnings.simplefilter("error", category=SAWarning)

        stmt = _get_latest_runs_stmt(["fake-dag"])
        compiled_stmt = str(stmt.compile())
        actual = [x.strip() for x in compiled_stmt.splitlines()]
        expected = [
            "SELECT dag_run.id, dag_run.dag_id, dag_run.logical_date, dag_run.data_interval_start, "
            "dag_run.data_interval_end, dag_run.run_after, dag_run.partition_key, "
            "dag_run.partition_date",
            "FROM dag_run",
            "WHERE dag_run.dag_id = :dag_id_1 AND dag_run.logical_date = ("
            "SELECT max(dag_run.logical_date) AS max_logical_date",
            "FROM dag_run",
            "WHERE dag_run.dag_id = :dag_id_2 AND dag_run.run_type IN (__[POSTCOMPILE_run_type_1]))",
        ]
        assert actual == expected, compiled_stmt


@pytest.mark.db_test
def test_statement_latest_runs_loads_timetable_fields(dag_maker, session):
    with dag_maker("fake-dag", schedule=None):
        pass
    dag_maker.sync_dagbag_to_db()

    logical_date = tz.datetime(2025, 1, 1)
    run_after = tz.datetime(2025, 1, 2)

    dag_maker.create_dagrun(
        run_id="latest-run",
        logical_date=logical_date,
        data_interval=(logical_date, run_after),
        run_type=DagRunType.SCHEDULED,
        run_after=run_after,
        session=session,
    )
    session.flush()
    session.expunge_all()  # Ensure we load from DB, not from session cache

    latest = session.scalar(_get_latest_runs_stmt("fake-dag"))
    assert latest is not None
    assert {"run_after", "partition_date", "partition_key"}.isdisjoint(sa_inspect(latest).unloaded)
    assert latest.run_after == run_after
    assert latest.partition_key is None
    assert latest.partition_date is None


@pytest.mark.db_test
def test_statement_latest_runs_partitioned_sorted_by_partition_date(dag_maker, session):
    with dag_maker("fake-dag", schedule=PartitionedAtRuntime()):
        pass
    dag_maker.sync_dagbag_to_db()
    for i, (run_id, partition_key, partition_date) in enumerate(
        (
            ("newest-partition-date", "2025-01-02", tz.datetime(2025, 1, 2)),
            ("older-partition-date", "2025-01-01", tz.datetime(2025, 1, 1)),
            ("null-partition-date", "not-a-time-based-partition", None),
        )
    ):
        dag_maker.create_dagrun(
            run_id=run_id,
            logical_date=None,
            data_interval=None,
            run_type=DagRunType.SCHEDULED,
            run_after=tz.datetime(2025, 1, 1 + i),
            partition_key=partition_key,
            partition_date=partition_date,
            session=session,
        )

    session.flush()
    session.expunge_all()  # Ensure we load from DB, not from session cache

    latest = session.scalar(_get_latest_runs_stmt_partitioned("fake-dag"))
    assert latest is not None
    assert {"run_after", "partition_date", "partition_key"}.isdisjoint(sa_inspect(latest).unloaded)
    assert latest.run_after == tz.datetime(2025, 1, 1)
    assert latest.partition_key == "2025-01-02"
    assert latest.partition_date == tz.datetime(2025, 1, 2)


@pytest.mark.db_test
class TestAssetModelOperation:
    @staticmethod
    def clean_db():
        clear_db_dags()
        clear_db_assets()
        clear_db_triggers()

    @pytest.fixture(autouse=True)
    def per_test(self) -> Generator:
        self.clean_db()
        yield
        self.clean_db()

    @pytest.mark.usefixtures("testing_dag_bundle")
    def test_sync_assets_preserves_access_control_from_other_bundle(self, dag_maker, session):
        """When a producer bundle (without access_control) is synced after a consumer bundle
        (with access_control), the stored access control fields must not be wiped out."""
        from airflow.models.asset import DagScheduleAssetReference, TaskOutletAssetReference
        from airflow.sdk import AssetAccessControl

        # First sync: consumer bundle sets access_control on the asset (producer_teams on schedule side).
        consumer_asset = Asset(
            "shared_asset",
            access_control=AssetAccessControl(producer_teams=["team1", "team2"], allow_global=False),
        )
        with dag_maker(dag_id="consumer_dag", schedule=[consumer_asset]) as consumer_dag:
            EmptyOperator(task_id="mytask")

        consumer_dags = {consumer_dag.dag_id: LazyDeserializedDAG.from_dag(consumer_dag)}
        orm_dags = DagModelOperation(consumer_dags, "testing", None).add_dags(session=session)
        asset_op = AssetModelOperation.collect(consumer_dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        asset_op.add_dag_asset_references(orm_dags, orm_assets, session=session)
        session.flush()

        ref = session.scalar(
            select(DagScheduleAssetReference).where(DagScheduleAssetReference.dag_id == "consumer_dag")
        )
        assert ref.allow_producer_teams == ["team1", "team2"]
        assert ref.allow_global_producers is False

        # Second sync: producer bundle references the same asset with consumer_teams on the outlet.
        producer_asset = Asset(
            "shared_asset",
            access_control=AssetAccessControl(consumer_teams=["team_ml"], allow_global=False),
        )
        with dag_maker(dag_id="producer_dag", schedule="@once") as producer_dag:
            EmptyOperator(task_id="produce", outlets=[producer_asset])

        producer_dags = {producer_dag.dag_id: LazyDeserializedDAG.from_dag(producer_dag)}
        producer_orm_dags = DagModelOperation(producer_dags, "testing", None).add_dags(session=session)
        asset_op = AssetModelOperation.collect(producer_dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        asset_op.add_task_asset_references(producer_orm_dags, orm_assets, session=session)
        session.flush()

        # Consumer's schedule-side access control must still be preserved.
        session.expire(ref)
        assert ref.allow_producer_teams == ["team1", "team2"]
        assert ref.allow_global_producers is False

        # Producer's outlet-side access control must be stored.
        outlet_ref = session.scalar(
            select(TaskOutletAssetReference).where(TaskOutletAssetReference.dag_id == "producer_dag")
        )
        assert outlet_ref.allow_consumer_teams == ["team_ml"]
        assert outlet_ref.allow_global_consumers is False

    @pytest.mark.usefixtures("testing_dag_bundle")
    def test_add_task_outlet_asset_references_updates_consumer_teams_on_change(self, dag_maker, session):
        """When access_control changes, existing outlet references are updated in place."""
        from airflow.models.asset import TaskOutletAssetReference
        from airflow.sdk import AssetAccessControl

        asset = Asset(
            "evolving_asset",
            access_control=AssetAccessControl(consumer_teams=["team_old"], allow_global=True),
        )

        with dag_maker(dag_id="evolving_producer", schedule="@once") as dag:
            EmptyOperator(task_id="produce", outlets=[asset])

        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).add_dags(session=session)
        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        asset_op.add_task_asset_references(orm_dags, orm_assets, session=session)
        session.flush()

        ref = session.scalar(
            select(TaskOutletAssetReference).where(TaskOutletAssetReference.dag_id == "evolving_producer")
        )
        assert ref.allow_consumer_teams == ["team_old"]
        assert ref.allow_global_consumers is True

        # Change access_control and re-sync.
        asset.access_control = AssetAccessControl(consumer_teams=["team_new"], allow_global=False)
        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).find_orm_dags(session=session)
        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        asset_op.add_task_asset_references(orm_dags, orm_assets, session=session)
        session.flush()

        session.expire(ref)
        assert ref.allow_consumer_teams == ["team_new"]
        assert ref.allow_global_consumers is False

    @pytest.mark.usefixtures("testing_dag_bundle")
    def test_add_task_outlet_asset_references_defaults_when_no_access_control(self, dag_maker, session):
        """Outlet references default to None consumer_teams and allow_global_consumers=True."""
        from airflow.models.asset import TaskOutletAssetReference

        asset = Asset("plain_asset")

        with dag_maker(dag_id="plain_producer_dag", schedule="@once") as dag:
            EmptyOperator(task_id="produce", outlets=[asset])

        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).add_dags(session=session)
        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        asset_op.add_task_asset_references(orm_dags, orm_assets, session=session)
        session.flush()

        ref = session.scalar(
            select(TaskOutletAssetReference).where(TaskOutletAssetReference.dag_id == "plain_producer_dag")
        )
        assert ref is not None
        assert ref.allow_consumer_teams is None
        assert ref.allow_global_consumers is True

    @pytest.mark.parametrize(
        ("is_active", "is_paused", "expected_num_triggers"),
        [
            (True, True, 0),
            (True, False, 1),
            (False, True, 0),
            (False, False, 0),
        ],
    )
    @pytest.mark.usefixtures("testing_dag_bundle")
    def test_add_asset_trigger_references(
        self, dag_maker, session, is_active, is_paused, expected_num_triggers
    ):
        asset = Asset(
            "test_add_asset_trigger_references_asset",
            watchers=[AssetWatcher(name="test", trigger=FileDeleteTrigger(mock.Mock()))],
        )

        with dag_maker(dag_id="test_add_asset_trigger_references_dag", schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).add_dags(session=session)

        # Simulate dag unpause and deletion.
        dag_model = orm_dags[dag.dag_id]
        dag_model.is_stale = not is_active
        dag_model.is_paused = is_paused

        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()

        asset_op.add_dag_asset_references(orm_dags, orm_assets, session=session)
        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
        asset_op.add_asset_trigger_references(orm_assets, session=session)
        session.flush()

        asset_model = session.scalars(select(AssetModel)).one()
        assert len(asset_model.triggers) == expected_num_triggers

    @pytest.mark.usefixtures("testing_dag_bundle")
    @pytest.mark.parametrize(
        ("use_team", "expected"),
        [
            pytest.param(True, "testing", id="with-team"),
            pytest.param(False, None, id="no-team"),
        ],
    )
    def test_add_asset_trigger_references_populates_team_name(
        self, dag_maker, session, testing_team, use_team, expected
    ):
        asset = Asset(
            "test_trigger_team_asset",
            watchers=[AssetWatcher(name="watcher", trigger=FileDeleteTrigger(mock.Mock()))],
        )

        with dag_maker(dag_id="test_trigger_team_dag", schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        # Use raw DagModelOperation (not dag_maker's bulk_write_to_db) to control team_name
        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).add_dags(session=session)
        orm_dags[dag.dag_id].is_stale = False
        orm_dags[dag.dag_id].is_paused = False
        session.flush()

        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        asset_op.add_dag_asset_references(orm_dags, orm_assets, session=session)
        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
        session.flush()

        # Clear any triggers created by dag_maker's bulk_write_to_db
        session.execute(delete(Trigger))
        for asset_model in orm_assets.values():
            asset_model.watchers = []
        session.flush()

        team_name = testing_team.name if use_team else None
        asset_op.add_asset_trigger_references(orm_assets, team_name=team_name, session=session)
        session.flush()

        triggers = session.scalars(select(Trigger)).all()
        assert len(triggers) == 1
        assert triggers[0].team_name == expected

    @pytest.mark.usefixtures("testing_dag_bundle")
    def test_add_asset_trigger_references_hash_consistency(self, dag_maker, session):
        """Trigger hash from the DAG-parsed path must equal the hash computed
        from the DB-stored Trigger row.  A mismatch causes the scheduler to
        recreate trigger rows on every heartbeat.
        """
        trigger = FileDeleteTrigger(filepath="/tmp/test.txt", poke_interval=5.0)
        asset = Asset(
            "test_hash_consistency_asset",
            watchers=[AssetWatcher(name="file_watcher", trigger=trigger)],
        )

        with dag_maker(dag_id="test_hash_consistency_dag", schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).add_dags(session=session)
        orm_dags[dag.dag_id].is_paused = False

        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()

        asset_op.add_dag_asset_references(orm_dags, orm_assets, session=session)
        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
        asset_op.add_asset_trigger_references(orm_assets, session=session)
        session.flush()

        # DAG-side hash (same computation as add_asset_trigger_references line 1025)
        encoded = encode_trigger(trigger)
        dag_hash = BaseEventTrigger.hash(encoded["classpath"], encoded["kwargs"])

        # DB-side: expire and re-load the Trigger row to force a real DB read
        asset_model = session.scalars(select(AssetModel)).one()
        assert len(asset_model.triggers) == 1
        orm_trigger = asset_model.triggers[0]
        trigger_id = orm_trigger.id
        session.expire(orm_trigger)
        reloaded = session.get(Trigger, trigger_id)

        # DB-side hash (same computation as add_asset_trigger_references line 1033)
        db_hash = BaseEventTrigger.hash(reloaded.classpath, reloaded.kwargs)

        assert dag_hash == db_hash

    @pytest.mark.usefixtures("testing_dag_bundle")
    def test_add_asset_trigger_references_idempotent(self, dag_maker, session):
        """Calling add_asset_trigger_references twice with the same trigger
        must not create duplicate rows.
        """
        trigger = FileDeleteTrigger(filepath="/tmp/test.txt", poke_interval=5.0)
        asset = Asset(
            "test_idempotent_asset",
            watchers=[AssetWatcher(name="file_watcher", trigger=trigger)],
        )

        with dag_maker(dag_id="test_idempotent_dag", schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        dags = {dag.dag_id: LazyDeserializedDAG.from_dag(dag)}
        orm_dags = DagModelOperation(dags, "testing", None).add_dags(session=session)
        orm_dags[dag.dag_id].is_paused = False

        asset_op = AssetModelOperation.collect(dags)
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()

        asset_op.add_dag_asset_references(orm_dags, orm_assets, session=session)
        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)

        # First call — creates the trigger
        asset_op.add_asset_trigger_references(orm_assets, session=session)
        session.flush()
        count_after_first = session.scalar(select(func.count(Trigger.id)))

        # Second call — should be a no-op (hashes match, no diff)
        asset_op.add_asset_trigger_references(orm_assets, session=session)
        session.flush()
        count_after_second = session.scalar(select(func.count(Trigger.id)))

        assert count_after_first == count_after_second

    @pytest.mark.parametrize(
        ("schedule", "model", "columns", "expected"),
        [
            pytest.param(
                Asset.ref(name="name1"),
                DagScheduleAssetNameReference,
                (DagScheduleAssetNameReference.name, DagScheduleAssetNameReference.dag_id),
                [("name1", "test")],
                id="name-ref",
            ),
            pytest.param(
                Asset.ref(uri="foo://1"),
                DagScheduleAssetUriReference,
                (DagScheduleAssetUriReference.uri, DagScheduleAssetUriReference.dag_id),
                [("foo://1", "test")],
                id="uri-ref",
            ),
        ],
    )
    def test_add_dag_asset_name_uri_references(self, dag_maker, session, schedule, model, columns, expected):
        with dag_maker(dag_id="test", schedule=schedule, session=session) as dag:
            pass

        op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        op.add_dag_asset_name_uri_references(session=session)
        assert session.execute(select(*columns)).all() == expected

    def test_change_asset_property_sync_group(self, dag_maker, session):
        asset = Asset("myasset", group="old_group")
        with dag_maker(schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        assert len(orm_assets) == 1
        assert next(iter(orm_assets.values())).group == "old_group"

        # Parser should pick up group change.
        asset.group = "new_group"
        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        assert len(orm_assets) == 1
        assert next(iter(orm_assets.values())).group == "new_group"

    def test_change_asset_property_sync_extra(self, dag_maker, session):
        asset = Asset("myasset", extra={"foo": "old"})
        with dag_maker(schedule=asset) as dag:
            EmptyOperator(task_id="mytask")

        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        assert len(orm_assets) == 1
        assert next(iter(orm_assets.values())).extra == {"foo": "old"}

        # Parser should pick up extra change.
        asset.extra = {"foo": "new"}
        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        assert len(orm_assets) == 1
        assert next(iter(orm_assets.values())).extra == {"foo": "new"}

    def test_change_asset_alias_property_sync_group(self, dag_maker, session):
        alias = AssetAlias("myalias", group="old_group")
        with dag_maker(schedule=alias) as dag:
            EmptyOperator(task_id="mytask")

        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_aliases = asset_op.sync_asset_aliases(session=session)
        assert len(orm_aliases) == 1
        assert next(iter(orm_aliases.values())).group == "old_group"

        # Parser should pick up group change.
        alias.group = "new_group"
        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_aliases = asset_op.sync_asset_aliases(session=session)
        assert len(orm_aliases) == 1
        assert next(iter(orm_aliases.values())).group == "new_group"


@pytest.mark.db_test
@pytest.mark.want_activate_assets(False)
class TestAssetModelOperationSyncAssetActive:
    @staticmethod
    def clean_db():
        clear_db_dags()
        clear_db_assets()
        clear_db_triggers()

    @pytest.fixture(autouse=True)
    def per_test(self) -> Generator:
        self.clean_db()
        yield
        self.clean_db()

    def test_add_asset_activate(self, dag_maker, session):
        asset = Asset("myasset", "file://myasset/", group="old_group")
        with dag_maker(schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        assert len(orm_assets) == 1

        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
        session.flush()
        assert orm_assets["myasset", "file://myasset/"].active is not None

    def test_add_asset_activate_already_exists(self, dag_maker, session):
        asset = Asset(name="myasset", uri="file://myasset/", group="old_group")

        # Set up existing parsing result.
        serialized_asset = ensure_serialized_asset(asset)
        orm_asset = AssetModel.from_serialized(serialized_asset)
        session.add(orm_asset)
        session.add(AssetActive.for_asset(serialized_asset))
        session.flush()

        with dag_maker(schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        assert orm_assets == {("myasset", "file://myasset/"): orm_asset}

        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
        session.flush()
        assert orm_assets["myasset", "file://myasset/"].active is not None, "should pick up existing active"

    @pytest.mark.parametrize(
        "existing_assets",
        [
            pytest.param([SerializedAsset("myasset", "file://different/asset", "", {}, [])], id="name"),
            pytest.param([SerializedAsset("another", "file://myasset/", "", {}, [])], id="uri"),
        ],
    )
    def test_add_asset_activate_conflict(self, dag_maker, session, existing_assets):
        session.add_all(AssetModel.from_serialized(a) for a in existing_assets)
        session.flush()
        session.add_all(AssetActive.for_asset(a) for a in existing_assets)
        session.flush()

        asset = Asset(name="myasset", uri="file://myasset/", group="old_group")
        with dag_maker(schedule=[asset]) as dag:
            EmptyOperator(task_id="mytask")

        asset_op = AssetModelOperation.collect({dag.dag_id: LazyDeserializedDAG.from_dag(dag)})
        orm_assets = asset_op.sync_assets(session=session)
        session.flush()
        assert len(orm_assets) == 1

        asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
        session.flush()
        assert orm_assets["myasset", "file://myasset/"].active is None, "should not activate due to conflict"


@pytest.mark.need_serialized_dag
@pytest.mark.db_test
class TestUpdateDagParsingResults:
    """Tests centred around the ``update_dag_parsing_results_in_db`` function."""

    @pytest.fixture
    def clean_db(self, session):
        yield
        clear_db_serialized_dags()
        clear_db_dags()
        clear_db_import_errors()

    @pytest.fixture(name="dag_import_error_listener")
    def _dag_import_error_listener(self, listener_manager):
        from unit.listeners import dag_import_error_listener

        listener_manager(dag_import_error_listener)
        yield dag_import_error_listener
        dag_import_error_listener.clear()

    @mark_fab_auth_manager_test
    @conf_vars({("core", "min_serialized_dag_update_interval"): "5"})
    @pytest.mark.usefixtures("clean_db")  # sync_perms in fab has bad session commit hygiene
    def test_sync_perms_syncs_dag_specific_perms_on_update(
        self, monkeypatch, spy_agency: SpyAgency, session, time_machine, testing_dag_bundle
    ):
        """Test DAG-specific permissions are synced when a DAG is new or updated"""
        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert serialized_dags_count == 0

        time_machine.move_to(tz.datetime(2020, 1, 5, 0, 0, 0), tick=False)

        dag = DAG(dag_id="test")

        sync_perms_spy = spy_agency.spy_on(
            airflow.dag_processing.collection._sync_dag_perms,
            call_original=False,
        )

        def _sync_to_db():
            sync_perms_spy.reset_calls()
            time_machine.shift(20)

            update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session)

        _sync_to_db()
        spy_agency.assert_spy_called_with(sync_perms_spy, dag, session=session)

        # DAG isn't updated
        _sync_to_db()
        # `_sync_dag_perms` should be called even the DAG isn't updated. Otherwise, any import error will not show up until DAG is updated.
        spy_agency.assert_spy_called_with(sync_perms_spy, dag, session=session)

        # DAG is updated
        dag.tags = {"new_tag"}
        _sync_to_db()
        spy_agency.assert_spy_called_with(sync_perms_spy, dag, session=session)

        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))

    @patch.object(SerializedDagModel, "write_dag")
    @patch("airflow.serialization.definitions.dag.SerializedDAG.bulk_write_to_db")
    def test_sync_to_db_is_retried(
        self, mock_bulk_write_to_db, mock_s10n_write_dag, testing_dag_bundle, session
    ):
        """Test that important DB operations in db sync are retried on OperationalError"""
        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert serialized_dags_count == 0
        mock_dag = mock.MagicMock()
        dags = [mock_dag]

        op_error = OperationalError(statement=mock.ANY, params=mock.ANY, orig=mock.ANY)

        # Mock error for the first 2 tries and a successful third try
        side_effect = [op_error, op_error, mock.ANY]

        mock_bulk_write_to_db.side_effect = side_effect

        mock_session = mock.MagicMock()
        update_dag_parsing_results_in_db(
            "testing",
            None,
            dags=dags,
            import_errors={},
            parse_duration=None,
            warnings=set(),
            session=mock_session,
        )

        # Test that 3 attempts were made to run 'DAG.bulk_write_to_db' successfully
        mock_bulk_write_to_db.assert_has_calls(
            [
                mock.call("testing", None, mock.ANY, None, session=mock.ANY),
                mock.call("testing", None, mock.ANY, None, session=mock.ANY),
                mock.call("testing", None, mock.ANY, None, session=mock.ANY),
            ]
        )
        # Assert that rollback is called twice (i.e. whenever OperationalError occurs)
        mock_session.rollback.assert_has_calls([mock.call(), mock.call()])
        # Check that 'SerializedDagModel.write_dag' is also called
        # Only called once since the other two times the 'DAG.bulk_write_to_db' error'd
        # and the session was roll-backed before even reaching 'SerializedDagModel.write_dag'
        mock_s10n_write_dag.assert_has_calls(
            [
                mock.call(
                    mock_dag,
                    bundle_name="testing",
                    bundle_version=None,
                    version_data=None,
                    min_update_interval=mock.ANY,
                    session=mock_session,
                    _prefetched=mock.ANY,
                ),
            ]
        )

        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert serialized_dags_count == 0

    def test_serialized_dags_are_written_to_db_on_sync(self, testing_dag_bundle, session):
        """Test DAGs are Serialized and written to DB when parsing result is updated"""
        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert serialized_dags_count == 0

        dag = DAG(dag_id="test")

        update_dag_parsing_results_in_db(
            bundle_name="testing",
            bundle_version=None,
            dags=[LazyDeserializedDAG.from_dag(dag)],
            import_errors={},
            parse_duration=None,
            warnings=set(),
            session=session,
        )

        new_serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert new_serialized_dags_count == 1

    def test_parse_time_written_to_db_on_sync(self, testing_dag_bundle, session):
        """Test that the parse time is correctly written to the DB after parsing"""

        parse_duration = 1.25
        dag = DAG(dag_id="test")
        update_dag_parsing_results_in_db("testing", None, [dag], dict(), parse_duration, set(), session)

        dag_model: DagModel = session.get(DagModel, (dag.dag_id,))
        assert dag_model.last_parse_duration == parse_duration

    @patch.object(ParseImportError, "full_file_path")
    @patch.object(SerializedDagModel, "write_dag")
    @pytest.mark.usefixtures("clean_db")
    def test_serialized_dag_errors_are_import_errors(
        self, mock_serialize, mock_full_path, caplog, session, dag_import_error_listener, testing_dag_bundle
    ):
        """
        Test that errors serializing a DAG are recorded as import_errors in the DB
        """
        mock_serialize.side_effect = SerializationError
        caplog.set_level(logging.ERROR)

        dag = DAG(dag_id="test")
        dag.fileloc = "abc.py"
        dag.relative_fileloc = "abc.py"
        mock_full_path.return_value = "abc.py"

        import_errors = {}
        update_dag_parsing_results_in_db(
            "testing", None, [dag], import_errors, None, set(), session, files_parsed={("testing", "abc.py")}
        )
        assert "SerializationError" in caplog.text

        # Should have been edited in place
        err = import_errors.get(("testing", dag.relative_fileloc))
        assert "SerializationError" in err
        dag_model: DagModel = session.get(DagModel, (dag.dag_id,))
        assert dag_model.has_import_errors is True

        import_errors = session.scalars(select(ParseImportError)).all()

        assert len(import_errors) == 1
        import_error = import_errors[0]
        assert import_error.filename == dag.relative_fileloc
        assert "SerializationError" in import_error.stacktrace

        # Ensure the listener was notified
        assert len(dag_import_error_listener.new) == 1
        assert len(dag_import_error_listener.existing) == 0
        assert dag_import_error_listener.new["abc.py"] == import_error.stacktrace

    @patch.object(ParseImportError, "full_file_path")
    @mark_fab_auth_manager_test
    @conf_vars({("core", "min_serialized_dag_update_interval"): "5"})
    @pytest.mark.usefixtures("clean_db")
    def test_import_error_persist_for_invalid_access_control_role(
        self,
        mock_full_path,
        monkeypatch,
        dag_maker,
        session,
        time_machine,
        dag_import_error_listener,
        testing_dag_bundle,
    ):
        """
        Test that import errors related to invalid access control role are tracked in the DB until being fixed.
        """
        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert serialized_dags_count == 0
        time_machine.move_to(tz.datetime(2020, 1, 5, 0, 0, 0), tick=False)

        # create a DAG and assign it a non-exist role.
        with dag_maker(
            dag_id="test_nonexist_access_control",
            access_control={
                "non_existing_role": {"can_edit", "can_read", "can_delete"},
            },
        ) as dag:
            pass
        dag.fileloc = "test_nonexist_access_control.py"
        dag.relative_fileloc = "test_nonexist_access_control.py"
        mock_full_path.return_value = "test_nonexist_access_control.py"

        # the DAG processor should raise an import error when processing the DAG above.
        import_errors = {}
        # run the DAG parsing.
        update_dag_parsing_results_in_db("testing", None, [dag], import_errors, None, set(), session)
        # expect to get an error with "role does not exist" message.
        err = import_errors.get(("testing", dag.relative_fileloc))
        assert "AirflowException" in err
        assert "role does not exist" in err
        dag_model: DagModel = session.get(DagModel, (dag.dag_id,))
        # the DAG should contain an import error.
        assert dag_model.has_import_errors is True

        prev_import_errors = session.scalars(select(ParseImportError)).all()
        # the import error message should match.
        assert len(prev_import_errors) == 1
        prev_import_error = prev_import_errors[0]
        assert prev_import_error.filename == dag.relative_fileloc
        assert "AirflowException" in prev_import_error.stacktrace
        assert "role does not exist" in prev_import_error.stacktrace

        # this is a new import error.
        assert len(dag_import_error_listener.new) == 1
        assert len(dag_import_error_listener.existing) == 0
        assert (
            dag_import_error_listener.new["test_nonexist_access_control.py"] == prev_import_error.stacktrace
        )

        # the DAG is serialized into the DB.
        serialized_dags_count = session.scalar(select(func.count(SerializedDagModel.dag_id)))
        assert serialized_dags_count == 1

        # run the update again. Even though the DAG is not updated, the processor should raise import error since the access control is not fixed.
        time_machine.move_to(tz.datetime(2020, 1, 5, 0, 0, 5), tick=False)
        update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session)

        dag_model: DagModel = session.get(DagModel, (dag.dag_id,))
        # the DAG should contain an import error.
        assert dag_model.has_import_errors is True

        import_errors = session.scalars(select(ParseImportError)).all()
        # the import error should still in the DB.
        assert len(import_errors) == 1
        import_error = import_errors[0]
        assert import_error.filename == dag.relative_fileloc
        assert "AirflowException" in import_error.stacktrace
        assert "role does not exist" in import_error.stacktrace

        # the new import error should be the same as the previous one
        assert len(import_errors) == len(prev_import_errors)
        assert import_error.filename == prev_import_error.filename
        assert import_error.filename == dag.relative_fileloc
        assert import_error.stacktrace == prev_import_error.stacktrace

        # there is a new error and an existing error.
        assert len(dag_import_error_listener.new) == 1
        assert len(dag_import_error_listener.existing) == 1
        assert (
            dag_import_error_listener.new["test_nonexist_access_control.py"] == prev_import_error.stacktrace
        )

        # run the update again, but the incorrect access control configuration is removed.
        time_machine.move_to(tz.datetime(2020, 1, 5, 0, 0, 10), tick=False)
        dag.access_control = None
        update_dag_parsing_results_in_db("testing", None, [dag], dict(), None, set(), session)

        dag_model: DagModel = session.get(DagModel, (dag.dag_id,))
        # the import error should be cleared.
        assert dag_model.has_import_errors is False

        import_errors = session.scalars(select(ParseImportError)).all()
        # the import error should be cleared.
        assert len(import_errors) == 0

        # no import error should be introduced.
        assert len(dag_import_error_listener.new) == 1
        assert len(dag_import_error_listener.existing) == 1

    @patch.object(ParseImportError, "full_file_path")
    @pytest.mark.usefixtures("clean_db")
    def test_new_import_error_replaces_old(
        self, mock_full_file_path, session, dag_import_error_listener, testing_dag_bundle
    ):
        """
        Test that existing import error is updated and new record not created
        for a dag with the same filename
        """
        bundle_name = "testing"
        filename = "abc.py"
        mock_full_file_path.return_value = filename
        prev_error = ParseImportError(
            filename=filename,
            bundle_name=bundle_name,
            timestamp=tz.utcnow(),
            stacktrace="Some error",
        )
        session.add(prev_error)
        session.flush()
        prev_error_id = prev_error.id

        update_dag_parsing_results_in_db(
            bundle_name=bundle_name,
            bundle_version=None,
            dags=[],
            import_errors={("testing", "abc.py"): "New error"},
            parse_duration=None,
            warnings=set(),
            session=session,
            files_parsed={("testing", "abc.py")},
        )

        import_error = session.scalar(
            select(ParseImportError).where(
                ParseImportError.filename == filename, ParseImportError.bundle_name == bundle_name
            )
        )

        # assert that the ID of the import error did not change
        assert import_error.id == prev_error_id
        assert import_error.stacktrace == "New error"

        # Ensure the listener was notified
        assert len(dag_import_error_listener.new) == 0
        assert len(dag_import_error_listener.existing) == 1
        assert dag_import_error_listener.existing["abc.py"] == prev_error.stacktrace

    @pytest.mark.usefixtures("clean_db")
    def test_remove_error_clears_import_error(self, testing_dag_bundle, session):
        # Pre-condition: there is an import error for the dag file
        bundle_name = "testing"
        filename = "abc.py"
        prev_error = ParseImportError(
            filename=filename,
            bundle_name=bundle_name,
            timestamp=tz.utcnow(),
            stacktrace="Some error",
        )
        session.add(prev_error)

        # And one for another file we haven't been given results for -- this shouldn't be deleted
        session.add(
            ParseImportError(
                filename="def.py",
                bundle_name=bundle_name,
                timestamp=tz.utcnow(),
                stacktrace="Some error",
            )
        )
        session.flush()

        # Sanity check of pre-condition
        import_errors = set(session.execute(select(ParseImportError.filename, ParseImportError.bundle_name)))
        assert import_errors == {("abc.py", bundle_name), ("def.py", bundle_name)}

        dag = DAG(dag_id="test")
        dag.fileloc = filename
        dag.relative_fileloc = filename

        import_errors = {}
        update_dag_parsing_results_in_db(
            bundle_name,
            bundle_version=None,
            dags=[LazyDeserializedDAG.from_dag(dag)],
            import_errors=dict.fromkeys(import_errors),
            parse_duration=None,
            warnings=set(),
            session=session,
            files_parsed={(bundle_name, "abc.py")},
        )
        dag_model: DagModel = session.get(DagModel, (dag.dag_id,))
        assert dag_model.has_import_errors is False

        import_errors = set(session.execute(select(ParseImportError.filename, ParseImportError.bundle_name)))

        assert import_errors == {("def.py", bundle_name)}

    @pytest.mark.usefixtures("clean_db")
    def test_remove_error_updates_loaded_dag_model(self, testing_dag_bundle, session):
        bundle_name = "testing"
        filename = "abc.py"
        session.add(
            ParseImportError(
                filename=filename,
                bundle_name=bundle_name,
                timestamp=tz.utcnow(),
                stacktrace="Some error",
            )
        )
        session.add(
            ParseImportError(
                filename="def.py",
                bundle_name=bundle_name,
                timestamp=tz.utcnow(),
                stacktrace="Some error",
            )
        )
        session.flush()

        dag = DAG(dag_id="test")
        dag.fileloc = filename
        dag.relative_fileloc = filename
        lazy_deserialized_dags = [LazyDeserializedDAG.from_dag(dag)]

        import_errors = {(bundle_name, filename): "Some error"}
        update_dag_parsing_results_in_db(
            bundle_name,
            bundle_version=None,
            dags=lazy_deserialized_dags,
            import_errors=import_errors,
            parse_duration=None,
            warnings=set(),
            session=session,
            files_parsed={(bundle_name, "abc.py")},
        )
        dag_model = session.get(DagModel, (dag.dag_id,))
        assert dag_model.has_import_errors is True

        import_errors = {}
        update_dag_parsing_results_in_db(
            bundle_name,
            bundle_version=None,
            dags=lazy_deserialized_dags,
            import_errors=import_errors,
            parse_duration=None,
            warnings=set(),
            session=session,
        )
        assert dag_model.has_import_errors is False

    @pytest.mark.usefixtures("clean_db")
    def test_clear_import_error_for_file_without_dags(self, testing_dag_bundle, session):
        """
        Test that import errors are cleared for files that were parsed but no longer contain DAGs.
        """
        bundle_name = "testing"
        filename = "no_dags.py"

        prev_error = ParseImportError(
            filename=filename,
            bundle_name=bundle_name,
            timestamp=tz.utcnow(),
            stacktrace="Previous import error",
        )
        session.add(prev_error)

        # And import error for another file we haven't parsed (this shouldn't be deleted)
        other_file_error = ParseImportError(
            filename="other.py",
            bundle_name=bundle_name,
            timestamp=tz.utcnow(),
            stacktrace="Some error",
        )
        session.add(other_file_error)
        session.flush()

        import_errors = set(session.execute(select(ParseImportError.filename, ParseImportError.bundle_name)))
        assert import_errors == {("no_dags.py", bundle_name), ("other.py", bundle_name)}

        # Simulate parsing the file: it was parsed successfully (no import errors),
        # but it no longer contains any DAGs. By passing files_parsed, we ensure
        # the import error is cleared even though there are no DAGs.
        files_parsed = {(bundle_name, filename)}
        update_dag_parsing_results_in_db(
            bundle_name=bundle_name,
            bundle_version=None,
            dags=[],  # No DAGs in this file
            import_errors={},  # No import errors
            parse_duration=None,
            warnings=set(),
            session=session,
            files_parsed=files_parsed,
        )

        import_errors = set(session.execute(select(ParseImportError.filename, ParseImportError.bundle_name)))
        assert import_errors == {("other.py", bundle_name)}, "Import error for parsed file should be cleared"

    @pytest.mark.usefixtures("clean_db")
    def test_import_error_update_does_not_touch_other_bundle_with_same_relative_fileloc(self, session):
        relative_fileloc = "example_dag.py"
        session.add_all([DagBundleModel(name="bundle_a"), DagBundleModel(name="bundle_b")])
        session.flush()
        session.add_all(
            [
                DagModel(dag_id="dag_in_bundle_a", relative_fileloc=relative_fileloc, bundle_name="bundle_a"),
                DagModel(dag_id="dag_in_bundle_b", relative_fileloc=relative_fileloc, bundle_name="bundle_b"),
            ]
        )
        session.flush()

        update_dag_parsing_results_in_db(
            bundle_name="bundle_a",
            bundle_version=None,
            dags=[],
            import_errors={("bundle_a", relative_fileloc): "Import failed in bundle_a"},
            parse_duration=None,
            warnings=set(),
            session=session,
            files_parsed={("bundle_a", relative_fileloc)},
        )
        session.flush()

        dag_in_bundle_a = session.get(DagModel, "dag_in_bundle_a")
        dag_in_bundle_b = session.get(DagModel, "dag_in_bundle_b")

        assert dag_in_bundle_a is not None
        assert dag_in_bundle_b is not None
        assert dag_in_bundle_a.bundle_name == "bundle_a"
        assert dag_in_bundle_b.bundle_name == "bundle_b"
        assert dag_in_bundle_a.has_import_errors is True
        assert dag_in_bundle_b.has_import_errors is False

    @pytest.mark.need_serialized_dag(False)
    @pytest.mark.parametrize(
        ("attrs", "expected"),
        [
            pytest.param(
                {
                    "_tasks_": [
                        EmptyOperator(task_id="task", owner="owner1"),
                        EmptyOperator(task_id="task2", owner="owner2"),
                        EmptyOperator(task_id="task3"),
                        EmptyOperator(task_id="task4", owner="owner2"),
                    ]
                },
                {"owners": ["owner1", "owner2"]},
                id="tasks-multiple-owners",
            ),
            pytest.param(
                {"is_paused_upon_creation": True},
                {"is_paused": True},
                id="is_paused_upon_creation",
            ),
            pytest.param(
                {},
                {"owners": ["airflow"]},
                id="default-owner",
            ),
            pytest.param(
                {},
                {"fail_fast": False},
                id="default-fail-fast",
            ),
            pytest.param(
                {"fail_fast": True},
                {"fail_fast": True},
                id="fail-fast-true",
            ),
            pytest.param(
                {
                    "_tasks_": [
                        EmptyOperator(task_id="task", owner="owner1"),
                        EmptyOperator(task_id="task2", owner="owner2"),
                        EmptyOperator(task_id="task3"),
                        EmptyOperator(task_id="task4", owner="owner2"),
                    ],
                    "schedule": "0 0 * * *",
                    "catchup": False,
                },
                {
                    "owners": ["owner1", "owner2"],
                    "next_dagrun": tz.datetime(2020, 1, 5, 0, 0, 0),
                    "next_dagrun_data_interval_start": tz.datetime(2020, 1, 5, 0, 0, 0),
                    "next_dagrun_data_interval_end": tz.datetime(2020, 1, 6, 0, 0, 0),
                    "next_dagrun_create_after": tz.datetime(2020, 1, 6, 0, 0, 0),
                },
                id="with-scheduled-dagruns",
            ),
        ],
    )
    @pytest.mark.usefixtures("clean_db")
    def test_dagmodel_properties(self, attrs, expected, session, time_machine, testing_dag_bundle, dag_maker):
        """Test that properties on the dag model are correctly set when dealing with a LazySerializedDag"""
        dt = tz.datetime(2020, 1, 6, 0, 0, 0)
        time_machine.move_to(dt, tick=False)

        tasks = attrs.pop("_tasks_", None)
        with dag_maker("dag", **attrs) as dag:
            ...
        if tasks:
            dag.add_tasks(tasks)

        if attrs.pop("schedule", None):
            dr_kwargs = {
                "dag_id": "dag",
                "run_type": "scheduled",
                "data_interval": (dt, dt + timedelta(minutes=5)),
            }
            dr1 = DagRun(logical_date=dt, run_id="test_run_id_1", **dr_kwargs, start_date=dt)
            session.add(dr1)
        update_dag_parsing_results_in_db(
            bundle_name="testing",
            bundle_version=None,
            dags=[LazyDeserializedDAG.from_dag(dag)],
            import_errors={},
            parse_duration=None,
            warnings=set(),
            session=session,
        )

        orm_dag = session.get(DagModel, ("dag",))

        for attrname, expected_value in expected.items():
            if attrname == "owners":
                assert sorted(orm_dag.owners.split(", ")) == expected_value
            else:
                assert getattr(orm_dag, attrname) == expected_value

        assert orm_dag.last_parsed_time == dt

    def test_existing_dag_is_paused_upon_creation(self, testing_dag_bundle, session, dag_maker):
        with dag_maker("dag_paused", schedule=None) as dag:
            ...
        update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
        orm_dag = session.get(DagModel, ("dag_paused",))
        assert orm_dag.is_paused is False

        with dag_maker("dag_paused", schedule=None, is_paused_upon_creation=True) as dag:
            ...
        update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
        # Since the dag existed before, it should not follow the pause flag upon creation
        orm_dag = session.get(DagModel, ("dag_paused",))
        assert orm_dag.is_paused is False

    def test_bundle_name_and_version_are_stored(self, testing_dag_bundle, session, dag_maker):
        with dag_maker("mydag", schedule=None) as dag:
            ...
        update_dag_parsing_results_in_db("testing", "1.0", [dag], {}, 0.1, set(), session)
        orm_dag = session.get(DagModel, "mydag")
        assert orm_dag.bundle_name == "testing"
        assert orm_dag.bundle_version == "1.0"

    def test_max_active_tasks_explicit_value_is_used(self, testing_dag_bundle, session, dag_maker):
        with dag_maker("dag_max_tasks", schedule=None, max_active_tasks=5) as dag:
            ...
        update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
        orm_dag = session.get(DagModel, "dag_max_tasks")
        assert orm_dag.max_active_tasks == 5

    def test_max_active_tasks_defaults_from_conf_when_none(self, testing_dag_bundle, session, dag_maker):
        # Override config so that when DAG.max_active_tasks is None, DagModel gets the configured default
        with conf_vars({("core", "max_active_tasks_per_dag"): "7"}):
            with dag_maker("dag_max_tasks_default", schedule=None) as dag:
                ...
            update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
            orm_dag = session.get(DagModel, "dag_max_tasks_default")
            assert orm_dag.max_active_tasks == 7

    def test_max_active_runs_explicit_value_is_used(self, testing_dag_bundle, session, dag_maker):
        with dag_maker("dag_max_runs", schedule=None, max_active_runs=3) as dag:
            ...
        update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
        orm_dag = session.get(DagModel, "dag_max_runs")
        assert orm_dag.max_active_runs == 3

    @pytest.mark.parametrize(
        ("field", "cfg_key", "schema_default"),
        [
            ("max_active_runs", "max_active_runs_per_dag", 16),
            ("max_active_tasks", "max_active_tasks_per_dag", 16),
            ("max_consecutive_failed_dag_runs", "max_consecutive_failed_dag_runs_per_dag", 0),
        ],
    )
    def test_config_driven_field_equal_to_schema_default_not_overridden_by_conf(
        self, testing_dag_bundle, session, dag_maker, field, cfg_key, schema_default
    ):
        with conf_vars({("core", cfg_key): "1"}):
            with dag_maker(f"dag_{field}_schema_default", schedule=None, **{field: schema_default}) as dag:
                ...
            update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
            orm_dag = session.get(DagModel, f"dag_{field}_schema_default")
            assert getattr(orm_dag, field) == schema_default

    def test_max_active_runs_defaults_from_conf_when_none(self, testing_dag_bundle, session, dag_maker):
        with conf_vars({("core", "max_active_runs_per_dag"): "4"}):
            with dag_maker("dag_max_runs_default", schedule=None) as dag:
                ...
            update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
            orm_dag = session.get(DagModel, "dag_max_runs_default")
            assert orm_dag.max_active_runs == 4

    def test_max_consecutive_failed_dag_runs_explicit_value_is_used(
        self, testing_dag_bundle, session, dag_maker
    ):
        with dag_maker("dag_max_failed_runs", schedule=None, max_consecutive_failed_dag_runs=2) as dag:
            ...
        update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
        orm_dag = session.get(DagModel, "dag_max_failed_runs")
        assert orm_dag.max_consecutive_failed_dag_runs == 2

    def test_max_consecutive_failed_dag_runs_defaults_from_conf_when_none(
        self, testing_dag_bundle, session, dag_maker
    ):
        with conf_vars({("core", "max_consecutive_failed_dag_runs_per_dag"): "6"}):
            with dag_maker("dag_max_failed_runs_default", schedule=None) as dag:
                ...
            update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
            orm_dag = session.get(DagModel, "dag_max_failed_runs_default")
            assert orm_dag.max_consecutive_failed_dag_runs == 6


@pytest.mark.db_test
class TestUpdateDagTags:
    @pytest.fixture(autouse=True)
    def setup_teardown(self, session):
        yield
        session.execute(delete(DagModel).where(DagModel.dag_id == "test_dag"))
        session.commit()

    @pytest.mark.parametrize(
        ("initial_tags", "new_tags", "expected_tags"),
        [
            (["dangerous"], {"DANGEROUS"}, {"DANGEROUS"}),
            (["existing"], {"existing", "new"}, {"existing", "new"}),
            (["tag1", "tag2"], {"tag1"}, {"tag1"}),
            (["keep", "remove", "lowercase"], {"keep", "LOWERCASE", "new"}, {"keep", "LOWERCASE", "new"}),
            (["tag1", "tag2"], set(), set()),
        ],
    )
    def test_update_dag_tags(self, testing_dag_bundle, session, initial_tags, new_tags, expected_tags):
        dag_model = DagModel(dag_id="test_dag", bundle_name="testing")
        dag_model.tags = [DagTag(name=tag, dag_id="test_dag") for tag in initial_tags]
        session.add(dag_model)
        session.commit()

        _update_dag_tags(new_tags, dag_model, session=session)
        session.commit()

        assert {t.name for t in dag_model.tags} == expected_tags


@pytest.mark.db_test
class TestPartitionMapperInfoSync:
    """Verify partition_mapper_info is populated on DagModel during Dag sync."""

    @pytest.fixture(autouse=True)
    def clean_db_around_test(self) -> Generator:
        def reset() -> None:
            clear_db_dags()
            clear_db_assets()
            clear_db_serialized_dags()

        reset()
        yield
        reset()

    def test_partitioned_dag_with_rollup_mapper(self, dag_maker, session):
        """Cover regular Asset, name ref, and uri ref entries in partition_mapper_info."""
        rollup_asset = Asset(uri="s3://bucket/rollup", name="rollup")
        name_ref = Asset.ref(name="ref_by_name")
        uri_ref = Asset.ref(uri="s3://ref")
        rollup_mapper = RollupMapper(upstream_mapper=StartOfDayMapper(), window=DayWindow())

        with dag_maker(
            dag_id="partitioned_with_rollup",
            schedule=PartitionedAssetTimetable(
                assets=AssetAll(rollup_asset, name_ref, uri_ref),
                partition_mapper_config={
                    rollup_asset: rollup_mapper,
                    name_ref: rollup_mapper,
                    uri_ref: rollup_mapper,
                },
            ),
            serialized=True,
        ):
            EmptyOperator(task_id="t")

        dag_model = session.get(DagModel, "partitioned_with_rollup")
        assert dag_model.partition_mapper_info == [
            {"name": "rollup", "uri": "s3://bucket/rollup", "is_rollup": True},
            {"name": "ref_by_name", "is_rollup": True},
            {"uri": "s3://ref", "is_rollup": True},
        ]
        assert dag_model.has_rollup_mappers is True
        assert dag_model.is_rollup_asset(name="rollup", uri="s3://bucket/rollup") is True
        assert dag_model.is_rollup_asset(name="ref_by_name", uri="") is True
        assert dag_model.is_rollup_asset(name="", uri="s3://ref") is True

    def test_partitioned_dag_with_default_rollup_mapper(self, dag_maker, session):
        """
        Using only ``default_partition_mapper=RollupMapper(...)`` (the primary
        documented pattern, see ``example_asset_partition.py``) must still
        produce a non-empty ``partition_mapper_info`` with ``is_rollup=True``,
        so the UI's ``has_rollup_mappers`` / ``is_rollup_asset`` checks return
        the right values without inspecting ``partition_mapper_config``.
        """
        rollup_asset = Asset(uri="s3://bucket/rollup", name="rollup")

        with dag_maker(
            dag_id="partitioned_with_default_rollup",
            schedule=PartitionedAssetTimetable(
                assets=rollup_asset,
                default_partition_mapper=RollupMapper(upstream_mapper=StartOfDayMapper(), window=DayWindow()),
            ),
            serialized=True,
        ):
            EmptyOperator(task_id="t")

        dag_model = session.get(DagModel, "partitioned_with_default_rollup")
        assert dag_model.partition_mapper_info == [
            {"name": "rollup", "uri": "s3://bucket/rollup", "is_rollup": True},
        ]
        assert dag_model.has_rollup_mappers is True
        assert dag_model.is_rollup_asset(name="rollup", uri="s3://bucket/rollup") is True

    def test_non_partitioned_dag_leaves_info_empty(self, dag_maker, session):
        with dag_maker(
            dag_id="non_partitioned_dag",
            schedule=[Asset(uri="s3://bucket/A", name="A")],
            serialized=True,
        ):
            EmptyOperator(task_id="t")

        dag_model = session.get(DagModel, "non_partitioned_dag")
        assert dag_model.partition_mapper_info == []
        assert dag_model.has_rollup_mappers is False
