Repository: apache/airflow

Issue:
Triggering a DAG with the same run_id as a scheduled one causes the scheduler to crash
### Apache Airflow version

2.5.0

### What happened

A user with access to manually triggering DAGs can trigger a DAG. provide a run_id that matches the pattern used when creating scheduled runs and cause the scheduler to crash due to database unique key violation:
```
2022-12-12 12:58:00,793] {scheduler_job.py:776} ERROR - Exception when executing SchedulerJob._run_scheduler_loop
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 759, in _execute
    self._run_scheduler_loop()
  File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 885, in _run_scheduler_loop
    num_queued_tis = self._do_scheduling(session)
  File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 956, in _do_scheduling
    self._create_dagruns_for_dags(guard, session)
  File "/usr/local/lib/python3.8/site-packages/airflow/utils/retries.py", line 78, in wrapped_function
    for attempt in run_with_db_retries(max_retries=retries, logger=logger, **retry_kwargs):
  File "/usr/local/lib/python3.8/site-packages/tenacity/__init__.py", line 384, in __iter__
    do = self.iter(retry_state=retry_state)
  File "/usr/local/lib/python3.8/site-packages/tenacity/__init__.py", line 351, in iter
    return fut.result()
  File "/usr/local/lib/python3.8/concurrent/futures/_base.py", line 437, in result
    return self.__get_result()
  File "/usr/local/lib/python3.8/concurrent/futures/_base.py", line 389, in __get_result
    raise self._exception
  File "/usr/local/lib/python3.8/site-packages/airflow/utils/retries.py", line 87, in wrapped_function
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 1018, in _create_dagruns_for_dags
    query, dataset_triggered_dag_info = DagModel.dags_needing_dagruns(session)
  File "/usr/local/lib/python3.8/site-packages/airflow/models/dag.py", line 3341, in dags_needing_dagruns
    for x in session.query(
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/orm/query.py", line 2773, in all
    return self._iter().all()
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/orm/query.py", line 2916, in _iter
    result = self.session.execute(
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/orm/session.py", line 1713, in execute
    conn = self._connection_for_bind(bind)
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/orm/session.py", line 1552, in _connection_for_bind
    return self._transaction._connection_for_bind(
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/orm/session.py", line 721, in _connection_for_bind
    self._assert_active()
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/orm/session.py", line 601, in _assert_active
    raise sa_exc.PendingRollbackError(
sqlalchemy.exc.PendingRollbackError: This Session's transaction has been rolled back due to a previous exception during flush. To begin a new transaction with this Session, first issue Session.rollback(). Original exception was: (psycopg2.errors.UniqueViolation) duplicate key value violates unique constraint "dag_run_dag_id_run_id_key"
DETAIL:  Key (dag_id, run_id)=(example_branch_dop_operator_v3, scheduled__2022-12-12T12:57:00+00:00) already exists.

[SQL: INSERT INTO dag_run (dag_id, queued_at, execution_date, start_date, end_date, state, run_id, creating_job_id, external_trigger, run_type, conf, data_interval_start, data_interval_end, last_scheduling_decision, dag_hash, log_template_id, updated_at) VALUES (%(dag_id)s, %(queued_at)s, %(execution_date)s, %(start_date)s, %(end_date)s, %(state)s, %(run_id)s, %(creating_job_id)s, %(external_trigger)s, %(run_type)s, %(conf)s, %(data_interval_start)s, %(data_interval_end)s, %(last_scheduling_decision)s, %(dag_hash)s, (SELECT max(log_template.id) AS max_1 
FROM log_template), %(updated_at)s) RETURNING dag_run.id]
[parameters: {'dag_id': 'example_branch_dop_operator_v3', 'queued_at': datetime.datetime(2022, 12, 12, 12, 58, 0, 435945, tzinfo=Timezone('UTC')), 'execution_date': DateTime(2022, 12, 12, 12, 57, 0, tzinfo=Timezone('UTC')), 'start_date': None, 'end_date': None, 'state': <DagRunState.QUEUED: 'queued'>, 'run_id': 'scheduled__2022-12-12T12:57:00+00:00', 'creating_job_id': 1, 'external_trigger': False, 'run_type': <DagRunType.SCHEDULED: 'scheduled'>, 'conf': <psycopg2.extensions.Binary object at 0x7f283a82af60>, 'data_interval_start': DateTime(2022, 12, 12, 12, 57, 0, tzinfo=Timezone('UTC')), 'data_interval_end': DateTime(2022, 12, 12, 12, 58, 0, tzinfo=Timezone('UTC')), 'last_scheduling_decision': None, 'dag_hash': '1653a588de69ed25c5b1dcfef928479c', 'updated_at': datetime.datetime(2022, 12, 12, 12, 58, 0, 436871, tzinfo=Timezone('UTC'))}]
(Background on this error at: https://sqlalche.me/e/14/gkpj) (Background on this error at: https://sqlalche.me/e/14/7s2a)
```

Worse yet, the scheduler will keep crashing after a restart with the same exception.

### What you think should happen instead

A user should not be able to crash the scheduler from the UI.  
I see 2 alternatives for solving this:

1. Reject custom run_id that would (or could) collide with a scheduled one, preventing this situation from happening.
2. Handle the database error and assign a different run_id to the scheduled run.

### How to reproduce

1. Find an unpaused DAG.
2. Trigger DAG w/ config, set the run id to something like scheduled__2022-11-21T12:00:00+00:00 (adjust the time to be in the future where there is no run yet).
3. Let the manual DAG run finish.
4. Wait for the scheduler to try to schedule another DAG run with the same run id.
5.  :boom: 
6. Attempt to restart the scheduler.
7. :boom: 

### Operating System

Debian GNU/Linux 11 (bullseye)

### Versions of Apache Airflow Providers

apache-airflow-providers-postgres==5.3.1

### Deployment

Docker-Compose

### Deployment details

I'm using a Postgres docker container as a metadata database that is linked via docker networking to the scheduler and the rest of the components. Scheduler, workers and webserver are all running in separate containers (using CeleryExecutor backed by a Redis container), though I do not think it is relevant in this case. 

### Anything else

_No response_

### Are you willing to submit PR?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Smth to take a look at 2.5.1
Code around this has actually changed quite a bit due to datasets. A reproduction on latest Airflow would be awesome.
Yep. @uranusjr  is right. We are releasing 2.5.0 in coming days - can you please try to reproduce it when we do @saulbein ?
I reproduced the same bug with 2.4.3, will try when 2.5.0 is in PyPI to see if it's reproducible there as well :slightly_smiling_face: 
> I reproduced the same bug with 2.4.3, will try when 2.5.0 is in PyPI to see if it's reproducible there as well 🙂

2.5.0 is released. If you see that there is the same error in 2.5.0 - can you please post stacktrace and any extra infoy you will find @saulbein ?
Verified this still happens on 2.5.0 and updated the stack trace.
Coool . Thanks!
To make sure, is the log attached above from 2.5.0 or 2.4.3?
The log is from 2.5.0
Awesome, that should make debugging easier. Thanks!

Resolving patch:
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -2549,7 +2549,7 @@ def create_dagrun(
         external_trigger: bool | None = False,
         conf: dict | None = None,
         run_type: DagRunType | None = None,
-        session=NEW_SESSION,
+        session: Session = NEW_SESSION,
         dag_hash: str | None = None,
         creating_job_id: int | None = None,
         data_interval: tuple[datetime, datetime] | None = None,
@@ -2586,14 +2586,27 @@ def create_dagrun(
             else:
                 data_interval = self.infer_automated_data_interval(logical_date)
 
+        if run_type is None or isinstance(run_type, DagRunType):
+            pass
+        elif isinstance(run_type, str):  # Compatibility: run_type used to be a str.
+            run_type = DagRunType(run_type)
+        else:
+            raise ValueError(f"`run_type` should be a DagRunType, not {type(run_type)}")
+
         if run_id:  # Infer run_type from run_id if needed.
             if not isinstance(run_id, str):
                 raise ValueError(f"`run_id` should be a str, not {type(run_id)}")
-            if not run_type:
-                run_type = DagRunType.from_run_id(run_id)
+            inferred_run_type = DagRunType.from_run_id(run_id)
+            if run_type is None:
+                # No explicit type given, use the inferred type.
+                run_type = inferred_run_type
+            elif run_type == DagRunType.MANUAL and inferred_run_type != DagRunType.MANUAL:
+                # Prevent a manual run from using an ID that looks like a scheduled run.
+                raise ValueError(
+                    f"A {run_type.value} DAG run cannot use ID {run_id!r} since it "
+                    f"is reserved for {inferred_run_type.value} runs"
+                )
         elif run_type and logical_date is not None:  # Generate run_id from run_type and execution_date.
-            if not isinstance(run_type, DagRunType):
-                raise ValueError(f"`run_type` should be a DagRunType, not {type(run_type)}")
             run_id = self.timetable.generate_run_id(
                 run_type=run_type, logical_date=logical_date, data_interval=data_interval
             )


Test patch:
diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -3147,3 +3147,23 @@ def generate_run_id(self, *, run_type, logical_date, data_interval, **extra) ->
     )
 
     assert dag_run.run_id == "abc"
+
+
+@pytest.mark.parametrize(
+    "run_id_type",
+    [DagRunType.BACKFILL_JOB, DagRunType.SCHEDULED, DagRunType.DATASET_TRIGGERED],
+)
+def test_create_dagrun_disallow_manual_to_use_automated_run_id(run_id_type: DagRunType) -> None:
+    dag = DAG(dag_id="test", start_date=DEFAULT_DATE, schedule="@daily")
+    run_id = run_id_type.generate_run_id(DEFAULT_DATE)
+    with pytest.raises(ValueError) as ctx:
+        dag.create_dagrun(
+            run_type=DagRunType.MANUAL,
+            run_id=run_id,
+            execution_date=DEFAULT_DATE,
+            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
+            state=DagRunState.QUEUED,
+        )
+    assert str(ctx.value) == (
+        f"A manual DAG run cannot use ID {run_id!r} since it is reserved for {run_id_type.value} runs"
+    )