Lee-W commented on code in PR #37176:
URL: https://github.com/apache/airflow/pull/37176#discussion_r1480883003
##########
airflow/api_connexion/endpoints/dataset_endpoint.py:
##########
@@ -124,3 +129,150 @@ def get_dataset_events(
return dataset_event_collection_schema.dump(
DatasetEventCollection(dataset_events=events,
total_entries=total_entries)
)
+
+
+def _get_ddrq(dag_id: str, uri: str, session: Session, before: str | None =
None) -> DatasetDagRunQueue:
+ """Get DatasetDagRunQueue."""
+ dataset = session.scalar(select(DatasetModel).where(DatasetModel.uri ==
uri))
+ dataset_id = dataset.id
+
+ where_clause_conditions = [
+ DatasetDagRunQueue.target_dag_id == dag_id,
+ DatasetDagRunQueue.dataset_id == dataset_id,
+ ]
+ if before is not None:
+ where_clause_conditions.append(DatasetDagRunQueue.created_at <
format_datetime(before))
+
+ return
session.scalar(select(DatasetDagRunQueue).where(*where_clause_conditions))
+
+
+@security.requires_access_dataset("GET")
+@security.requires_access_dag("GET")
+@provide_session
+def get_dag_dataset_pending_event(
+ *, dag_id: str, uri: str, before: str | None = None, session: Session =
NEW_SESSION
+) -> APIResponse:
+ """Get a pending dataset event for a DAG."""
+ ddrq = _get_ddrq(dag_id=dag_id, uri=uri, session=session, before=before)
+ if ddrq is None:
+ raise NotFound(
+ "Pending event not found",
+ detail=f"Pending event with dag_id: `{dag_id}` and daatset uri:
`{uri}` was not found",
Review Comment:
fixed. Thanks!
##########
airflow/api_connexion/endpoints/dataset_endpoint.py:
##########
@@ -124,3 +129,150 @@ def get_dataset_events(
return dataset_event_collection_schema.dump(
DatasetEventCollection(dataset_events=events,
total_entries=total_entries)
)
+
+
+def _get_ddrq(dag_id: str, uri: str, session: Session, before: str | None =
None) -> DatasetDagRunQueue:
+ """Get DatasetDagRunQueue."""
+ dataset = session.scalar(select(DatasetModel).where(DatasetModel.uri ==
uri))
+ dataset_id = dataset.id
+
+ where_clause_conditions = [
+ DatasetDagRunQueue.target_dag_id == dag_id,
+ DatasetDagRunQueue.dataset_id == dataset_id,
+ ]
+ if before is not None:
+ where_clause_conditions.append(DatasetDagRunQueue.created_at <
format_datetime(before))
+
+ return
session.scalar(select(DatasetDagRunQueue).where(*where_clause_conditions))
+
+
+@security.requires_access_dataset("GET")
+@security.requires_access_dag("GET")
+@provide_session
+def get_dag_dataset_pending_event(
+ *, dag_id: str, uri: str, before: str | None = None, session: Session =
NEW_SESSION
+) -> APIResponse:
+ """Get a pending dataset event for a DAG."""
+ ddrq = _get_ddrq(dag_id=dag_id, uri=uri, session=session, before=before)
+ if ddrq is None:
+ raise NotFound(
+ "Pending event not found",
+ detail=f"Pending event with dag_id: `{dag_id}` and daatset uri:
`{uri}` was not found",
+ )
+ return dataset_dag_run_queue_schema.dump(ddrq)
+
+
+@security.requires_access_dataset("DELETE")
+@security.requires_access_dag("DELETE")
+@provide_session
+def delete_dag_dataset_pending_event(
+ *, dag_id: str, uri: str, before: str | None = None, session: Session =
NEW_SESSION
+) -> APIResponse:
+ """Delete a pending dataset event for a DAG."""
+ ddrq = _get_ddrq(dag_id=dag_id, uri=uri, session=session, before=before)
+ if ddrq is None:
+ raise NotFound(
+ "Pending event not found",
+ detail=f"Pending event with dag_id: `{dag_id}` and daatset uri:
`{uri}` was not found",
Review Comment:
fixed. Thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
