refactor: add sqlalchemy ulid type to all relevant models

Author: Panaetius

bases/renku_data_services/background_jobs/core.py

@@ -11,6 +11,7 @@
     SubjectFilter,
     WriteRelationshipsRequest,
 )
+from ulid import ULID
 
 from renku_data_services.authz.authz import Authz, ResourceType, _AuthzConverter, _Relation
 from renku_data_services.authz.models import Scope
@@ -117,7 +118,7 @@ async def fix_mismatched_project_namespace_ids(config: SyncConfig) -> None:
                                 relation=rel.relationship.relation,
                                 subject=SubjectReference(
                                     object=ObjectReference(
-                                        object_type=ResourceType.group.value, object_id=correct_group_id
+                                        object_type=ResourceType.group.value, object_id=str(correct_group_id)
                                     )
                                 ),
                             ),
@@ -169,7 +170,7 @@ async def migrate_groups_make_all_public(config: SyncConfig) -> None:
     all_users = SubjectReference(object=_AuthzConverter.all_users())
     all_anon_users = SubjectReference(object=_AuthzConverter.anonymous_users())
     for group_id in groups_to_process:
-        group_res = _AuthzConverter.group(group_id)
+        group_res = _AuthzConverter.group(ULID.from_str(group_id))
         all_users_are_viewers = Relationship(
             resource=group_res,
             relation=_Relation.public_viewer.value,
@@ -228,7 +229,7 @@ async def migrate_user_namespaces_make_all_public(config: SyncConfig) -> None:
     all_users = SubjectReference(object=_AuthzConverter.all_users())
     all_anon_users = SubjectReference(object=_AuthzConverter.anonymous_users())
     for ns_id in namespaces_to_process:
-        namespace_res = _AuthzConverter.user_namespace(ns_id)
+        namespace_res = _AuthzConverter.user_namespace(ULID.from_str(ns_id))
         all_users_are_viewers = Relationship(
             resource=namespace_res,
             relation=_Relation.public_viewer.value,

components/renku_data_services/authz/authz.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass, field
 from enum import StrEnum
 from functools import wraps
-from typing import ClassVar, Concatenate, ParamSpec, Protocol, TypeVar
+from typing import ClassVar, Concatenate, ParamSpec, Protocol, TypeVar, cast
 
 from authzed.api.v1 import AsyncClient
 from authzed.api.v1.core_pb2 import ObjectReference, Relationship, RelationshipUpdate, SubjectReference, ZedToken
@@ -163,12 +163,12 @@ def all_users() -> ObjectReference:
         return ObjectReference(object_type=ResourceType.user, object_id="*")
 
     @staticmethod
-    def group(id: str) -> ObjectReference:
-        return ObjectReference(object_type=ResourceType.group, object_id=id)
+    def group(id: ULID) -> ObjectReference:
+        return ObjectReference(object_type=ResourceType.group, object_id=str(id))
 
     @staticmethod
-    def user_namespace(id: str) -> ObjectReference:
-        return ObjectReference(object_type=ResourceType.user_namespace, object_id=id)
+    def user_namespace(id: ULID) -> ObjectReference:
+        return ObjectReference(object_type=ResourceType.user_namespace, object_id=str(id))
 
     @staticmethod
     def to_object(resource_type: ResourceType, resource_id: str | ULID | int) -> ObjectReference:
@@ -179,9 +179,9 @@ def to_object(resource_type: ResourceType, resource_id: str | ULID | int) -> Obj
                 return _AuthzConverter.user(sid)
             case (ResourceType.anonymous_user, _):
                 return _AuthzConverter.anonymous_users()
-            case (ResourceType.user_namespace, rid) if isinstance(rid, str):
+            case (ResourceType.user_namespace, rid) if isinstance(rid, ULID):
                 return _AuthzConverter.user_namespace(rid)
-            case (ResourceType.group, rid) if isinstance(rid, str):
+            case (ResourceType.group, rid) if isinstance(rid, ULID):
                 return _AuthzConverter.group(rid)
         raise errors.ProgrammingError(
             message=f"Unexpected or unknown resource type when checking permissions {resource_type}"
@@ -580,7 +580,7 @@ def _add_project(self, project: Project) -> _AuthzChange:
             object=(
                 _AuthzConverter.user_namespace(project.namespace.id)
                 if project.namespace.kind == NamespaceKind.user
-                else _AuthzConverter.group(project.namespace.underlying_resource_id)
+                else _AuthzConverter.group(cast(ULID, project.namespace.underlying_resource_id))
             )
         )
         project_in_platform = Relationship(
@@ -759,10 +759,14 @@ async def _update_project_namespace(
             else SubjectReference(object=_AuthzConverter.user_namespace(project.namespace.id))
         )
         old_namespace_sub = (
-            SubjectReference(object=_AuthzConverter.group(current_namespace.relationship.subject.object.object_id))
+            SubjectReference(
+                object=_AuthzConverter.group(ULID.from_str(current_namespace.relationship.subject.object.object_id))
+            )
             if current_namespace.relationship.subject.object.object_type == ResourceType.group.value
             else SubjectReference(
-                object=_AuthzConverter.user_namespace(current_namespace.relationship.subject.object.object_id)
+                object=_AuthzConverter.user_namespace(
+                    ULID.from_str(current_namespace.relationship.subject.object.object_id)
+                )
             )
         )
         new_namespace = Relationship(
@@ -1093,7 +1097,7 @@ async def _remove_group(
                 message="Cannot remove a group in the authorization database if the group has no ID"
             )
         consistency = Consistency(at_least_as_fresh=zed_token) if zed_token else Consistency(fully_consistent=True)
-        rel_filter = RelationshipFilter(resource_type=ResourceType.group.value, optional_resource_id=group.id)
+        rel_filter = RelationshipFilter(resource_type=ResourceType.group.value, optional_resource_id=str(group.id))
         responses = self.client.ReadRelationships(
             ReadRelationshipsRequest(consistency=consistency, relationship_filter=rel_filter)
         )
@@ -1113,7 +1117,7 @@ async def upsert_group_members(
         self,
         user: base_models.APIUser,
         resource_type: ResourceType,
-        resource_id: str,
+        resource_id: ULID,
         members: list[Member],
         *,
         zed_token: ZedToken | None = None,
@@ -1124,8 +1128,9 @@ async def upsert_group_members(
         add_members: list[RelationshipUpdate] = []
         undo: list[RelationshipUpdate] = []
         output: list[MembershipChange] = []
+        resource_id_str = str(resource_id)
         expected_user_roles = {_Relation.viewer.value, _Relation.owner.value, _Relation.editor.value}
-        existing_owners_rels = await self._get_resource_owners(resource_type, resource_id, consistency)
+        existing_owners_rels = await self._get_resource_owners(resource_type, resource_id_str, consistency)
         n_existing_owners = len(existing_owners_rels)
         for member in members:
             rel = Relationship(
@@ -1135,7 +1140,7 @@ async def upsert_group_members(
             )
             existing_rel_filter = RelationshipFilter(
                 resource_type=resource_type.value,
-                optional_resource_id=resource_id,
+                optional_resource_id=resource_id_str,
                 optional_subject_filter=SubjectFilter(
                     subject_type=ResourceType.user, optional_subject_id=member.user_id
                 ),
@@ -1231,7 +1236,7 @@ async def remove_group_members(
         self,
         user: base_models.APIUser,
         resource_type: ResourceType,
-        resource_id: str,
+        resource_id: ULID,
         user_ids: list[str],
         *,
         zed_token: ZedToken | None = None,
@@ -1242,12 +1247,13 @@ async def remove_group_members(
         remove_members: list[RelationshipUpdate] = []
         output: list[MembershipChange] = []
         existing_owners_rels: list[ReadRelationshipsResponse] | None = None
+        resource_id_str = str(resource_id)
         for user_id in user_ids:
             if user_id == "*":
                 raise errors.ValidationError(message="Cannot remove a group member with ID '*'")
             existing_rel_filter = RelationshipFilter(
                 resource_type=resource_type.value,
-                optional_resource_id=resource_id,
+                optional_resource_id=resource_id_str,
                 optional_subject_filter=SubjectFilter(subject_type=ResourceType.user, optional_subject_id=user_id),
             )
             existing_rels: AsyncIterable[ReadRelationshipsResponse] = self.client.ReadRelationships(
@@ -1258,7 +1264,9 @@ async def remove_group_members(
             async for existing_rel in existing_rels:
                 if existing_rel.relationship.relation == _Relation.owner.value:
                     if existing_owners_rels is None:
-                        existing_owners_rels = await self._get_resource_owners(resource_type, resource_id, consistency)
+                        existing_owners_rels = await self._get_resource_owners(
+                            resource_type, resource_id_str, consistency
+                        )
                     if len(existing_owners_rels) == 1:
                         raise errors.ValidationError(
                             message="You are trying to remove the single last owner of the group, "

components/renku_data_services/authz/models.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from enum import Enum
 
+from ulid import ULID
+
 from renku_data_services.errors import errors
 from renku_data_services.namespace.apispec import GroupRole
 
@@ -56,7 +58,7 @@ class Member:
 
     role: Role
     user_id: str
-    resource_id: str
+    resource_id: str | ULID
 
 
 class Change(Enum):

components/renku_data_services/connected_services/apispec_base.py

@@ -1,6 +1,7 @@
 """Base models for API specifications."""
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
+from ulid import ULID
 
 
 class BaseAPISpec(BaseModel):
@@ -14,6 +15,12 @@ class Config:
         # this rust crate does not support lookahead regex syntax but we need it in this component
         regex_engine = "python-re"
 
+    @field_validator("id", mode="before", check_fields=False)
+    @classmethod
+    def serialize_id(cls, id: str | ULID) -> str:
+        """Custom serializer that can handle ULIDs."""
+        return str(id)
+
 
 class AuthorizeParams(BaseAPISpec):
     """The schema for the query parameters used in the authorize request."""

components/renku_data_services/connected_services/blueprints.py

@@ -7,6 +7,7 @@
 from sanic.log import logger
 from sanic.response import JSONResponse
 from sanic_ext import validate
+from ulid import ULID
 
 import renku_data_services.base_models as base_models
 from renku_data_services.base_api.auth import authenticate, only_admins, only_authenticated
@@ -150,34 +151,34 @@ def get_one(self) -> BlueprintFactoryResponse:
         """Get a specific OAuth2 connection."""
 
         @authenticate(self.authenticator)
-        async def _get_one(_: Request, user: base_models.APIUser, connection_id: str) -> JSONResponse:
+        async def _get_one(_: Request, user: base_models.APIUser, connection_id: ULID) -> JSONResponse:
             connection = await self.connected_services_repo.get_oauth2_connection(
                 connection_id=connection_id, user=user
             )
             return validated_json(apispec.Connection, connection)
 
-        return "/oauth2/connections/<connection_id>", ["GET"], _get_one
+        return "/oauth2/connections/<connection_id:ulid>", ["GET"], _get_one
 
     def get_account(self) -> BlueprintFactoryResponse:
         """Get the account information for a specific OAuth2 connection."""
 
         @authenticate(self.authenticator)
-        async def _get_account(_: Request, user: base_models.APIUser, connection_id: str) -> JSONResponse:
+        async def _get_account(_: Request, user: base_models.APIUser, connection_id: ULID) -> JSONResponse:
             account = await self.connected_services_repo.get_oauth2_connected_account(
                 connection_id=connection_id, user=user
             )
             return validated_json(apispec.ConnectedAccount, account)
 
-        return "/oauth2/connections/<connection_id>/account", ["GET"], _get_account
+        return "/oauth2/connections/<connection_id:ulid>/account", ["GET"], _get_account
 
     def get_token(self) -> BlueprintFactoryResponse:
         """Get the access token for a specific OAuth2 connection."""
 
         @authenticate(self.authenticator)
-        async def _get_token(_: Request, user: base_models.APIUser, connection_id: str) -> JSONResponse:
+        async def _get_token(_: Request, user: base_models.APIUser, connection_id: ULID) -> JSONResponse:
             token = await self.connected_services_repo.get_oauth2_connection_token(
                 connection_id=connection_id, user=user
             )
             return json(token.dump_for_api())
 
-        return "/oauth2/connections/<connection_id>/token", ["GET"], _get_token
+        return "/oauth2/connections/<connection_id:ulid>/token", ["GET"], _get_token

components/renku_data_services/connected_services/db.py

@@ -11,6 +11,7 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
+from ulid import ULID
 
 import renku_data_services.base_models as base_models
 from renku_data_services import errors
@@ -282,7 +283,7 @@ async def get_oauth2_connections(
             connections = result.all()
             return [c.dump() for c in connections]
 
-    async def get_oauth2_connection(self, connection_id: str, user: base_models.APIUser) -> models.OAuth2Connection:
+    async def get_oauth2_connection(self, connection_id: ULID, user: base_models.APIUser) -> models.OAuth2Connection:
         """Get one OAuth2 connection from the database."""
         if not user.is_authenticated or user.id is None:
             raise errors.MissingResourceError(
@@ -303,7 +304,7 @@ async def get_oauth2_connection(self, connection_id: str, user: base_models.APIU
             return connection.dump()
 
     async def get_oauth2_connected_account(
-        self, connection_id: str, user: base_models.APIUser
+        self, connection_id: ULID, user: base_models.APIUser
     ) -> models.ConnectedAccount:
         """Get the account information from a OAuth2 connection."""
         async with self.get_async_oauth2_client(connection_id=connection_id, user=user) as (oauth2_client, _, adapter):
@@ -316,7 +317,9 @@ async def get_oauth2_connected_account(
             account = adapter.api_validate_account_response(response)
             return account
 
-    async def get_oauth2_connection_token(self, connection_id: str, user: base_models.APIUser) -> models.OAuth2TokenSet:
+    async def get_oauth2_connection_token(
+        self, connection_id: ULID, user: base_models.APIUser
+    ) -> models.OAuth2TokenSet:
         """Get the OAuth2 access token from one connection from the database."""
         async with self.get_async_oauth2_client(connection_id=connection_id, user=user) as (oauth2_client, _, _):
             await oauth2_client.ensure_active_token(oauth2_client.token)
@@ -325,7 +328,7 @@ async def get_oauth2_connection_token(self, connection_id: str, user: base_model
 
     @asynccontextmanager
     async def get_async_oauth2_client(
-        self, connection_id: str, user: base_models.APIUser
+        self, connection_id: ULID, user: base_models.APIUser
     ) -> AsyncGenerator[tuple[AsyncOAuth2Client, schemas.OAuth2ConnectionORM, ProviderAdapter], None]:
         """Get the AsyncOAuth2Client for the given connection_id and user."""
         if not user.is_authenticated or user.id is None:

components/renku_data_services/connected_services/models.py

@@ -4,6 +4,8 @@
 from datetime import UTC, datetime
 from typing import Any
 
+from ulid import ULID
+
 from renku_data_services.connected_services.apispec import ConnectionStatus, ProviderKind
 
 
@@ -28,7 +30,7 @@ class OAuth2Client:
 class OAuth2Connection:
     """OAuth2 connection model."""
 
-    id: str
+    id: ULID
     provider_id: str
     status: ConnectionStatus
 

components/renku_data_services/connected_services/orm.py

@@ -11,6 +11,7 @@
 
 from renku_data_services.connected_services import models
 from renku_data_services.connected_services.apispec import ConnectionStatus, ProviderKind
+from renku_data_services.utils.sqlalchemy import ULIDType
 
 JSONVariant = JSON().with_variant(JSONB(), "postgresql")
 
@@ -72,7 +73,7 @@ class OAuth2ConnectionORM(BaseORM):
     """An OAuth2 connection."""
 
     __tablename__ = "oauth2_connections"
-    id: Mapped[str] = mapped_column("id", String(26), primary_key=True, default_factory=lambda: str(ULID()), init=False)
+    id: Mapped[ULID] = mapped_column("id", ULIDType, primary_key=True, default_factory=lambda: str(ULID()), init=False)
     user_id: Mapped[str] = mapped_column("user_id", String())
     client_id: Mapped[str] = mapped_column(ForeignKey(OAuth2ClientORM.id, ondelete="CASCADE"), index=True)
     client: Mapped[OAuth2ClientORM] = relationship(init=False, repr=False)