diff options
| -rw-r--r-- | gn_auth/auth/authorisation/resources/system/models.py | 17 |
1 files changed, 12 insertions, 5 deletions
diff --git a/gn_auth/auth/authorisation/resources/system/models.py b/gn_auth/auth/authorisation/resources/system/models.py index 303b0ac..25089fa 100644 --- a/gn_auth/auth/authorisation/resources/system/models.py +++ b/gn_auth/auth/authorisation/resources/system/models.py @@ -1,9 +1,10 @@ """Base functions and utilities for system resources.""" from uuid import UUID from functools import reduce -from typing import Sequence +from typing import Union, Sequence + +from gn_libs import sqlite3 as db -from gn_auth.auth.db import sqlite3 as db from gn_auth.auth.errors import NotFoundError from gn_auth.auth.authentication.users import User @@ -52,9 +53,9 @@ def user_roles_on_system(conn: db.DbConnection, user: User) -> Sequence[Role]: return tuple() -def system_resource(conn: db.DbConnection) -> Resource: +def system_resource(conn: Union[db.DbConnection, db.DbCursor]) -> Resource: """Retrieve the system resource.""" - with db.cursor(conn) as cursor: + def __fetch_sys_resource__(cursor: db.DbCursor) -> Resource: cursor.execute( "SELECT resource_categories.*, resources.resource_id, " "resources.resource_name, resources.public " @@ -65,4 +66,10 @@ def system_resource(conn: db.DbConnection) -> Resource: if row: return resource_from_dbrow(row) - raise NotFoundError("Could not find a system resource!") + raise NotFoundError("Could not find a system resource!") + + if hasattr(conn, "cursor"): # is connection + with db.cursor(conn) as cursor: + return __fetch_sys_resource__(cursor) + else: + return __fetch_sys_resource__(conn) |