Merge branch 'main' into feat-49796-azure_virtual_machines_operator

Author: nailo2c

PROVIDERS.rst

@@ -74,22 +74,26 @@ def exception_handler(self, request: Request, exc: IntegrityError):
             for tb in traceback.format_tb(exc.__traceback__):
                 stacktrace += tb
 
-            log_message = f"Error with id {exception_id}\n{stacktrace}"
+            log_message = f"Error with id {exception_id}, statement: {exc.statement}\n{stacktrace}"
             log.error(log_message)
             if conf.get("api", "expose_stacktrace") == "True":
                 message = log_message
+                statement = str(exc.statement)
+                orig_error = str(exc.orig)
             else:
                 message = (
                     "Serious error when handling your request. Check logs for more details - "
                     f"you will find it in api server when you look for ID {exception_id}"
                 )
+                statement = "hidden"
+                orig_error = "hidden"
 
             raise HTTPException(
                 status_code=status.HTTP_409_CONFLICT,
                 detail={
                     "reason": "Unique constraint violation",
-                    "statement": str(exc.statement),
-                    "orig_error": str(exc.orig),
+                    "statement": statement,
+                    "orig_error": orig_error,
                     "message": message,
                 },
             )

airflow-core/src/airflow/api_fastapi/common/exceptions.py

@@ -30,7 +30,7 @@
 from cadwyn import VersionedAPIRouter
 from fastapi import Body, HTTPException, Query, Security, status
 from pydantic import JsonValue
-from sqlalchemy import func, or_, tuple_, update
+from sqlalchemy import and_, func, or_, tuple_, update
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import NoResultFound, SQLAlchemyError
 from sqlalchemy.orm import joinedload
@@ -64,6 +64,7 @@
 from airflow.models.asset import AssetActive
 from airflow.models.dag import DagModel
 from airflow.models.dagrun import DagRun as DR
+from airflow.models.log import Log
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
@@ -136,10 +137,14 @@ def ti_run(
             # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the
             # client
             column("next_kwargs", JSON),
+            DR.logical_date,
+            DagModel.owners,
         )
         .select_from(TI)
+        .join(DR, and_(TI.dag_id == DR.dag_id, TI.run_id == DR.run_id))
+        .join(DagModel, TI.dag_id == DagModel.dag_id)
         .where(TI.id == task_instance_id)
-        .with_for_update()
+        .with_for_update(of=TI)
     )
     try:
         ti = session.execute(old).one()
@@ -195,6 +200,19 @@ def ti_run(
         )
     else:
         log.info("Task started", previous_state=previous_state, hostname=ti_run_payload.hostname)
+        session.add(
+            Log(
+                event=TaskInstanceState.RUNNING.value,
+                task_id=ti.task_id,
+                dag_id=ti.dag_id,
+                run_id=ti.run_id,
+                map_index=ti.map_index,
+                try_number=ti.try_number,
+                logical_date=ti.logical_date,
+                owner=ti.owners,
+                extra=json.dumps({"host_name": ti_run_payload.hostname}) if ti_run_payload.hostname else None,
+            )
+        )
     # Ensure there is no end date set.
     query = query.values(
         end_date=None,
@@ -297,16 +315,36 @@ def ti_update_state(
     log.debug("Updating task instance state", new_state=ti_patch_payload.state)
 
     old = (
-        select(TI.state, TI.try_number, TI.max_tries, TI.dag_id)
+        select(
+            TI.state,
+            TI.try_number,
+            TI.max_tries,
+            TI.dag_id,
+            TI.task_id,
+            TI.run_id,
+            TI.map_index,
+            TI.hostname,
+            DR.logical_date,
+            DagModel.owners,
+        )
+        .select_from(TI)
+        .join(DR, and_(TI.dag_id == DR.dag_id, TI.run_id == DR.run_id))
+        .join(DagModel, TI.dag_id == DagModel.dag_id)
         .where(TI.id == task_instance_id)
-        .with_for_update()
+        .with_for_update(of=TI)
     )
     try:
         (
             previous_state,
             try_number,
             max_tries,
             dag_id,
+            task_id,
+            run_id,
+            map_index,
+            hostname,
+            logical_date,
+            owners,
         ) = session.execute(old).one()
         log.debug(
             "Retrieved current task instance state",
@@ -373,6 +411,19 @@ def ti_update_state(
             new_state=updated_state,
             rows_affected=getattr(result, "rowcount", 0),
         )
+        session.add(
+            Log(
+                event=updated_state.value,
+                task_id=task_id,
+                dag_id=dag_id,
+                run_id=run_id,
+                map_index=map_index,
+                try_number=try_number,
+                logical_date=logical_date,
+                owner=owners,
+                extra=json.dumps({"host_name": hostname}) if hostname else None,
+            )
+        )
     except SQLAlchemyError as e:
         log.error("Error updating Task Instance state", error=str(e))
         raise HTTPException(

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -1414,9 +1414,21 @@ traces:
       description: |
         If True, then traces from Airflow internal methods are exported. Defaults to False.
       version_added: 3.1.0
+      version_deprecated: 3.2.0
+      deprecation_reason: |
+        This parameter is no longer used.
       type: string
       example: ~
       default: "False"
+    task_runner_flush_timeout_milliseconds:
+      description: |
+        Timeout in milliseconds to wait for the OpenTelemetry span exporter to flush pending spans
+        when a task runner process exits. If the exporter does not finish within this time, any
+        buffered spans may be dropped.
+      version_added: 3.2.0
+      type: integer
+      example: ~
+      default: "30000"
 secrets:
   description: ~
   options:

airflow-core/src/airflow/config_templates/config.yml

@@ -32,14 +32,11 @@
 from airflow.configuration import conf
 from airflow.executors import workloads
 from airflow.executors.executor_loader import ExecutorLoader
-from airflow.executors.workloads.task import TaskInstanceDTO
 from airflow.models import Log
 from airflow.models.callback import CallbackKey
 from airflow.observability.metrics import stats_utils
-from airflow.observability.trace import Trace
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import TaskInstanceState
-from airflow.utils.thread_safe_dict import ThreadSafeDict
 
 PARALLELISM: int = conf.getint("core", "PARALLELISM")
 
@@ -143,8 +140,6 @@ class BaseExecutor(LoggingMixin):
     :param parallelism: how many jobs should run at one time.
     """
 
-    active_spans = ThreadSafeDict()
-
     supports_ad_hoc_ti_run: bool = False
     supports_callbacks: bool = False
     supports_multi_team: bool = False
@@ -217,10 +212,6 @@ def __repr__(self):
         _repr += ")"
         return _repr
 
-    @classmethod
-    def set_active_spans(cls, active_spans: ThreadSafeDict):
-        cls.active_spans = active_spans
-
     def start(self):  # pragma: no cover
         """Executors may need to get things started."""
 
@@ -340,17 +331,6 @@ def _emit_metrics(self, open_slots, num_running_tasks, num_queued_tasks):
         queued_tasks_metric_name = self._get_metric_name("executor.queued_tasks")
         running_tasks_metric_name = self._get_metric_name("executor.running_tasks")
 
-        span = Trace.get_current_span()
-        if span.is_recording():
-            span.add_event(
-                name="executor",
-                attributes={
-                    open_slots_metric_name: open_slots,
-                    queued_tasks_metric_name: num_queued_tasks,
-                    running_tasks_metric_name: num_running_tasks,
-                },
-            )
-
         self.log.debug("%s running task instances for executor %s", num_running_tasks, name)
         self.log.debug("%s in queue for executor %s", num_queued_tasks, name)
         if open_slots == 0:
@@ -415,30 +395,6 @@ def trigger_tasks(self, open_slots: int) -> None:
             if key in self.attempts:
                 del self.attempts[key]
 
-            if isinstance(workload, workloads.ExecuteTask) and hasattr(workload, "ti"):
-                ti = workload.ti
-
-                # If it's None, then the span for the current id hasn't been started.
-                if self.active_spans is not None and self.active_spans.get("ti:" + str(ti.id)) is None:
-                    if isinstance(ti, TaskInstanceDTO):
-                        parent_context = Trace.extract(ti.parent_context_carrier)
-                    else:
-                        parent_context = Trace.extract(ti.dag_run.context_carrier)
-                    # Start a new span using the context from the parent.
-                    # Attributes will be set once the task has finished so that all
-                    # values will be available (end_time, duration, etc.).
-
-                    span = Trace.start_child_span(
-                        span_name=f"{ti.task_id}",
-                        parent_context=parent_context,
-                        component="task",
-                        start_as_current=False,
-                    )
-                    self.active_spans.set("ti:" + str(ti.id), span)
-                    # Inject the current context into the carrier.
-                    carrier = Trace.inject()
-                    ti.context_carrier = carrier
-
             workload_list.append(workload)
 
         if workload_list:

airflow-core/src/airflow/executors/base_executor.py

@@ -90,7 +90,7 @@ def make(
                 name=dag_run.dag_model.bundle_name,
                 version=dag_run.bundle_version,
             )
-        fname = f"executor_callbacks/{callback.id}"  # TODO: better log file template
+        fname = f"executor_callbacks/{dag_run.dag_id}/{dag_run.run_id}/{callback.id}"
 
         return cls(
             callback=CallbackDTO.model_validate(callback, from_attributes=True),

airflow-core/src/airflow/executors/workloads/callback.py

@@ -86,7 +86,7 @@ def make(
         from airflow.utils.helpers import log_filename_template_renderer
 
         ser_ti = TaskInstanceDTO.model_validate(ti, from_attributes=True)
-        ser_ti.parent_context_carrier = ti.dag_run.context_carrier
+        ser_ti.context_carrier = ti.dag_run.context_carrier
         if not bundle_info:
             bundle_info = BundleInfo(
                 name=ti.dag_model.bundle_name,