aminghadersohi commented on code in PR #35018:
URL: https://github.com/apache/superset/pull/35018#discussion_r2334351225
##########
superset/daos/base.py:
##########
@@ -83,39 +196,138 @@ def find_by_id_or_uuid(
return None
@classmethod
- def find_by_id(
- cls,
- model_id: str | int,
- skip_base_filter: bool = False,
- ) -> T | None:
+ def _apply_base_filter(
+ cls, query: Any, skip_base_filter: bool = False, data_model: Any = None
+ ) -> Any:
"""
- Find a model by id, if defined applies `base_filter`
+ Apply the base_filter to the query if it exists and skip_base_filter
is False.
"""
- query = db.session.query(cls.model_cls)
if cls.base_filter and not skip_base_filter:
- data_model = SQLAInterface(cls.model_cls, db.session)
+ if data_model is None:
+ data_model = SQLAInterface(cls.model_cls, db.session)
query = cls.base_filter( # pylint: disable=not-callable
cls.id_column_name, data_model
).apply(query, None)
- id_column = getattr(cls.model_cls, cls.id_column_name)
+ return query
+
+ @classmethod
+ def _convert_value_for_column(cls, column: Any, value: Any) -> Any:
+ """
+ Convert a value to the appropriate type for a given SQLAlchemy column.
+
+ Args:
+ column: SQLAlchemy column object
+ value: Value to convert
+
+ Returns:
+ Converted value or None if conversion fails
+ """
+ if (
+ hasattr(column.type, "python_type")
+ and column.type.python_type == uuid_lib.UUID
+ ):
+ if isinstance(value, str):
+ try:
+ return uuid_lib.UUID(value)
+ except (ValueError, AttributeError):
+ return None
+ return value
+
+ @classmethod
+ def _find_by_column(
+ cls,
+ column_name: str,
+ value: str | int,
+ skip_base_filter: bool = False,
+ ) -> T | None:
+ """
+ Private method to find a model by any column value.
+
+ Args:
+ column_name: Name of the column to search by
+ value: Value to search for
+ skip_base_filter: Whether to skip base filtering
+
+ Returns:
+ Model instance or None if not found
+ """
+ query = db.session.query(cls.model_cls)
+ query = cls._apply_base_filter(query, skip_base_filter)
+
+ if not hasattr(cls.model_cls, column_name):
+ return None
+
+ column = getattr(cls.model_cls, column_name)
+ converted_value = cls._convert_value_for_column(column, value)
+ if converted_value is None:
+ return None
+
try:
- return query.filter(id_column == model_id).one_or_none()
+ return query.filter(column == converted_value).one_or_none()
except StatementError:
# can happen if int is passed instead of a string or similar
return None
+ @classmethod
+ def find_by_id(
+ cls,
+ model_id: str | int,
+ skip_base_filter: bool = False,
+ id_column: str | None = None,
+ ) -> T | None:
+ """
+ Find a model by ID using specified or default ID column.
+
+ Args:
+ model_id: ID value to search for
+ skip_base_filter: Whether to skip base filtering
+ id_column: Column name to use (defaults to cls.id_column_name)
+
+ Returns:
+ Model instance or None if not found
+ """
+ column = id_column or cls.id_column_name
+ return cls._find_by_column(column, model_id, skip_base_filter)
+
@classmethod
def find_by_ids(
cls,
- model_ids: list[str] | list[int],
+ model_ids: Sequence[str | int],
skip_base_filter: bool = False,
+ id_column: str | None = None,
) -> list[T]:
"""
Find a List of models by a list of ids, if defined applies
`base_filter`
+
+ :param model_ids: List of IDs to find
+ :param skip_base_filter: If true, skip applying the base filter
+ :param id_column: Optional column name to use for ID lookup
+ (defaults to id_column_name)
"""
- id_col = getattr(cls.model_cls, cls.id_column_name, None)
+ column = id_column or cls.id_column_name
+ id_col = getattr(cls.model_cls, column, None)
if id_col is None or not model_ids:
return []
+
+ # Convert IDs to appropriate types based on column type
+ converted_ids: list[str | int | uuid_lib.UUID] = []
+ for id_val in model_ids:
+ converted_value = cls._convert_value_for_column(id_col, id_val)
+ if converted_value is not None:
+ converted_ids.append(converted_value)
Review Comment:
Good catch! You're absolutely right - we should validate type consistency.
I'll add a check to ensure all converted IDs are of the same type, with a
warning logged for mixed types. This will help catch bugs early while
maintaining backward compatibility.
--
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.
To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org
For additional commands, e-mail: notifications-h...@superset.apache.org
