diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index e68001d5b539a..ba8352b6afc17 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -171,43 +171,68 @@ def __init__( self.to_replace = to_replace def to_orm(self, select: Select) -> Select: + MAX_SORT_PARAMS = 5 if self.skip_none is False: raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}") if self.value is None: self.value = self.get_primary_key_string() - - lstriped_orderby = self.value.lstrip("-") - column: Column | None = None - if self.to_replace: - replacement = self.to_replace.get(lstriped_orderby, lstriped_orderby) - if isinstance(replacement, str): - lstriped_orderby = replacement - else: - column = replacement - - if (self.allowed_attrs and lstriped_orderby not in self.allowed_attrs) and column is None: + + order_by_list = self.value + if len(order_by_list) > MAX_SORT_PARAMS: raise HTTPException( - 400, - f"Ordering with '{lstriped_orderby}' is disallowed or " - f"the attribute does not exist on the model", + 400, + f"Ordering with more than {MAX_SORT_PARAMS} parameters is not allowed. Provided: {order_by_list}" ) - if column is None: - column = getattr(self.model, lstriped_orderby) - - # MySQL does not support `nullslast`, and True/False ordering depends on the - # database implementation. - nullscheck = case((column.isnot(None), 0), else_=1) + order_by_columns = [] + + for order_by in order_by_list: + lstriped_orderby = order_by.lstrip("-") + column: Column | None = None + if self.to_replace: + replacement = self.to_replace.get(lstriped_orderby, lstriped_orderby) + if isinstance(replacement, str): + lstriped_orderby = replacement + else: + column = replacement + + if (self.allowed_attrs and lstriped_orderby not in self.allowed_attrs) and column is None: + raise HTTPException( + 400, + f"Ordering with '{lstriped_orderby}' is disallowed or " + f"the attribute does not exist on the model", + ) + if column is None: + column = getattr(self.model, lstriped_orderby) + + # MySQL does not support nullslast, and True/False ordering depends on the + # database implementation. + nullscheck = case((column.isnot(None), 0), else_=1) + + if order_by[0] == "-": + order_by_columns.append((nullscheck, column.desc())) + else: + order_by_columns.append((nullscheck, column.asc())) # Reset default sorting select = select.order_by(None) primary_key_column = self.get_primary_key_column() + primary_key_name = primary_key_column.name - if self.value[0] == "-": - return select.order_by(nullscheck, column.desc(), primary_key_column.desc()) + primary_key_included = any(order_by.lstrip("-") == primary_key_name for order_by in order_by_list) + + if order_by_columns: + if not primary_key_included: + first_sort_desc = order_by_list[0].startswith("-") if order_by_list else False + primary_key_sort = primary_key_column.desc() if first_sort_desc else primary_key_column.asc() + select = select.order_by(*[col for pair in order_by_columns for col in pair], primary_key_sort) + else: + select = select.order_by(*[col for pair in order_by_columns for col in pair]) else: - return select.order_by(nullscheck, column.asc(), primary_key_column.asc()) + select = select.order_by(primary_key_sort) + + return select def get_primary_key_column(self) -> Column: """Get the primary key column of the model of SortParam object.""" @@ -221,9 +246,13 @@ def get_primary_key_string(self) -> str: def depends(cls, *args: Any, **kwargs: Any) -> Self: raise NotImplementedError("Use dynamic_depends, depends not implemented.") - def dynamic_depends(self, default: str | None = None) -> Callable: - def inner(order_by: str = default or self.get_primary_key_string()) -> SortParam: - return self.set_value(self.get_primary_key_string() if order_by == "" else order_by) + def dynamic_depends(self, default: list[str] | None = None) -> Callable: + def inner(order_by: list[str] | str | None = Query(default)) -> SortParam: + if order_by is None: + order_by = [self.get_primary_key_string()] + elif isinstance(order_by, str): + order_by = [order_by] + return self.set_value(order_by) return inner diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py index 95693fdcbf643..9de1fd31a2aba 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py @@ -230,6 +230,22 @@ class TestGetDags(TestDagEndpoint): 3, [DAG3_ID, DAG1_ID, DAG2_ID], ), + ({"order_by": ["last_run_state", "dag_display_name"], "only_active": False}, 3, [DAG1_ID, DAG3_ID, DAG2_ID]), + ({"order_by": ["-last_run_state", "dag_display_name"], "only_active": False}, 3, [DAG3_ID, DAG1_ID, DAG2_ID]), + ({"order_by": ["last_run_start_date", "dag_display_name"], "only_active": False}, 3, [DAG1_ID, DAG3_ID, DAG2_ID]), + ({"order_by": ["-last_run_start_date", "dag_display_name"], "only_active": False}, 3, [DAG3_ID, DAG1_ID, DAG2_ID]), + ({"order_by": ["-dag_id", "dag_display_name"]}, 2, [DAG2_ID, DAG1_ID]), + ({"order_by": ["dag_id", "dag_display_name"]}, 2, [DAG1_ID, DAG2_ID]), + ({"order_by": ["dag_display_name", "dag_id"]}, 2, [DAG1_ID, DAG2_ID]), + ({"order_by": ["last_run_state"], "only_active": False}, 3, [DAG1_ID, DAG3_ID, DAG2_ID]), + ({"order_by": ["-last_run_state"], "only_active": False}, 3, [DAG3_ID, DAG1_ID, DAG2_ID]), + ({"order_by": ["last_run_start_date"], "only_active": False}, 3, [DAG1_ID, DAG3_ID, DAG2_ID]), + ({"order_by": ["-last_run_start_date"], "only_active": False}, 3, [DAG3_ID, DAG1_ID, DAG2_ID]) + + ({"order_by": ["dag_display_name", "-next_dagrun"]}, 2, [DAG1_ID, DAG2_ID]), + ({"order_by": ["last_run_state", "-dag_display_name", "dag_id"]}, 2, [DAG1_ID, DAG2_ID]), + ({"order_by": ["-last_run_start_date", "dag_display_name", "next_dagrun", "dag_id"]}, 2, [ DAG1_ID, DAG2_ID]), + ({"order_by": ["dag_display_name", "-last_run_state", "next_dagrun", "dag_id", "last_run_start_date"]}, 2, [DAG1_ID, DAG2_ID]), # Search ({"dag_id_pattern": "1"}, 1, [DAG1_ID]), ({"dag_display_name_pattern": "test_dag2"}, 1, [DAG2_ID]),