using NamedTuple instead of Tuple

apache · Jan 27, 2022 · b2b230c · b2b230c
1 parent c5ddcfc
commit b2b230c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 17 deletions.
diff --git a/superset/databases/filters.py b/superset/databases/filters.py
@@ -25,27 +25,22 @@
 
 class DatabaseFilter(BaseFilter):
     # TODO(bogdan): consider caching.
-    def schema_access_databases(self) -> Set[str]:  # noqa pylint: disable=no-self-use
-        return {
-            security_manager.unpack_perm(vm)[0]
-            for vm in security_manager.user_view_menu_names("schema_access")
-        }
 
-    def datasource_access_databases(  # noqa pylint: disable=no-self-use
-        self,
+    def can_access_databases(  # noqa pylint: disable=no-self-use
+        self, view_menu_name: str,
     ) -> Set[str]:
         return {
-            security_manager.unpack_perm(vm)[0]
-            for vm in security_manager.user_view_menu_names("datasource_access")
+            security_manager.unpack_database_and_schema(vm).database
+            for vm in security_manager.user_view_menu_names(view_menu_name)
         }
 
     def apply(self, query: Query, value: Any) -> Query:
         if security_manager.can_access_all_databases():
             return query
         database_perms = security_manager.user_view_menu_names("database_access")
-        schema_access_databases = self.schema_access_databases()
+        schema_access_databases = self.can_access_databases("schema_access")
 
-        datasource_access_databases = self.datasource_access_databases()
+        datasource_access_databases = self.can_access_databases("datasource_access")
 
         return query.filter(
             or_(

diff --git a/superset/security/manager.py b/superset/security/manager.py
@@ -25,9 +25,9 @@
     cast,
     Dict,
     List,
+    NamedTuple,
     Optional,
     Set,
-    Tuple,
     TYPE_CHECKING,
     Union,
 )
@@ -78,6 +78,11 @@
 logger = logging.getLogger(__name__)
 
 
+class DatabaseAndSchema(NamedTuple):
+    database: str
+    schema: str
+
+
 class SupersetSecurityListWidget(ListWidget):  # pylint: disable=too-few-public-methods
     """
     Redeclaring to avoid circular imports
@@ -237,13 +242,14 @@ def get_schema_perm(  # pylint: disable=no-self-use
 
         return None
 
-    def unpack_perm(  # pylint: disable=no-self-use
+    def unpack_database_and_schema(  # pylint: disable=no-self-use
         self, schema_permission: str
-    ) -> Tuple[str, str]:
-        # [database_name].[schema_name]
+    ) -> DatabaseAndSchema:
+        # [database_name].[schema|table]
+
         schema_name = schema_permission.split(".")[1][1:-1]
         database_name = schema_permission.split(".")[0][1:-1]
-        return database_name, schema_name
+        return DatabaseAndSchema(database_name, schema_name)
 
     def can_access(self, permission_name: str, view_name: str) -> bool:
         """
@@ -532,7 +538,7 @@ def get_schemas_accessible_by_user(
 
         # schema_access
         accessible_schemas = {
-            self.unpack_perm(s)[1]
+            self.unpack_database_and_schema(s).schema
             for s in self.user_view_menu_names("schema_access")
             if s.startswith(f"[{database}].")
         }