mik-laj commented on a change in pull request #9597: URL: https://github.com/apache/airflow/pull/9597#discussion_r448590758
########## File path: airflow/api_connexion/endpoints/task_instance_endpoint.py ########## @@ -14,23 +14,110 @@ # 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 -# TODO(mik-laj): We have to implement it. -# Do you want to help? Please look at: https://github.com/apache/airflow/issues/8132 +from sqlalchemy import and_, func, or_ +from airflow.api_connexion.exceptions import NotFound +from airflow.api_connexion.parameters import format_datetime, format_parameters +from airflow.api_connexion.schemas.task_instance_schema import ( + TaskInstanceCollection, task_instance_collection_schema, task_instance_schema, +) +from airflow.models.dagrun import DagRun as DR +from airflow.models.taskinstance import TaskInstance as TI +from airflow.utils.session import provide_session + + +@provide_session +def get_task_instance(dag_id: str, dag_run_id: str, task_id: str, session=None): + 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) + ) + + 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 get_task_instance(): - """ - Get a task instance - """ - raise NotImplementedError("Not implemented yet.") +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 get_task_instances(): + +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( + { + 'start_date_gte': format_datetime, + 'start_date_lte': format_datetime, + 'execution_date_gte': format_datetime, + 'execution_date_lte': format_datetime, + 'end_date_gte': format_datetime, + 'end_date_lte': format_datetime, + } +) +@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, Review comment: It is not correct with pylint if I remember correctly. I didn't specify the type here because SQLAlchemy does a lot of magic, and mypy likes to get lost when he is aware of its existence. ---------------------------------------------------------------- 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: us...@infra.apache.org
