We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
DryRun is not optional for patch task instance
## Summary

According to the [REST api docs](https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html#operation/patch_mapped_task_instance).
you can patch a task instance state. When you hit the api without sending a
"dry_run" variable, you get a KeyError (This is from a server running version 2.5.3):

```
Traceback (most recent call last):
  File "/home/airflow/.local/lib/python3.10/site-packages/flask/app.py", line 2528, in wsgi_app
    response = self.full_dispatch_request()
  File "/home/airflow/.local/lib/python3.10/site-packages/flask/app.py", line 1825, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/home/airflow/.local/lib/python3.10/site-packages/flask/app.py", line 1823, in full_dispatch_request
    rv = self.dispatch_request()
  File "/home/airflow/.local/lib/python3.10/site-packages/flask/app.py", line 1799, in dispatch_request
    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)
  File "/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/decorator.py", line 68, in wrapper
    response = function(request)
  File "/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/uri_parsing.py", line 149, in wrapper
    response = function(request)
  File "/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/validation.py", line 196, in wrapper
    response = function(request)
  File "/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/validation.py", line 399, in wrapper
    return function(request)
  File "/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/response.py", line 112, in wrapper
    response = function(request)
  File "/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/parameter.py", line 120, in wrapper
    return function(**kwargs)
  File "/home/airflow/.local/lib/python3.10/site-packages/airflow/api_connexion/security.py", line 51, in decorated
    return func(*args, **kwargs)
  File "/home/airflow/.local/lib/python3.10/site-packages/airflow/utils/session.py", line 75, in wrapper
    return func(*args, session=session, **kwargs)
  File "/home/airflow/.local/lib/python3.10/site-packages/airflow/api_connexion/endpoints/task_instance_endpoint.py", line 594, in patch_task_instance
    if not data["dry_run"]:
KeyError: 'dry_run'
```

The API docs state that dry_run is not required and that it is defaulted to false.

This can be reproduced in `main` with the tests by commenting out line 1699 in
[test_task_instance_endpoint.py](https://github.com/apache/airflow/blob/5b0ce3db4d36e2a7f20a78903daf538bbde5e38a/tests/api_connexion/endpoints/test_task_instance_endpoint.py#L1695-L1709)





--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `airflow/api_connexion/schemas/task_instance_schema.py`
Content:
```
1 # Licensed to the Apache Software Foundation (ASF) under one
2 # or more contributor license agreements.  See the NOTICE file
3 # distributed with this work for additional information
4 # regarding copyright ownership.  The ASF licenses this file
5 # to you under the Apache License, Version 2.0 (the
6 # "License"); you may not use this file except in compliance
7 # with the License.  You may obtain a copy of the License at
8 #
9 #   http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing,
12 # software distributed under the License is distributed on an
13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 # KIND, either express or implied.  See the License for the
15 # specific language governing permissions and limitations
16 # under the License.
17 from __future__ import annotations
18 
19 from typing import TYPE_CHECKING, NamedTuple
20 
21 from marshmallow import Schema, ValidationError, fields, validate, validates_schema
22 from marshmallow.utils import get_value
23 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
24 
25 from airflow.api_connexion.parameters import validate_istimezone
26 from airflow.api_connexion.schemas.common_schema import JsonObjectField
27 from airflow.api_connexion.schemas.enum_schemas import TaskInstanceStateField
28 from airflow.api_connexion.schemas.job_schema import JobSchema
29 from airflow.api_connexion.schemas.sla_miss_schema import SlaMissSchema
30 from airflow.api_connexion.schemas.trigger_schema import TriggerSchema
31 from airflow.models import TaskInstance
32 from airflow.utils.helpers import exactly_one
33 from airflow.utils.state import TaskInstanceState
34 
35 if TYPE_CHECKING:
36     from airflow.models import SlaMiss
37 
38 
39 class TaskInstanceSchema(SQLAlchemySchema):
40     """Task instance schema."""
41 
42     class Meta:
43         """Meta."""
44 
45         model = TaskInstance
46 
47     task_id = auto_field()
48     dag_id = auto_field()
49     run_id = auto_field(data_key="dag_run_id")
50     map_index = auto_field()
51     execution_date = auto_field()
52     start_date = auto_field()
53     end_date = auto_field()
54     duration = auto_field()
55     state = TaskInstanceStateField()
56     _try_number = auto_field(data_key="try_number")
57     max_tries = auto_field()
58     hostname = auto_field()
59     unixname = auto_field()
60     pool = auto_field()
61     pool_slots = auto_field()
62     queue = auto_field()
63     priority_weight = auto_field()
64     operator = auto_field()
65     queued_dttm = auto_field(data_key="queued_when")
66     pid = auto_field()
67     executor_config = auto_field()
68     note = auto_field()
69     sla_miss = fields.Nested(SlaMissSchema, dump_default=None)
70     rendered_fields = JsonObjectField(dump_default={})
71     trigger = fields.Nested(TriggerSchema)
72     triggerer_job = fields.Nested(JobSchema)
73 
74     def get_attribute(self, obj, attr, default):
75         if attr == "sla_miss":
76             # Object is a tuple of task_instance and slamiss
77             # and the get_value expects a dict with key, value
78             # corresponding to the attr.
79             slamiss_instance = {"sla_miss": obj[1]}
80             return get_value(slamiss_instance, attr, default)
81         elif attr == "rendered_fields":
82             return get_value(obj[0], "rendered_task_instance_fields.rendered_fields", default)
83         return get_value(obj[0], attr, default)
84 
85 
86 class TaskInstanceCollection(NamedTuple):
87     """List of task instances with metadata."""
88 
89     task_instances: list[tuple[TaskInstance, SlaMiss | None]]
90     total_entries: int
91 
92 
93 class TaskInstanceCollectionSchema(Schema):
94     """Task instance collection schema."""
95 
96     task_instances = fields.List(fields.Nested(TaskInstanceSchema))
97     total_entries = fields.Int()
98 
99 
100 class TaskInstanceBatchFormSchema(Schema):
101     """Schema for the request form passed to Task Instance Batch endpoint."""
102 
103     page_offset = fields.Int(load_default=0, validate=validate.Range(min=0))
104     page_limit = fields.Int(load_default=100, validate=validate.Range(min=1))
105     dag_ids = fields.List(fields.Str(), load_default=None)
106     dag_run_ids = fields.List(fields.Str(), load_default=None)
107     task_ids = fields.List(fields.Str(), load_default=None)
108     execution_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)
109     execution_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)
110     start_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)
111     start_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)
112     end_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)
113     end_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)
114     duration_gte = fields.Int(load_default=None)
115     duration_lte = fields.Int(load_default=None)
116     state = fields.List(fields.Str(allow_none=True), load_default=None)
117     pool = fields.List(fields.Str(), load_default=None)
118     queue = fields.List(fields.Str(), load_default=None)
119 
120 
121 class ClearTaskInstanceFormSchema(Schema):
122     """Schema for handling the request of clearing task instance of a Dag."""
123 
124     dry_run = fields.Boolean(load_default=True)
125     start_date = fields.DateTime(load_default=None, validate=validate_istimezone)
126     end_date = fields.DateTime(load_default=None, validate=validate_istimezone)
127     only_failed = fields.Boolean(load_default=True)
128     only_running = fields.Boolean(load_default=False)
129     include_subdags = fields.Boolean(load_default=False)
130     include_parentdag = fields.Boolean(load_default=False)
131     reset_dag_runs = fields.Boolean(load_default=False)
132     task_ids = fields.List(fields.String(), validate=validate.Length(min=1))
133     dag_run_id = fields.Str(load_default=None)
134     include_upstream = fields.Boolean(load_default=False)
135     include_downstream = fields.Boolean(load_default=False)
136     include_future = fields.Boolean(load_default=False)
137     include_past = fields.Boolean(load_default=False)
138 
139     @validates_schema
140     def validate_form(self, data, **kwargs):
141         """Validate clear task instance form."""
142         if data["only_failed"] and data["only_running"]:
143             raise ValidationError("only_failed and only_running both are set to True")
144         if data["start_date"] and data["end_date"]:
145             if data["start_date"] > data["end_date"]:
146                 raise ValidationError("end_date is sooner than start_date")
147         if data["start_date"] and data["end_date"] and data["dag_run_id"]:
148             raise ValidationError("Exactly one of dag_run_id or (start_date and end_date) must be provided")
149         if data["start_date"] and data["dag_run_id"]:
150             raise ValidationError("Exactly one of dag_run_id or start_date must be provided")
151         if data["end_date"] and data["dag_run_id"]:
152             raise ValidationError("Exactly one of dag_run_id or end_date must be provided")
153 
154 
155 class SetTaskInstanceStateFormSchema(Schema):
156     """Schema for handling the request of setting state of task instance of a DAG."""
157 
158     dry_run = fields.Boolean(dump_default=True)
159     task_id = fields.Str(required=True)
160     execution_date = fields.DateTime(validate=validate_istimezone)
161     dag_run_id = fields.Str()
162     include_upstream = fields.Boolean(required=True)
163     include_downstream = fields.Boolean(required=True)
164     include_future = fields.Boolean(required=True)
165     include_past = fields.Boolean(required=True)
166     new_state = TaskInstanceStateField(
167         required=True,
168         validate=validate.OneOf(
169             [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED]
170         ),
171     )
172 
173     @validates_schema
174     def validate_form(self, data, **kwargs):
175         """Validate set task instance state form."""
176         if not exactly_one(data.get("execution_date"), data.get("dag_run_id")):
177             raise ValidationError("Exactly one of execution_date or dag_run_id must be provided")
178 
179 
180 class SetSingleTaskInstanceStateFormSchema(Schema):
181     """Schema for handling the request of updating state of a single task instance."""
182 
183     dry_run = fields.Boolean(dump_default=True)
184     new_state = TaskInstanceStateField(
185         required=True,
186         validate=validate.OneOf(
187             [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED]
188         ),
189     )
190 
191 
192 class TaskInstanceReferenceSchema(Schema):
193     """Schema for the task instance reference schema."""
194 
195     task_id = fields.Str()
196     run_id = fields.Str(data_key="dag_run_id")
197     dag_id = fields.Str()
198     execution_date = fields.DateTime()
199 
200 
201 class TaskInstanceReferenceCollection(NamedTuple):
202     """List of objects with metadata about taskinstance and dag_run_id."""
203 
204     task_instances: list[tuple[TaskInstance, str]]
205 
206 
207 class TaskInstanceReferenceCollectionSchema(Schema):
208     """Collection schema for task reference."""
209 
210     task_instances = fields.List(fields.Nested(TaskInstanceReferenceSchema))
211 
212 
213 class SetTaskInstanceNoteFormSchema(Schema):
214     """Schema for settings a note for a TaskInstance."""
215 
216     # Note: We can't add map_index to the url as subpaths can't start with dashes.
217     map_index = fields.Int(allow_none=False)
218     note = fields.String(allow_none=True, validate=validate.Length(max=1000))
219 
220 
221 task_instance_schema = TaskInstanceSchema()
222 task_instance_collection_schema = TaskInstanceCollectionSchema()
223 task_instance_batch_form = TaskInstanceBatchFormSchema()
224 clear_task_instance_form = ClearTaskInstanceFormSchema()
225 set_task_instance_state_form = SetTaskInstanceStateFormSchema()
226 set_single_task_instance_state_form = SetSingleTaskInstanceStateFormSchema()
227 task_instance_reference_schema = TaskInstanceReferenceSchema()
228 task_instance_reference_collection_schema = TaskInstanceReferenceCollectionSchema()
229 set_task_instance_note_form_schema = SetTaskInstanceNoteFormSchema()
230
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/airflow/api_connexion/schemas/task_instance_schema.py b/airflow/api_connexion/schemas/task_instance_schema.py
--- a/airflow/api_connexion/schemas/task_instance_schema.py
+++ b/airflow/api_connexion/schemas/task_instance_schema.py
@@ -180,7 +180,7 @@
 class SetSingleTaskInstanceStateFormSchema(Schema):
     """Schema for handling the request of updating state of a single task instance."""
 
-    dry_run = fields.Boolean(dump_default=True)
+    dry_run = fields.Boolean(load_default=True)
     new_state = TaskInstanceStateField(
         required=True,
         validate=validate.OneOf(

{"golden_diff": "diff --git a/airflow/api_connexion/schemas/task_instance_schema.py b/airflow/api_connexion/schemas/task_instance_schema.py\n--- a/airflow/api_connexion/schemas/task_instance_schema.py\n+++ b/airflow/api_connexion/schemas/task_instance_schema.py\n@@ -180,7 +180,7 @@\n class SetSingleTaskInstanceStateFormSchema(Schema):\n     \"\"\"Schema for handling the request of updating state of a single task instance.\"\"\"\n \n-    dry_run = fields.Boolean(dump_default=True)\n+    dry_run = fields.Boolean(load_default=True)\n     new_state = TaskInstanceStateField(\n         required=True,\n         validate=validate.OneOf(\n", "issue": "DryRun is not optional for patch task instance\n## Summary\n\nAccording to the [REST api docs](https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html#operation/patch_mapped_task_instance).\nyou can patch a task instance state. When you hit the api without sending a\n\"dry_run\" variable, you get a KeyError (This is from a server running version 2.5.3):\n\n```\nTraceback (most recent call last):\n  File \"/home/airflow/.local/lib/python3.10/site-packages/flask/app.py\", line 2528, in wsgi_app\n    response = self.full_dispatch_request()\n  File \"/home/airflow/.local/lib/python3.10/site-packages/flask/app.py\", line 1825, in full_dispatch_request\n    rv = self.handle_user_exception(e)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/flask/app.py\", line 1823, in full_dispatch_request\n    rv = self.dispatch_request()\n  File \"/home/airflow/.local/lib/python3.10/site-packages/flask/app.py\", line 1799, in dispatch_request\n    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/decorator.py\", line 68, in wrapper\n    response = function(request)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/uri_parsing.py\", line 149, in wrapper\n    response = function(request)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/validation.py\", line 196, in wrapper\n    response = function(request)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/validation.py\", line 399, in wrapper\n    return function(request)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/response.py\", line 112, in wrapper\n    response = function(request)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/connexion/decorators/parameter.py\", line 120, in wrapper\n    return function(**kwargs)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/airflow/api_connexion/security.py\", line 51, in decorated\n    return func(*args, **kwargs)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/airflow/utils/session.py\", line 75, in wrapper\n    return func(*args, session=session, **kwargs)\n  File \"/home/airflow/.local/lib/python3.10/site-packages/airflow/api_connexion/endpoints/task_instance_endpoint.py\", line 594, in patch_task_instance\n    if not data[\"dry_run\"]:\nKeyError: 'dry_run'\n```\n\nThe API docs state that dry_run is not required and that it is defaulted to false.\n\nThis can be reproduced in `main` with the tests by commenting out line 1699 in\n[test_task_instance_endpoint.py](https://github.com/apache/airflow/blob/5b0ce3db4d36e2a7f20a78903daf538bbde5e38a/tests/api_connexion/endpoints/test_task_instance_endpoint.py#L1695-L1709)\n\n\n\n\n", "before_files": [{"content": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements.  See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership.  The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License.  You may obtain a copy of the License at\n#\n#   http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied.  See the License for the\n# specific language governing permissions and limitations\n# under the License.\nfrom __future__ import annotations\n\nfrom typing import TYPE_CHECKING, NamedTuple\n\nfrom marshmallow import Schema, ValidationError, fields, validate, validates_schema\nfrom marshmallow.utils import get_value\nfrom marshmallow_sqlalchemy import SQLAlchemySchema, auto_field\n\nfrom airflow.api_connexion.parameters import validate_istimezone\nfrom airflow.api_connexion.schemas.common_schema import JsonObjectField\nfrom airflow.api_connexion.schemas.enum_schemas import TaskInstanceStateField\nfrom airflow.api_connexion.schemas.job_schema import JobSchema\nfrom airflow.api_connexion.schemas.sla_miss_schema import SlaMissSchema\nfrom airflow.api_connexion.schemas.trigger_schema import TriggerSchema\nfrom airflow.models import TaskInstance\nfrom airflow.utils.helpers import exactly_one\nfrom airflow.utils.state import TaskInstanceState\n\nif TYPE_CHECKING:\n    from airflow.models import SlaMiss\n\n\nclass TaskInstanceSchema(SQLAlchemySchema):\n    \"\"\"Task instance schema.\"\"\"\n\n    class Meta:\n        \"\"\"Meta.\"\"\"\n\n        model = TaskInstance\n\n    task_id = auto_field()\n    dag_id = auto_field()\n    run_id = auto_field(data_key=\"dag_run_id\")\n    map_index = auto_field()\n    execution_date = auto_field()\n    start_date = auto_field()\n    end_date = auto_field()\n    duration = auto_field()\n    state = TaskInstanceStateField()\n    _try_number = auto_field(data_key=\"try_number\")\n    max_tries = auto_field()\n    hostname = auto_field()\n    unixname = auto_field()\n    pool = auto_field()\n    pool_slots = auto_field()\n    queue = auto_field()\n    priority_weight = auto_field()\n    operator = auto_field()\n    queued_dttm = auto_field(data_key=\"queued_when\")\n    pid = auto_field()\n    executor_config = auto_field()\n    note = auto_field()\n    sla_miss = fields.Nested(SlaMissSchema, dump_default=None)\n    rendered_fields = JsonObjectField(dump_default={})\n    trigger = fields.Nested(TriggerSchema)\n    triggerer_job = fields.Nested(JobSchema)\n\n    def get_attribute(self, obj, attr, default):\n        if attr == \"sla_miss\":\n            # Object is a tuple of task_instance and slamiss\n            # and the get_value expects a dict with key, value\n            # corresponding to the attr.\n            slamiss_instance = {\"sla_miss\": obj[1]}\n            return get_value(slamiss_instance, attr, default)\n        elif attr == \"rendered_fields\":\n            return get_value(obj[0], \"rendered_task_instance_fields.rendered_fields\", default)\n        return get_value(obj[0], attr, default)\n\n\nclass TaskInstanceCollection(NamedTuple):\n    \"\"\"List of task instances with metadata.\"\"\"\n\n    task_instances: list[tuple[TaskInstance, SlaMiss | None]]\n    total_entries: int\n\n\nclass TaskInstanceCollectionSchema(Schema):\n    \"\"\"Task instance collection schema.\"\"\"\n\n    task_instances = fields.List(fields.Nested(TaskInstanceSchema))\n    total_entries = fields.Int()\n\n\nclass TaskInstanceBatchFormSchema(Schema):\n    \"\"\"Schema for the request form passed to Task Instance Batch endpoint.\"\"\"\n\n    page_offset = fields.Int(load_default=0, validate=validate.Range(min=0))\n    page_limit = fields.Int(load_default=100, validate=validate.Range(min=1))\n    dag_ids = fields.List(fields.Str(), load_default=None)\n    dag_run_ids = fields.List(fields.Str(), load_default=None)\n    task_ids = fields.List(fields.Str(), load_default=None)\n    execution_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    execution_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    start_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    start_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    end_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    end_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    duration_gte = fields.Int(load_default=None)\n    duration_lte = fields.Int(load_default=None)\n    state = fields.List(fields.Str(allow_none=True), load_default=None)\n    pool = fields.List(fields.Str(), load_default=None)\n    queue = fields.List(fields.Str(), load_default=None)\n\n\nclass ClearTaskInstanceFormSchema(Schema):\n    \"\"\"Schema for handling the request of clearing task instance of a Dag.\"\"\"\n\n    dry_run = fields.Boolean(load_default=True)\n    start_date = fields.DateTime(load_default=None, validate=validate_istimezone)\n    end_date = fields.DateTime(load_default=None, validate=validate_istimezone)\n    only_failed = fields.Boolean(load_default=True)\n    only_running = fields.Boolean(load_default=False)\n    include_subdags = fields.Boolean(load_default=False)\n    include_parentdag = fields.Boolean(load_default=False)\n    reset_dag_runs = fields.Boolean(load_default=False)\n    task_ids = fields.List(fields.String(), validate=validate.Length(min=1))\n    dag_run_id = fields.Str(load_default=None)\n    include_upstream = fields.Boolean(load_default=False)\n    include_downstream = fields.Boolean(load_default=False)\n    include_future = fields.Boolean(load_default=False)\n    include_past = fields.Boolean(load_default=False)\n\n    @validates_schema\n    def validate_form(self, data, **kwargs):\n        \"\"\"Validate clear task instance form.\"\"\"\n        if data[\"only_failed\"] and data[\"only_running\"]:\n            raise ValidationError(\"only_failed and only_running both are set to True\")\n        if data[\"start_date\"] and data[\"end_date\"]:\n            if data[\"start_date\"] > data[\"end_date\"]:\n                raise ValidationError(\"end_date is sooner than start_date\")\n        if data[\"start_date\"] and data[\"end_date\"] and data[\"dag_run_id\"]:\n            raise ValidationError(\"Exactly one of dag_run_id or (start_date and end_date) must be provided\")\n        if data[\"start_date\"] and data[\"dag_run_id\"]:\n            raise ValidationError(\"Exactly one of dag_run_id or start_date must be provided\")\n        if data[\"end_date\"] and data[\"dag_run_id\"]:\n            raise ValidationError(\"Exactly one of dag_run_id or end_date must be provided\")\n\n\nclass SetTaskInstanceStateFormSchema(Schema):\n    \"\"\"Schema for handling the request of setting state of task instance of a DAG.\"\"\"\n\n    dry_run = fields.Boolean(dump_default=True)\n    task_id = fields.Str(required=True)\n    execution_date = fields.DateTime(validate=validate_istimezone)\n    dag_run_id = fields.Str()\n    include_upstream = fields.Boolean(required=True)\n    include_downstream = fields.Boolean(required=True)\n    include_future = fields.Boolean(required=True)\n    include_past = fields.Boolean(required=True)\n    new_state = TaskInstanceStateField(\n        required=True,\n        validate=validate.OneOf(\n            [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED]\n        ),\n    )\n\n    @validates_schema\n    def validate_form(self, data, **kwargs):\n        \"\"\"Validate set task instance state form.\"\"\"\n        if not exactly_one(data.get(\"execution_date\"), data.get(\"dag_run_id\")):\n            raise ValidationError(\"Exactly one of execution_date or dag_run_id must be provided\")\n\n\nclass SetSingleTaskInstanceStateFormSchema(Schema):\n    \"\"\"Schema for handling the request of updating state of a single task instance.\"\"\"\n\n    dry_run = fields.Boolean(dump_default=True)\n    new_state = TaskInstanceStateField(\n        required=True,\n        validate=validate.OneOf(\n            [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED]\n        ),\n    )\n\n\nclass TaskInstanceReferenceSchema(Schema):\n    \"\"\"Schema for the task instance reference schema.\"\"\"\n\n    task_id = fields.Str()\n    run_id = fields.Str(data_key=\"dag_run_id\")\n    dag_id = fields.Str()\n    execution_date = fields.DateTime()\n\n\nclass TaskInstanceReferenceCollection(NamedTuple):\n    \"\"\"List of objects with metadata about taskinstance and dag_run_id.\"\"\"\n\n    task_instances: list[tuple[TaskInstance, str]]\n\n\nclass TaskInstanceReferenceCollectionSchema(Schema):\n    \"\"\"Collection schema for task reference.\"\"\"\n\n    task_instances = fields.List(fields.Nested(TaskInstanceReferenceSchema))\n\n\nclass SetTaskInstanceNoteFormSchema(Schema):\n    \"\"\"Schema for settings a note for a TaskInstance.\"\"\"\n\n    # Note: We can't add map_index to the url as subpaths can't start with dashes.\n    map_index = fields.Int(allow_none=False)\n    note = fields.String(allow_none=True, validate=validate.Length(max=1000))\n\n\ntask_instance_schema = TaskInstanceSchema()\ntask_instance_collection_schema = TaskInstanceCollectionSchema()\ntask_instance_batch_form = TaskInstanceBatchFormSchema()\nclear_task_instance_form = ClearTaskInstanceFormSchema()\nset_task_instance_state_form = SetTaskInstanceStateFormSchema()\nset_single_task_instance_state_form = SetSingleTaskInstanceStateFormSchema()\ntask_instance_reference_schema = TaskInstanceReferenceSchema()\ntask_instance_reference_collection_schema = TaskInstanceReferenceCollectionSchema()\nset_task_instance_note_form_schema = SetTaskInstanceNoteFormSchema()\n", "path": "airflow/api_connexion/schemas/task_instance_schema.py"}], "after_files": [{"content": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements.  See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership.  The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License.  You may obtain a copy of the License at\n#\n#   http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied.  See the License for the\n# specific language governing permissions and limitations\n# under the License.\nfrom __future__ import annotations\n\nfrom typing import TYPE_CHECKING, NamedTuple\n\nfrom marshmallow import Schema, ValidationError, fields, validate, validates_schema\nfrom marshmallow.utils import get_value\nfrom marshmallow_sqlalchemy import SQLAlchemySchema, auto_field\n\nfrom airflow.api_connexion.parameters import validate_istimezone\nfrom airflow.api_connexion.schemas.common_schema import JsonObjectField\nfrom airflow.api_connexion.schemas.enum_schemas import TaskInstanceStateField\nfrom airflow.api_connexion.schemas.job_schema import JobSchema\nfrom airflow.api_connexion.schemas.sla_miss_schema import SlaMissSchema\nfrom airflow.api_connexion.schemas.trigger_schema import TriggerSchema\nfrom airflow.models import TaskInstance\nfrom airflow.utils.helpers import exactly_one\nfrom airflow.utils.state import TaskInstanceState\n\nif TYPE_CHECKING:\n    from airflow.models import SlaMiss\n\n\nclass TaskInstanceSchema(SQLAlchemySchema):\n    \"\"\"Task instance schema.\"\"\"\n\n    class Meta:\n        \"\"\"Meta.\"\"\"\n\n        model = TaskInstance\n\n    task_id = auto_field()\n    dag_id = auto_field()\n    run_id = auto_field(data_key=\"dag_run_id\")\n    map_index = auto_field()\n    execution_date = auto_field()\n    start_date = auto_field()\n    end_date = auto_field()\n    duration = auto_field()\n    state = TaskInstanceStateField()\n    _try_number = auto_field(data_key=\"try_number\")\n    max_tries = auto_field()\n    hostname = auto_field()\n    unixname = auto_field()\n    pool = auto_field()\n    pool_slots = auto_field()\n    queue = auto_field()\n    priority_weight = auto_field()\n    operator = auto_field()\n    queued_dttm = auto_field(data_key=\"queued_when\")\n    pid = auto_field()\n    executor_config = auto_field()\n    note = auto_field()\n    sla_miss = fields.Nested(SlaMissSchema, dump_default=None)\n    rendered_fields = JsonObjectField(dump_default={})\n    trigger = fields.Nested(TriggerSchema)\n    triggerer_job = fields.Nested(JobSchema)\n\n    def get_attribute(self, obj, attr, default):\n        if attr == \"sla_miss\":\n            # Object is a tuple of task_instance and slamiss\n            # and the get_value expects a dict with key, value\n            # corresponding to the attr.\n            slamiss_instance = {\"sla_miss\": obj[1]}\n            return get_value(slamiss_instance, attr, default)\n        elif attr == \"rendered_fields\":\n            return get_value(obj[0], \"rendered_task_instance_fields.rendered_fields\", default)\n        return get_value(obj[0], attr, default)\n\n\nclass TaskInstanceCollection(NamedTuple):\n    \"\"\"List of task instances with metadata.\"\"\"\n\n    task_instances: list[tuple[TaskInstance, SlaMiss | None]]\n    total_entries: int\n\n\nclass TaskInstanceCollectionSchema(Schema):\n    \"\"\"Task instance collection schema.\"\"\"\n\n    task_instances = fields.List(fields.Nested(TaskInstanceSchema))\n    total_entries = fields.Int()\n\n\nclass TaskInstanceBatchFormSchema(Schema):\n    \"\"\"Schema for the request form passed to Task Instance Batch endpoint.\"\"\"\n\n    page_offset = fields.Int(load_default=0, validate=validate.Range(min=0))\n    page_limit = fields.Int(load_default=100, validate=validate.Range(min=1))\n    dag_ids = fields.List(fields.Str(), load_default=None)\n    dag_run_ids = fields.List(fields.Str(), load_default=None)\n    task_ids = fields.List(fields.Str(), load_default=None)\n    execution_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    execution_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    start_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    start_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    end_date_gte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    end_date_lte = fields.DateTime(load_default=None, validate=validate_istimezone)\n    duration_gte = fields.Int(load_default=None)\n    duration_lte = fields.Int(load_default=None)\n    state = fields.List(fields.Str(allow_none=True), load_default=None)\n    pool = fields.List(fields.Str(), load_default=None)\n    queue = fields.List(fields.Str(), load_default=None)\n\n\nclass ClearTaskInstanceFormSchema(Schema):\n    \"\"\"Schema for handling the request of clearing task instance of a Dag.\"\"\"\n\n    dry_run = fields.Boolean(load_default=True)\n    start_date = fields.DateTime(load_default=None, validate=validate_istimezone)\n    end_date = fields.DateTime(load_default=None, validate=validate_istimezone)\n    only_failed = fields.Boolean(load_default=True)\n    only_running = fields.Boolean(load_default=False)\n    include_subdags = fields.Boolean(load_default=False)\n    include_parentdag = fields.Boolean(load_default=False)\n    reset_dag_runs = fields.Boolean(load_default=False)\n    task_ids = fields.List(fields.String(), validate=validate.Length(min=1))\n    dag_run_id = fields.Str(load_default=None)\n    include_upstream = fields.Boolean(load_default=False)\n    include_downstream = fields.Boolean(load_default=False)\n    include_future = fields.Boolean(load_default=False)\n    include_past = fields.Boolean(load_default=False)\n\n    @validates_schema\n    def validate_form(self, data, **kwargs):\n        \"\"\"Validate clear task instance form.\"\"\"\n        if data[\"only_failed\"] and data[\"only_running\"]:\n            raise ValidationError(\"only_failed and only_running both are set to True\")\n        if data[\"start_date\"] and data[\"end_date\"]:\n            if data[\"start_date\"] > data[\"end_date\"]:\n                raise ValidationError(\"end_date is sooner than start_date\")\n        if data[\"start_date\"] and data[\"end_date\"] and data[\"dag_run_id\"]:\n            raise ValidationError(\"Exactly one of dag_run_id or (start_date and end_date) must be provided\")\n        if data[\"start_date\"] and data[\"dag_run_id\"]:\n            raise ValidationError(\"Exactly one of dag_run_id or start_date must be provided\")\n        if data[\"end_date\"] and data[\"dag_run_id\"]:\n            raise ValidationError(\"Exactly one of dag_run_id or end_date must be provided\")\n\n\nclass SetTaskInstanceStateFormSchema(Schema):\n    \"\"\"Schema for handling the request of setting state of task instance of a DAG.\"\"\"\n\n    dry_run = fields.Boolean(dump_default=True)\n    task_id = fields.Str(required=True)\n    execution_date = fields.DateTime(validate=validate_istimezone)\n    dag_run_id = fields.Str()\n    include_upstream = fields.Boolean(required=True)\n    include_downstream = fields.Boolean(required=True)\n    include_future = fields.Boolean(required=True)\n    include_past = fields.Boolean(required=True)\n    new_state = TaskInstanceStateField(\n        required=True,\n        validate=validate.OneOf(\n            [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED]\n        ),\n    )\n\n    @validates_schema\n    def validate_form(self, data, **kwargs):\n        \"\"\"Validate set task instance state form.\"\"\"\n        if not exactly_one(data.get(\"execution_date\"), data.get(\"dag_run_id\")):\n            raise ValidationError(\"Exactly one of execution_date or dag_run_id must be provided\")\n\n\nclass SetSingleTaskInstanceStateFormSchema(Schema):\n    \"\"\"Schema for handling the request of updating state of a single task instance.\"\"\"\n\n    dry_run = fields.Boolean(load_default=True)\n    new_state = TaskInstanceStateField(\n        required=True,\n        validate=validate.OneOf(\n            [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED, TaskInstanceState.SKIPPED]\n        ),\n    )\n\n\nclass TaskInstanceReferenceSchema(Schema):\n    \"\"\"Schema for the task instance reference schema.\"\"\"\n\n    task_id = fields.Str()\n    run_id = fields.Str(data_key=\"dag_run_id\")\n    dag_id = fields.Str()\n    execution_date = fields.DateTime()\n\n\nclass TaskInstanceReferenceCollection(NamedTuple):\n    \"\"\"List of objects with metadata about taskinstance and dag_run_id.\"\"\"\n\n    task_instances: list[tuple[TaskInstance, str]]\n\n\nclass TaskInstanceReferenceCollectionSchema(Schema):\n    \"\"\"Collection schema for task reference.\"\"\"\n\n    task_instances = fields.List(fields.Nested(TaskInstanceReferenceSchema))\n\n\nclass SetTaskInstanceNoteFormSchema(Schema):\n    \"\"\"Schema for settings a note for a TaskInstance.\"\"\"\n\n    # Note: We can't add map_index to the url as subpaths can't start with dashes.\n    map_index = fields.Int(allow_none=False)\n    note = fields.String(allow_none=True, validate=validate.Length(max=1000))\n\n\ntask_instance_schema = TaskInstanceSchema()\ntask_instance_collection_schema = TaskInstanceCollectionSchema()\ntask_instance_batch_form = TaskInstanceBatchFormSchema()\nclear_task_instance_form = ClearTaskInstanceFormSchema()\nset_task_instance_state_form = SetTaskInstanceStateFormSchema()\nset_single_task_instance_state_form = SetSingleTaskInstanceStateFormSchema()\ntask_instance_reference_schema = TaskInstanceReferenceSchema()\ntask_instance_reference_collection_schema = TaskInstanceReferenceCollectionSchema()\nset_task_instance_note_form_schema = SetTaskInstanceNoteFormSchema()\n", "path": "airflow/api_connexion/schemas/task_instance_schema.py"}]}