Add dataset update endpoint

airflow/api_connexion/endpoints/dataset_endpoint.py

@@ -16,20 +16,30 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import timezone
+
+from flask import request
+from flask_login import current_user
+from marshmallow import ValidationError
 from sqlalchemy import func
 from sqlalchemy.orm import Session, joinedload, subqueryload
 
+from airflow import Dataset
 from airflow.api_connexion import security
-from airflow.api_connexion.exceptions import NotFound
+from airflow.api_connexion.endpoints.request_dict import get_json_request_dict
+from airflow.api_connexion.exceptions import BadRequest, NotFound
 from airflow.api_connexion.parameters import apply_sorting, check_limit, format_parameters
 from airflow.api_connexion.schemas.dataset_schema import (
     DatasetCollection,
     DatasetEventCollection,
+    dataset_change_schema,
     dataset_collection_schema,
     dataset_event_collection_schema,
+    dataset_event_schema,
     dataset_schema,
 )
 from airflow.api_connexion.types import APIResponse
+from airflow.datasets.manager import dataset_manager
 from airflow.models.dataset import DatasetEvent, DatasetModel
 from airflow.security import permissions
 from airflow.utils.session import NEW_SESSION, provide_session
@@ -120,3 +130,39 @@ def get_dataset_events(
     return dataset_event_collection_schema.dump(
         DatasetEventCollection(dataset_events=events, total_entries=total_entries)
     )
+
+
+@security.requires_access([(permissions.ACTION_CAN_CREATE, permissions.RESOURCE_DATASET)])
+@provide_session
+def post_dataset_event(session: Session = NEW_SESSION) -> APIResponse:
+    """Create an external dataset event. This endpoint is useful if you want to update a dataset and
+    trigger downstream DAG runs from external services.
+    """
+    try:
+        json_body = dataset_change_schema.load(get_json_request_dict())
+    except ValidationError as err:
+        raise BadRequest(detail=str(err))
+    uri = json_body["dataset_uri"]
+    external_source = request.remote_addr
+    user_id = getattr(current_user, "id", None)
+    timestamp = json_body.get("timestamp").astimezone(timezone.utc)
+    extra = json_body["extra"]
+    dataset_event = dataset_manager.register_external_dataset_change(
+        dataset=Dataset(uri),
+        external_source=external_source,
+        user_id=user_id,
+        timestamp=timestamp,
+        extra=extra,
+        session=session,
+    )
+
+    if dataset_event:
+        event_json = dataset_event_schema.dump(dataset_event)
+        # removing created_dagruns, since they will be created asynchronously in the scheduler
+        event_json.pop("created_dagruns")
+        return event_json
+
+    raise BadRequest(
+        "Dataset not found",
+        detail=f"The Dataset with uri: `{uri}` was not found",
+    )

airflow/api_connexion/openapi/v1.yaml

@@ -1921,6 +1921,32 @@ paths:
           $ref: '#/components/responses/PermissionDenied'
         '404':
           $ref: '#/components/responses/NotFound'
+    post:
+      summary: Post dataset event
+      description: Post dataset event
+      x-openapi-router-controller: airflow.api_connexion.endpoints.dataset_endpoint
+      operationId: post_dataset_event
+      tags: [Dataset]
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/ExternalDatasetChange'
+      responses:
+        '200':
+          description: Success.
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/DatasetEvent'
+        '401':
+          $ref: '#/components/responses/Unauthenticated'
+        '403':
+          $ref: '#/components/responses/PermissionDenied'
+        '404':
+          $ref: '#/components/responses/NotFound'
+
 
   /config:
     get:
@@ -3980,6 +4006,22 @@ components:
               items:
                 $ref: '#/components/schemas/DatasetEvent'
         - $ref: '#/components/schemas/CollectionInfo'
+    ExternalDatasetChange:
+      description: |
+        A external dataset change that should create a DatasetEvent
+
+        *New in version 2.6.0*
+      type: object
+      properties:
+        dataset_uri:
+          type: string
+          description: The URI of the dataset
+          nullable: false
+        extra:
+          type: object
+          description: The dataset event extra
+          nullable: true
+
 
 
     # Configuration

airflow/api_connexion/schemas/dataset_schema.py

@@ -112,6 +112,14 @@ class Meta:
     data_interval_end = auto_field(dump_only=True)
 
 
+class ExternalDatasetChangeSchema(Schema):
+    """External Dataset change/update Schema"""
+
+    dataset_uri = fields.String()
+    timestamp = fields.DateTime()
+    extra = JsonObjectField()
+
+
 class DatasetEventSchema(SQLAlchemySchema):
     """Dataset Event DB schema."""
 
@@ -128,6 +136,8 @@ class Meta:
     source_dag_id = auto_field()
     source_run_id = auto_field()
     source_map_index = auto_field()
+    external_source = fields.String(dump_only=True)
+    user_id = auto_field()
     created_dagruns = fields.List(fields.Nested(BasicDAGRunSchema))
     timestamp = auto_field()
 
@@ -146,5 +156,6 @@ class DatasetEventCollectionSchema(Schema):
     total_entries = fields.Int()
 
 
+dataset_change_schema = ExternalDatasetChangeSchema()
 dataset_event_schema = DatasetEventSchema()
 dataset_event_collection_schema = DatasetEventCollectionSchema()

airflow/datasets/manager.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import datetime
 from typing import TYPE_CHECKING
 
 from sqlalchemy import exc
@@ -44,28 +45,78 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def register_dataset_change(
-        self, *, task_instance: TaskInstance, dataset: Dataset, extra=None, session: Session, **kwargs
-    ) -> None:
+        self,
+        *,
+        dataset: Dataset,
+        task_instance: TaskInstance,
+        extra=None,
+        session: Session,
+        **kwargs,
+    ) -> DatasetEvent | None:
         """
-        Register dataset related changes.
+        Register dataset related changes from a task instance.
 
         For local datasets, look them up, record the dataset event, queue dagruns, and broadcast
         the dataset event
         """
+        dataset_model = self._get_dataset_model(dataset=dataset, session=session)
+        if dataset_model is None:
+            return None
+
+        dataset_event = DatasetEvent(
+            dataset_id=dataset_model.id,
+            source_task_id=task_instance.task_id,
+            source_dag_id=task_instance.dag_id,
+            source_run_id=task_instance.run_id,
+            source_map_index=task_instance.map_index,
+            extra=extra,
+        )
+
+        self._save_dataset_event(dataset_event, dataset_model, session)
+
+        return dataset_event
+
+    def register_external_dataset_change(
+        self,
+        dataset: Dataset,
+        timestamp: datetime,
+        session: Session,
+        external_source: str | None,
+        user_id: int | None,
+        extra=None,
+        **kwargs,
+    ) -> DatasetEvent | None:
+        """
+        Register a dataset change from an external source (rather than task_instance)
+
+        For local datasets, look them up, record the dataset event, and queue dagruns.
+        """
+        dataset_model = self._get_dataset_model(dataset=dataset, session=session)
+        if dataset_model is None:
+            return None
+
+        # When an external dataset change is made through the API, it isn't triggered by a task instance,
+        # so we create a DatasetEvent without the task and dag data.
+        dataset_event = DatasetEvent(
+            dataset_id=dataset_model.id,
+            external_source=external_source,
+            user_id=user_id,
+            timestamp=timestamp,
+            extra=extra,
+        )
+
+        self._save_dataset_event(dataset_event, dataset_model, session)
+        return dataset_event
+
+    def _get_dataset_model(self, dataset: Dataset, session: Session) -> DatasetModel | None:
         dataset_model = session.query(DatasetModel).filter(DatasetModel.uri == dataset.uri).one_or_none()
         if not dataset_model:
             self.log.warning("DatasetModel %s not found", dataset)
-            return
-        session.add(
-            DatasetEvent(
-                dataset_id=dataset_model.id,
-                source_task_id=task_instance.task_id,
-                source_dag_id=task_instance.dag_id,
-                source_run_id=task_instance.run_id,
-                source_map_index=task_instance.map_index,
-                extra=extra,
-            )
-        )
+            return None
+        return dataset_model
+
+    def _save_dataset_event(self, dataset_event: DatasetEvent, dataset_model: DatasetModel, session: Session):
+        session.add(dataset_event)
         session.flush()
         Stats.incr("dataset.updates")
         if dataset_model.consuming_dags:

airflow/models/dataset.py

@@ -275,6 +275,8 @@ class DatasetEvent(Base):
     source_dag_id = Column(StringID(), nullable=True)
     source_run_id = Column(StringID(), nullable=True)
     source_map_index = Column(Integer, nullable=True, server_default=text("-1"))
+    external_source = Column(StringID(), nullable=True)
+    user_id = Column(StringID(), nullable=True)
     timestamp = Column(UtcDateTime, default=timezone.utcnow, nullable=False)
 
     __tablename__ = "dataset_event"

airflow/www/static/js/components/Table/Cells.tsx

@@ -148,6 +148,11 @@ export const TaskInstanceLink = ({ cell: { value, row } }: CellProps) => {
   const { sourceRunId, sourceDagId, sourceMapIndex } = row.original;
   const gridUrl = getMetaValue("grid_url");
   const dagId = getMetaValue("dag_id");
+
+  if (!value || !sourceRunId || !sourceDagId || !gridUrl) {
+    return null;
+  }
+
   const stringToReplace = dagId || "__DAG_ID__";
   const url = `${gridUrl?.replace(
     stringToReplace,
@@ -156,6 +161,7 @@ export const TaskInstanceLink = ({ cell: { value, row } }: CellProps) => {
     value
   )}`;
   const mapIndex = sourceMapIndex > -1 ? `[${sourceMapIndex}]` : "";
+
   return (
     <Box>
       <Link

airflow/www/static/js/types/api-generated.ts

@@ -584,6 +584,8 @@ export interface paths {
   "/datasets/events": {
     /** Get dataset events */
     get: operations["get_dataset_events"];
+    /** Post dataset event */
+    post: operations["post_dataset_event"];
     parameters: {
       query: {
         /** The numbers of items to return. */
@@ -1701,6 +1703,17 @@ export interface components {
     DatasetEventCollection: {
       dataset_events?: components["schemas"]["DatasetEvent"][];
     } & components["schemas"]["CollectionInfo"];
+    /**
+     * @description A external dataset change that should create a DatasetEvent
+     *
+     * *New in version 2.6.0*
+     */
+    ExternalDatasetChange: {
+      /** @description The URI of the dataset */
+      dataset_uri?: string;
+      /** @description The dataset event extra */
+      extra?: { [key: string]: unknown } | null;
+    };
     /** @description The option of configuration. */
     ConfigOption: {
       key?: string;
@@ -4119,6 +4132,50 @@ export interface operations {
       404: components["responses"]["NotFound"];
     };
   };
+  /** Post dataset event */
+  post_dataset_event: {
+    parameters: {
+      query: {
+        /** The numbers of items to return. */
+        limit?: components["parameters"]["PageLimit"];
+        /** The number of items to skip before starting to collect the result set. */
+        offset?: components["parameters"]["PageOffset"];
+        /**
+         * The name of the field to order the results by.
+         * Prefix a field name with `-` to reverse the sort order.
+         *
+         * *New in version 2.1.0*
+         */
+        order_by?: components["parameters"]["OrderBy"];
+        /** The Dataset ID that updated the dataset. */
+        dataset_id?: components["parameters"]["FilterDatasetID"];
+        /** The DAG ID that updated the dataset. */
+        source_dag_id?: components["parameters"]["FilterSourceDAGID"];
+        /** The task ID that updated the dataset. */
+        source_task_id?: components["parameters"]["FilterSourceTaskID"];
+        /** The DAG run ID that updated the dataset. */
+        source_run_id?: components["parameters"]["FilterSourceRunID"];
+        /** The map index that updated the dataset. */
+        source_map_index?: components["parameters"]["FilterSourceMapIndex"];
+      };
+    };
+    responses: {
+      /** Success. */
+      200: {
+        content: {
+          "application/json": components["schemas"]["DatasetEvent"];
+        };
+      };
+      401: components["responses"]["Unauthenticated"];
+      403: components["responses"]["PermissionDenied"];
+      404: components["responses"]["NotFound"];
+    };
+    requestBody: {
+      content: {
+        "application/json": components["schemas"]["ExternalDatasetChange"];
+      };
+    };
+  };
   get_config: {
     responses: {
       /** Success. */
@@ -4661,6 +4718,9 @@ export type BasicDAGRun = CamelCasedPropertiesDeep<
 export type DatasetEventCollection = CamelCasedPropertiesDeep<
   components["schemas"]["DatasetEventCollection"]
 >;
+export type ExternalDatasetChange = CamelCasedPropertiesDeep<
+  components["schemas"]["ExternalDatasetChange"]
+>;
 export type ConfigOption = CamelCasedPropertiesDeep<
   components["schemas"]["ConfigOption"]
 >;
@@ -4933,6 +4993,10 @@ export type GetDatasetVariables = CamelCasedPropertiesDeep<
 export type GetDatasetEventsVariables = CamelCasedPropertiesDeep<
   operations["get_dataset_events"]["parameters"]["query"]
 >;
+export type PostDatasetEventVariables = CamelCasedPropertiesDeep<
+  operations["post_dataset_event"]["parameters"]["query"] &
+    operations["post_dataset_event"]["requestBody"]["content"]["application/json"]
+>;
 export type GetPluginsVariables = CamelCasedPropertiesDeep<
   operations["get_plugins"]["parameters"]["query"]
 >;