code: def update_dag_run_state(
      *, dag_id: str, dag_run_id: str, session: Session = NEW_SESSION
  ) -> APIResponse:
      """Set a state of a dag run."""
      dag_run: DagRun | None = session.scalar(
          select(DagRun).where(
              DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id
          )
      )
      if dag_run is None:
          error_message = (
              f"Dag Run id {dag_run_id} not found in dag {dag_id}"
          )
          raise NotFound(error_message)
      try:
          post_body = set_dagrun_state_form_schema.load(
              get_json_request_dict()
          )
      except ValidationError as err:
          raise BadRequest(detail=str(err))
  
      state = post_body["state"]
      dag = get_airflow_app().dag_bag.get_dag(dag_id)
      if state == DagRunState.SUCCESS:
          set_dag_run_state_to_success(
              dag=dag, run_id=dag_run.run_id, commit=True
          )
      elif state == DagRunState.QUEUED:
          set_dag_run_state_to_queued(
              dag=dag, run_id=dag_run.run_id, commit=True
          )
      else:
          set_dag_run_state_to_failed(
              dag=dag, run_id=dag_run.run_id, commit=True
          )
      dag_run = session.get(DagRun, dag_run.id)
      return dagrun_schema.dump(dag_run)
  
imports: from __future__ import annotations
  from http import HTTPStatus
  from typing import TYPE_CHECKING, Collection
  import pendulum
  from connexion import NoContent
  from flask import g
  from marshmallow import ValidationError
  from sqlalchemy import delete, or_, select
  from airflow.api.common.mark_tasks import set_dag_run_state_to_failed, set_dag_run_state_to_queued, set_dag_run_state_to_success
  from airflow.api_connexion import security
  from airflow.api_connexion.endpoints.request_dict import get_json_request_dict
  from airflow.api_connexion.exceptions import AlreadyExists, BadRequest, NotFound
  from airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters
  from airflow.api_connexion.schemas.dag_run_schema import DAGRunCollection, DAGRunCollectionSchema, DAGRunSchema, clear_dagrun_form_schema, dagrun_collection_schema, dagrun_schema, dagruns_batch_form_schema, set_dagrun_note_form_schema, set_dagrun_state_form_schema
  from airflow.api_connexion.schemas.dataset_schema import DatasetEventCollection, dataset_event_collection_schema
  from airflow.api_connexion.schemas.task_instance_schema import TaskInstanceReferenceCollection, task_instance_reference_collection_schema
  from airflow.auth.managers.models.resource_details import DagAccessEntity
  from airflow.models import DagModel, DagRun
  from airflow.security import permissions
  from airflow.utils.airflow_flask_app import get_airflow_app
  from airflow.utils.db import get_query_count
  from airflow.utils.log.action_logger import action_event_from_permission
  from airflow.utils.session import NEW_SESSION, provide_session
  from airflow.utils.state import DagRunState
  from airflow.utils.types import DagRunType
  from airflow.www.decorators import action_logging
  from airflow.www.extensions.init_auth_manager import get_auth_manager
label_desc: Use image node_id 8 for calling a global function with example usage: update_dag_run_state() and returns: dagrun_schema