feat(api): dataset read API uuid support (#34836)

superset/daos/base.py

@@ -47,12 +47,41 @@ class BaseDAO(Generic[T]):
     Child classes can register base filtering to be applied to all filter methods
     """
     id_column_name = "id"
+    uuid_column_name = "uuid"
 
     def __init_subclass__(cls) -> None:
         cls.model_cls = get_args(
             cls.__orig_bases__[0]  # type: ignore  # pylint: disable=no-member
         )[0]
 
+    @classmethod
+    def find_by_id_or_uuid(
+        cls,
+        model_id_or_uuid: str,
+        skip_base_filter: bool = False,
+    ) -> T | None:
+        """
+        Find a model by id or uuid, if defined applies `base_filter`
+        """
+        query = db.session.query(cls.model_cls)
+        if cls.base_filter and not skip_base_filter:
+            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)
+        uuid_column = getattr(cls.model_cls, cls.uuid_column_name)
+
+        if model_id_or_uuid.isdigit():
+            filter = id_column == int(model_id_or_uuid)
+        else:
+            filter = uuid_column == model_id_or_uuid
+        try:
+            return query.filter(filter).one_or_none()
+        except StatementError:
+            # can happen if neither uuid nor int is passed
+            return None
+
     @classmethod
     def find_by_id(
         cls,

superset/daos/datasource.py

@@ -16,12 +16,17 @@
 # under the License.
 
 import logging
+import uuid
 from typing import Union
 
 from superset import db
 from superset.connectors.sqla.models import SqlaTable
 from superset.daos.base import BaseDAO
-from superset.daos.exceptions import DatasourceNotFound, DatasourceTypeNotSupportedError
+from superset.daos.exceptions import (
+    DatasourceNotFound,
+    DatasourceTypeNotSupportedError,
+    DatasourceValueIsIncorrect,
+)
 from superset.models.sql_lab import Query, SavedQuery
 from superset.utils.core import DatasourceType
 
@@ -41,22 +46,34 @@ class DatasourceDAO(BaseDAO[Datasource]):
     def get_datasource(
         cls,
         datasource_type: Union[DatasourceType, str],
-        datasource_id: int,
+        database_id_or_uuid: int | str,
     ) -> Datasource:
         if datasource_type not in cls.sources:
             raise DatasourceTypeNotSupportedError()
 
+        model = cls.sources[datasource_type]
+
+        if str(database_id_or_uuid).isdigit():
+            filter = model.id == int(database_id_or_uuid)
+        else:
+            try:
+                uuid.UUID(str(database_id_or_uuid))  # uuid validation
+                filter = model.uuid == database_id_or_uuid
+            except ValueError as err:
+                logger.warning(
+                    f"database_id_or_uuid {database_id_or_uuid} isn't valid uuid"
+                )
+                raise DatasourceValueIsIncorrect() from err
+
         datasource = (
-            db.session.query(cls.sources[datasource_type])
-            .filter_by(id=datasource_id)
-            .one_or_none()
+            db.session.query(cls.sources[datasource_type]).filter(filter).one_or_none()
         )
 
         if not datasource:
             logger.warning(
-                "Datasource not found datasource_type: %s, datasource_id: %s",
+                "Datasource not found datasource_type: %s, database_id_or_uuid: %s",
                 datasource_type,
-                datasource_id,
+                database_id_or_uuid,
             )
             raise DatasourceNotFound()
 

superset/daos/exceptions.py

@@ -68,3 +68,8 @@ class DatasourceTypeNotSupportedError(DAOException):
 class DatasourceNotFound(DAOException):
     status = 404
     message = "Datasource does not exist"
+
+
+class DatasourceValueIsIncorrect(DAOException):
+    status = 422
+    message = "Datasource value is neither id or uuid"