You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2020/10/08 01:15:26 UTC

[GitHub] [airflow] OmairK commented on a change in pull request #9597: [WIP] Add read-only endpoints for task instances

OmairK commented on a change in pull request #9597:
URL: https://github.com/apache/airflow/pull/9597#discussion_r501395345



##########
File path: airflow/api_connexion/endpoints/task_instance_endpoint.py
##########
@@ -14,39 +14,243 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from typing import Any, List, Optional, Tuple
+
+from flask import request
+from marshmallow import ValidationError
+from sqlalchemy import and_, func, or_
 
-# TODO(mik-laj): We have to implement it.
-#     Do you want to help? Please look at: https://github.com/apache/airflow/issues/8132
 from airflow.api_connexion import security
+from airflow.api_connexion.exceptions import NotFound, BadRequest
+from airflow.api_connexion.parameters import format_datetime, format_parameters
+from airflow.api_connexion.schemas.task_instance_schema import (
+    clear_task_instance_form,
+    TaskInstanceCollection,
+    task_instance_collection_schema,
+    task_instance_schema,
+    task_instance_batch_form,
+    task_instance_reference_collection_schema,
+    TaskInstanceReferenceCollection,
+)
+from airflow.configuration import conf
+from airflow.models.dagrun import DagRun as DR
+from airflow.models.taskinstance import clear_task_instances, TaskInstance as TI
+from airflow.models.dag import DagModel
+from airflow.models.dagbag import DagBag
+from airflow.models import SlaMiss
+from airflow.utils.state import State
+from airflow.utils.session import provide_session
 
 
 @security.requires_authentication
-def get_task_instance():
-    """
-    Get a task instance
-    """
-    raise NotImplementedError("Not implemented yet.")
+@provide_session
+def get_task_instance(dag_id: str, dag_run_id: str, task_id: str, session=None):
+    """Get task instance"""
+    query = (
+        session.query(TI)
+        .filter(TI.dag_id == dag_id)
+        .join(DR, and_(TI.dag_id == DR.dag_id, TI.execution_date == DR.execution_date))
+        .filter(DR.run_id == dag_run_id)
+        .filter(TI.task_id == task_id)
+        .join(
+            SlaMiss,
+            and_(
+                SlaMiss.dag_id == TI.dag_id,
+                SlaMiss.execution_date == TI.execution_date,
+                SlaMiss.task_id == TI.task_id,
+            ),
+            isouter=True,
+        )
+        .add_entity(SlaMiss)
+    )
+    task_instance = query.one_or_none()
+    if task_instance is None:
+        raise NotFound("Task instance not found")
+
+    return task_instance_schema.dump(task_instance)
+
+
+def _apply_array_filter(query, key, values):
+    if values is not None:
+        query = query.filter(or_(*[key == v for v in values]))
+    return query
+
+
+def _apply_range_filter(query, key, value_range: Tuple[Any, Any]):
+    gte_value, lte_value = value_range
+    if gte_value is not None:
+        query = query.filter(key >= gte_value)
+    if lte_value is not None:
+        query = query.filter(key <= lte_value)
+    return query
 
 
+@format_parameters(
+    {
+        "execution_date_gte": format_datetime,
+        "execution_date_lte": format_datetime,
+        "start_date_gte": format_datetime,
+        "start_date_lte": format_datetime,
+        "end_date_gte": format_datetime,
+        "end_date_lte": format_datetime,
+    }
+)
 @security.requires_authentication
-def get_task_instances():
+@provide_session
+def get_task_instances(
+    limit: int,
+    dag_id: Optional[str] = None,
+    dag_run_id: Optional[str] = None,
+    execution_date_gte: Optional[str] = None,
+    execution_date_lte: Optional[str] = None,
+    start_date_gte: Optional[str] = None,
+    start_date_lte: Optional[str] = None,
+    end_date_gte: Optional[str] = None,
+    end_date_lte: Optional[str] = None,
+    duration_gte: Optional[float] = None,
+    duration_lte: Optional[float] = None,
+    state: Optional[str] = None,
+    pool: Optional[List[str]] = None,
+    queue: Optional[List[str]] = None,
+    offset: Optional[int] = None,
+    session=None,
+):  # pylint: disable=too-many-arguments
     """
-    Get list of task instances of DAG.
+    Get list of task instances.
     """
-    raise NotImplementedError("Not implemented yet.")
+    query = session.query(TI)
+
+    if dag_id != "~":
+        query = query.filter(TI.dag_id == dag_id)
+    if dag_run_id != "~":
+        query = query.join(DR, and_(TI.dag_id == DR.dag_id, TI.execution_date == DR.execution_date))
+        query = query.filter(DR.run_id == dag_run_id)
+    query = _apply_range_filter(
+        query,
+        key=DR.execution_date,
+        value_range=(execution_date_gte, execution_date_lte),
+    )
+    query = _apply_range_filter(query, key=TI.start_date, value_range=(start_date_gte, start_date_lte))
+    query = _apply_range_filter(query, key=TI.end_date, value_range=(end_date_gte, end_date_lte))
+    query = _apply_range_filter(query, key=TI.duration, value_range=(duration_gte, duration_lte))
+    query = _apply_array_filter(query, key=TI.state, values=state)
+    query = _apply_array_filter(query, key=TI.pool, values=pool)
+    query = _apply_array_filter(query, key=TI.queue, values=queue)
+
+    query = query.join(
+        SlaMiss,
+        and_(
+            SlaMiss.dag_id == TI.dag_id,
+            SlaMiss.task_id == TI.task_id,
+            SlaMiss.execution_date == TI.execution_date,
+        ),
+        isouter=True,
+    ).add_entity(SlaMiss)
+    task_instances = query.offset(offset).limit(limit).all()
+    total_entries = session.query(
+        func.count(func.distinct(TI.dag_id, TI.execution_date, TI.task_id))
+    ).scalar()

Review comment:
       @mik-laj Passing multiple columns to distinct works perfectly with postgres but raises error in case of sqlite how should we handle this case. 




----------------------------------------------------------------
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org