From 59eab52f517904f42a0ad6ac7632bb08f8699f54 Mon Sep 17 00:00:00 2001 From: prasad-madine Date: Thu, 6 Mar 2025 13:21:39 +0530 Subject: [PATCH 1/5] fix:bug Add multisort to dags list request #46383 --- airflow/api_fastapi/common/parameters.py | 54 +++++++++++++++++++----- 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index f69d64bd5d031..11a009fb8d9f1 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -164,6 +164,7 @@ def depends_search(value: str | None = Query(alias=pattern_name, default=None)) class SortParam(BaseParam[str]): """Order result by the attribute.""" + MAX_SORT_PARAMS = 2 def __init__( self, allowed_attrs: list[str], model: Base, to_replace: dict[str, str | Column] | None = None @@ -173,14 +174,31 @@ def __init__( self.model = model self.to_replace = to_replace - def to_orm(self, select: Select) -> Select: - if self.skip_none is False: - raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}") + def get_order_by_columns(self, order_by_list: list[str]) -> list: + """Generates order_by conditions based on the given sorting parameters.""" + if len(order_by_list) > MAX_SORT_PARAMS: + raise HTTPException( + 400, + f"Ordering with more than two parameters is not allowed. Provided: {order_by_list}" + ) - if self.value is None: - self.value = self.get_primary_key_string() + order_by_columns = [] + + if len(order_by_list) >= 1: + order_by_columns.append(self.get_column_with_sort(order_by_list[0])) + + if len(order_by_list) == 2: + order_by_columns.append(self.get_column_with_sort(order_by_list[1])) - lstriped_orderby = self.value.lstrip("-") + primary_key_name = self.get_primary_key_string() + if not (order_by_list and primary_key_name in [item.lstrip("-") for item in order_by_list]): + order_by_columns.append(self.get_column_with_sort(primary_key_name)) + + return order_by_columns + + def get_column_with_sort(self, order_by: str): + """Helper function to process each order_by field with sorting and null handling.""" + lstriped_orderby = order_by.lstrip("-") column: Column | None = None if self.to_replace: replacement = self.to_replace.get(lstriped_orderby, lstriped_orderby) @@ -201,16 +219,30 @@ def to_orm(self, select: Select) -> Select: # MySQL does not support `nullslast`, and True/False ordering depends on the # database implementation. nullscheck = case((column.isnot(None), 0), else_=1) + + return (nullscheck, column.desc() if order_by.startswith("-") else column.asc()) + + def to_orm(self, select: Select) -> Select: + 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() + + order_by_list = self.value.split(",") if self.value else [] + + order_by_columns = self.get_order_by_columns(order_by_list) # Reset default sorting select = select.order_by(None) - primary_key_column = self.get_primary_key_column() - - if self.value[0] == "-": - return select.order_by(nullscheck, column.desc(), primary_key_column.desc()) + if order_by_columns: + 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()) + primary_key_column = self.get_primary_key_column() + select = select.order_by(primary_key_column.asc()) + + return select def get_primary_key_column(self) -> Column: """Get the primary key column of the model of SortParam object.""" From a03fb8ac66e197c686d8ec422b0224b39aa493bf Mon Sep 17 00:00:00 2001 From: prasad-madine Date: Fri, 7 Mar 2025 10:24:01 +0530 Subject: [PATCH 2/5] fix:review comments --- airflow/api_fastapi/common/parameters.py | 96 ++++++++----------- .../core_api/routes/public/dags.py | 2 +- 2 files changed, 41 insertions(+), 57 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 11a009fb8d9f1..032ee46755d53 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -164,7 +164,6 @@ def depends_search(value: str | None = Query(alias=pattern_name, default=None)) class SortParam(BaseParam[str]): """Order result by the attribute.""" - MAX_SORT_PARAMS = 2 def __init__( self, allowed_attrs: list[str], model: Base, to_replace: dict[str, str | Column] | None = None @@ -174,64 +173,49 @@ def __init__( self.model = model self.to_replace = to_replace - def get_order_by_columns(self, order_by_list: list[str]) -> list: - """Generates order_by conditions based on the given sorting parameters.""" - if len(order_by_list) > MAX_SORT_PARAMS: - raise HTTPException( - 400, - f"Ordering with more than two parameters is not allowed. Provided: {order_by_list}" - ) - - order_by_columns = [] - - if len(order_by_list) >= 1: - order_by_columns.append(self.get_column_with_sort(order_by_list[0])) - - if len(order_by_list) == 2: - order_by_columns.append(self.get_column_with_sort(order_by_list[1])) - - primary_key_name = self.get_primary_key_string() - if not (order_by_list and primary_key_name in [item.lstrip("-") for item in order_by_list]): - order_by_columns.append(self.get_column_with_sort(primary_key_name)) - - return order_by_columns - - def get_column_with_sort(self, order_by: str): - """Helper function to process each order_by field with sorting and null handling.""" - 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) - - return (nullscheck, column.desc() if order_by.startswith("-") else column.asc()) - 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() - - order_by_list = self.value.split(",") if self.value else [] - - order_by_columns = self.get_order_by_columns(order_by_list) + + order_by_list = self.value + if len(order_by_list) > MAX_SORT_PARAMS: + raise HTTPException( + 400, + f"Ordering with more than {MAX_SORT_PARAMS} parameters is not allowed. Provided: {order_by_list}" + ) + 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) @@ -256,9 +240,9 @@ 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] = Query(default or self.get_primary_key_string())) -> SortParam: + return self.set_value(order_by if order_by else [self.get_primary_key_string()]) return inner diff --git a/airflow/api_fastapi/core_api/routes/public/dags.py b/airflow/api_fastapi/core_api/routes/public/dags.py index 3eaf487948226..f3d6d232aa70c 100644 --- a/airflow/api_fastapi/core_api/routes/public/dags.py +++ b/airflow/api_fastapi/core_api/routes/public/dags.py @@ -95,7 +95,7 @@ def get_dags( ), ], order_by: Annotated[ - SortParam, + list[str], Depends( SortParam( ["dag_id", "dag_display_name", "next_dagrun", "state", "start_date"], From 96085fb46a8785c97711f7a97b0c9084f9b6d785 Mon Sep 17 00:00:00 2001 From: prasad-madine Date: Fri, 7 Mar 2025 10:28:46 +0530 Subject: [PATCH 3/5] fix:review comments --- airflow/api_fastapi/common/parameters.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 032ee46755d53..485446a23584c 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -221,6 +221,9 @@ def to_orm(self, select: Select) -> Select: select = select.order_by(None) if order_by_columns: + primary_key_column = self.get_primary_key_column() + primary_key_sort = primary_key_column.asc() + order_by_columns.append((case((primary_key_column.isnot(None), 0), else_=1), primary_key_sort)) select = select.order_by(*[col for pair in order_by_columns for col in pair]) else: primary_key_column = self.get_primary_key_column() From 05e51047f12a9ced4a655fd8a56741dfc00ffaa0 Mon Sep 17 00:00:00 2001 From: prasad-madine Date: Thu, 13 Mar 2025 12:30:56 +0530 Subject: [PATCH 4/5] fix: Added test cases and fixed logic --- airflow/api_fastapi/common/parameters.py | 10 +++++++--- airflow/api_fastapi/core_api/routes/public/dags.py | 2 +- tests/api_fastapi/core_api/routes/public/test_dags.py | 11 +++++++++++ 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 485446a23584c..7ee54f20d9766 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -208,7 +208,7 @@ def to_orm(self, select: Select) -> Select: if column is None: column = getattr(self.model, lstriped_orderby) - # MySQL does not support `nullslast`, and True/False ordering depends on the + # MySQL does not support nullslast, and True/False ordering depends on the # database implementation. nullscheck = case((column.isnot(None), 0), else_=1) @@ -244,8 +244,12 @@ def depends(cls, *args: Any, **kwargs: Any) -> Self: raise NotImplementedError("Use dynamic_depends, depends not implemented.") def dynamic_depends(self, default: list[str] | None = None) -> Callable: - def inner(order_by: list[str] = Query(default or self.get_primary_key_string())) -> SortParam: - return self.set_value(order_by if order_by else [self.get_primary_key_string()]) + 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/api_fastapi/core_api/routes/public/dags.py b/airflow/api_fastapi/core_api/routes/public/dags.py index f3d6d232aa70c..3eaf487948226 100644 --- a/airflow/api_fastapi/core_api/routes/public/dags.py +++ b/airflow/api_fastapi/core_api/routes/public/dags.py @@ -95,7 +95,7 @@ def get_dags( ), ], order_by: Annotated[ - list[str], + SortParam, Depends( SortParam( ["dag_id", "dag_display_name", "next_dagrun", "state", "start_date"], diff --git a/tests/api_fastapi/core_api/routes/public/test_dags.py b/tests/api_fastapi/core_api/routes/public/test_dags.py index 9c4650636b22f..0369ea893b4fb 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dags.py +++ b/tests/api_fastapi/core_api/routes/public/test_dags.py @@ -227,6 +227,17 @@ class TestGetDags(TestDagEndpoint): 3, [DAG3_ID, DAG1_ID, DAG2_ID], ), + ({"order_by": ["-dag_id", "dag_display_name"]}, 2, [DAG2_ID, DAG1_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]), + ({"order_by": ["dag_display_name", "dag_id"]}, 2, [DAG1_ID, DAG2_ID]), + ({"order_by": ["-dag_display_name", "-dag_id"]}, 2, [DAG2_ID, DAG1_ID]), + ({"order_by": ["last_run_state", "dag_id"], "only_active": False},3, [DAG1_ID, DAG3_ID, DAG2_ID]), + ({"order_by": ["-last_run_state", "-dag_id"], "only_active": False},3, [DAG3_ID, DAG1_ID, DAG2_ID]), + ({"order_by": ["-last_run_start_date", "dag_id"], "only_active": False},3, [DAG3_ID, DAG1_ID, DAG2_ID]), + ({"order_by": ["last_run_start_date", "-dag_id"], "only_active": False},3, [DAG1_ID, DAG3_ID, DAG2_ID]), # Search ({"dag_id_pattern": "1"}, 1, [DAG1_ID]), ({"dag_display_name_pattern": "test_dag2"}, 1, [DAG2_ID]), From d1a769c0eba50f5e3f8e1664eec44a8b527e0109 Mon Sep 17 00:00:00 2001 From: prasad-madine Date: Wed, 26 Mar 2025 17:57:15 +0530 Subject: [PATCH 5/5] fix: addressed the comments --- airflow/api_fastapi/common/parameters.py | 18 ++++++++++++------ .../core_api/routes/public/test_dags.py | 17 +++++++++++------ 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index fb96b94dcd16e..a5235d9180657 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -217,14 +217,20 @@ def to_orm(self, select: Select) -> Select: # Reset default sorting select = select.order_by(None) + primary_key_column = self.get_primary_key_column() + primary_key_name = primary_key_column.name + + primary_key_included = any(order_by.lstrip("-") == primary_key_name for order_by in order_by_list) + if order_by_columns: - primary_key_column = self.get_primary_key_column() - primary_key_sort = primary_key_column.asc() - order_by_columns.append((case((primary_key_column.isnot(None), 0), else_=1), primary_key_sort)) - select = select.order_by(*[col for pair in order_by_columns for col in pair]) + 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: - primary_key_column = self.get_primary_key_column() - select = select.order_by(primary_key_column.asc()) + select = select.order_by(primary_key_sort) return select diff --git a/tests/api_fastapi/core_api/routes/public/test_dags.py b/tests/api_fastapi/core_api/routes/public/test_dags.py index 5424eda2bff36..c0711023a7989 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dags.py +++ b/tests/api_fastapi/core_api/routes/public/test_dags.py @@ -228,17 +228,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]), - ({"order_by": ["dag_display_name", "dag_id"]}, 2, [DAG1_ID, DAG2_ID]), - ({"order_by": ["-dag_display_name", "-dag_id"]}, 2, [DAG2_ID, DAG1_ID]), - ({"order_by": ["last_run_state", "dag_id"], "only_active": False},3, [DAG1_ID, DAG3_ID, DAG2_ID]), - ({"order_by": ["-last_run_state", "-dag_id"], "only_active": False},3, [DAG3_ID, DAG1_ID, DAG2_ID]), - ({"order_by": ["-last_run_start_date", "dag_id"], "only_active": False},3, [DAG3_ID, DAG1_ID, DAG2_ID]), - ({"order_by": ["last_run_start_date", "-dag_id"], "only_active": False},3, [DAG1_ID, DAG3_ID, DAG2_ID]), # Search ({"dag_id_pattern": "1"}, 1, [DAG1_ID]), ({"dag_display_name_pattern": "test_dag2"}, 1, [DAG2_ID]),