uranusjr commented on a change in pull request #14895:
URL: https://github.com/apache/airflow/pull/14895#discussion_r599905501
##########
File path: airflow/api_connexion/parameters.py
##########
@@ -86,3 +87,32 @@ def wrapped_function(*args, **kwargs):
return cast(T, wrapped_function)
return format_parameters_decorator
+
+
+def apply_sorting(model, query, order_by, to_replace=None, allowed_attrs=None):
+ """Apply sorting to query"""
+ lstriped_orderby = order_by.lstrip('-')
+ if allowed_attrs and lstriped_orderby not in allowed_attrs:
+ modelname = model.__tablename__.capitalize()
+ model_mapping = {
+ "Ab_user": 'User',
+ "Slot_pool": "Pool",
+ "Dag_run": "DagRun",
+ "Dag": "DagModel",
+ "Ab_role": "Role",
+ "Import_error": "ImportError",
+ }
+ if model_mapping.get(modelname, None):
+ modelname = model_mapping[modelname]
+ raise BadRequest(
+ detail=f"Ordering with '{lstriped_orderby}' is disallowed or "
+ f"the attribute does not exist on {modelname} model"
+ )
+ if to_replace:
+ for key, value in to_replace.items():
+ if key == order_by:
+ order_by = value
Review comment:
```suggestion
order_by = to_replace.get(order_by, order_by)
```
Also I feel we don’t need to repeat the ascending and descnding variants in
the `to_replace` mapping. You can strip the `"-"` prefix before look up the
`to_replace` mapping instead.
Something like
```python
if order_by[0] == "-":
func = desc
order_by = order_by[1:]
else:
func = asc
if to_replace: # This does not need to include `"-id"` etc. anymore, only
`"id"`.
order_by = to_replace.get(order_by, order_by)
return query.order_by(func(order_by))
```
##########
File path: airflow/api_connexion/parameters.py
##########
@@ -86,3 +87,32 @@ def wrapped_function(*args, **kwargs):
return cast(T, wrapped_function)
return format_parameters_decorator
+
+
+def apply_sorting(model, query, order_by, to_replace=None, allowed_attrs=None):
+ """Apply sorting to query"""
+ lstriped_orderby = order_by.lstrip('-')
+ if allowed_attrs and lstriped_orderby not in allowed_attrs:
+ modelname = model.__tablename__.capitalize()
+ model_mapping = {
+ "Ab_user": 'User',
+ "Slot_pool": "Pool",
+ "Dag_run": "DagRun",
+ "Dag": "DagModel",
+ "Ab_role": "Role",
+ "Import_error": "ImportError",
+ }
+ if model_mapping.get(modelname, None):
+ modelname = model_mapping[modelname]
Review comment:
```suggestion
modelname = model_mapping.get(modelname, modelname)
```
It feels weird to me the field name replacements are declared outside of
this function, but model name replacements inside. Maybe it’d be a good idea to
somehow put them together? Maybe a few seperate classes (enums?) to hold those
definitions. Not entirely sure.
##########
File path: airflow/api_connexion/parameters.py
##########
@@ -86,3 +87,32 @@ def wrapped_function(*args, **kwargs):
return cast(T, wrapped_function)
return format_parameters_decorator
+
+
+def apply_sorting(model, query, order_by, to_replace=None, allowed_attrs=None):
+ """Apply sorting to query"""
+ lstriped_orderby = order_by.lstrip('-')
+ if allowed_attrs and lstriped_orderby not in allowed_attrs:
+ modelname = model.__tablename__.capitalize()
Review comment:
Why `.capitalize()`?
##########
File path: airflow/api_connexion/endpoints/connection_endpoint.py
##########
@@ -63,11 +63,15 @@ def get_connection(connection_id, session):
@security.requires_access([(permissions.ACTION_CAN_READ,
permissions.RESOURCE_CONNECTION)])
@format_parameters({'limit': check_limit})
@provide_session
-def get_connections(session, limit, offset=0):
+def get_connections(session, limit, offset=0, order_by="id"):
"""Get all connection entries"""
+ to_replace = {"connection_id": "conn_id", "-connection_id": "-conn_id"}
+ allowed_filter_attrs = ['connection_id', 'conn_type', 'description',
'host', 'port', 'id']
+
total_entries = session.query(func.count(Connection.id)).scalar()
query = session.query(Connection)
- connections =
query.order_by(Connection.id).offset(offset).limit(limit).all()
+ query = apply_sorting(Connection, query, order_by, to_replace,
allowed_filter_attrs)
Review comment:
I feel there should be a way to avoid repeating `Connection` here. Same
goes for other calls below.
--
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:
[email protected]
