|
17 | 17 |
|
18 | 18 | from __future__ import annotations
|
19 | 19 |
|
20 |
| -import importlib |
21 | 20 | from abc import ABC, abstractmethod
|
22 | 21 | from datetime import datetime
|
23 | 22 | from typing import TYPE_CHECKING, Any, Callable, Generic, List, TypeVar
|
@@ -195,32 +194,27 @@ def to_orm(self, select: Select) -> Select:
|
195 | 194 | # Reset default sorting
|
196 | 195 | select = select.order_by(None)
|
197 | 196 |
|
| 197 | + primary_key_column = self.get_primary_key_column() |
| 198 | + |
198 | 199 | if self.value[0] == "-":
|
199 |
| - return select.order_by(nullscheck, column.desc(), column.desc()) |
| 200 | + return select.order_by(nullscheck, column.desc(), primary_key_column) |
200 | 201 | else:
|
201 |
| - return select.order_by(nullscheck, column.asc(), column.asc()) |
202 |
| - |
203 |
| - def get_primary_key(self) -> str: |
204 |
| - """Get the primary key of the model of SortParam object.""" |
205 |
| - return inspect(self.model).primary_key[0].name |
| 202 | + return select.order_by(nullscheck, column.asc(), primary_key_column) |
206 | 203 |
|
207 |
| - @staticmethod |
208 |
| - def get_primary_key_of_given_model_string(model_string: str) -> str: |
209 |
| - """ |
210 |
| - Get the primary key of given 'airflow.models' class as a string. The class should have driven be from 'airflow.models.base'. |
| 204 | + def get_primary_key_column(self) -> Column: |
| 205 | + """Get the primary key column of the model of SortParam object.""" |
| 206 | + return inspect(self.model).primary_key[0] |
211 | 207 |
|
212 |
| - :param model_string: The string representation of the model class. |
213 |
| - :return: The primary key of the model class. |
214 |
| - """ |
215 |
| - dynamic_return_model = getattr(importlib.import_module("airflow.models"), model_string) |
216 |
| - return inspect(dynamic_return_model).primary_key[0].name |
| 208 | + def get_primary_key_string(self) -> str: |
| 209 | + """Get the primary key string of the model of SortParam object.""" |
| 210 | + return self.get_primary_key_column().name |
217 | 211 |
|
218 | 212 | def depends(self, *args: Any, **kwargs: Any) -> Self:
|
219 | 213 | raise NotImplementedError("Use dynamic_depends, depends not implemented.")
|
220 | 214 |
|
221 | 215 | def dynamic_depends(self) -> Callable:
|
222 |
| - def inner(order_by: str = self.get_primary_key()) -> SortParam: |
223 |
| - return self.set_value(self.get_primary_key() if order_by == "" else order_by) |
| 216 | + def inner(order_by: str = self.get_primary_key_string()) -> SortParam: |
| 217 | + return self.set_value(self.get_primary_key_string() if order_by == "" else order_by) |
224 | 218 |
|
225 | 219 | return inner
|
226 | 220 |
|
|
0 commit comments