paskia 0.9.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

paskia/fastapi/logging.py

@@ -0,0 +1,218 @@
+"""Custom access logging middleware for FastAPI/Uvicorn."""
+
+import logging
+import sys
+import time
+from ipaddress import IPv6Address
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+
+logger = logging.getLogger("paskia.access")
+
+_RESET = "\033[0m"
+_STATUS_INFO = "\033[32m"  # 1xx (green)
+_STATUS_OK = "\033[92m"  # 2xx (bright green)
+_STATUS_REDIRECT = "\033[32m"  # 3xx (green)
+_STATUS_CLIENT_ERR = "\033[0;31m"  # 4xx (red)
+_STATUS_SERVER_ERR = "\033[1;31m"  # 5xx (bright red)
+_METHOD_READ = "\033[0;34m"  # GET, HEAD, OPTIONS (blue)
+_METHOD_WRITE = "\033[1;34m"  # POST, PUT, DELETE, PATCH (bright blue)
+_HOST = "\033[1;30m"  # hostname (dark grey)
+_PATH = "\033[0m"  # path (default)
+_TIMING = "\033[2m"  # timing (dim)
+_WS_OPEN = "\033[1;33m"  # WebSocket connect (bright yellow)
+_WS_CLOSE = "\033[0;33m"  # WebSocket disconnect (yellow)
+_WS_STATUS = "\033[1;30m"  # WebSocket close status (dark grey)
+
+
+def format_ipv6_network(ip: str) -> str:
+    """Format IPv6 address to show only network part (first 64 bits)."""
+    try:
+        addr = IPv6Address(ip)
+        # Get the integer representation and mask to first 64 bits
+        network_int = int(addr) >> 64
+        # Format as IPv6 with trailing ::
+        # Split into 4 groups of 16 bits
+        groups = []
+        for _ in range(4):
+            groups.insert(0, format(network_int & 0xFFFF, "x"))
+            network_int >>= 16
+        # Compress consecutive zero groups
+        result = ":".join(groups) + "::"
+        # Simplify leading zeros in groups and compress
+        return str(IPv6Address(result + "0"))
+    except Exception:
+        return ip
+
+
+def format_client_ip(ip: str) -> str:
+    """Format client IP, compressing IPv6 to network part only."""
+    if not ip or ip == "-":
+        return "-"
+    if ":" in ip:
+        return format_ipv6_network(ip)
+    return ip
+
+
+def status_color(status: int) -> str:
+    """Return color code based on HTTP status."""
+    if status < 200:
+        return _STATUS_INFO
+    if status < 300:
+        return _STATUS_OK
+    if status < 400:
+        return _STATUS_REDIRECT
+    if status < 500:
+        return _STATUS_CLIENT_ERR
+    return _STATUS_SERVER_ERR
+
+
+def method_color(method: str) -> str:
+    """Return color code based on HTTP method."""
+    if method in ("GET", "HEAD", "OPTIONS"):
+        return _METHOD_READ
+    return _METHOD_WRITE
+
+
+def format_access_log(
+    client: str, status: int, method: str, host: str, path: str, duration_ms: float
+) -> str:
+    """Format access log line with colors and aligned fields."""
+    use_color = sys.stderr.isatty()
+
+    # Format components with fixed widths for alignment
+    ip = format_client_ip(client).ljust(15)  # IPv4 max 15 chars
+    timing = f"{duration_ms:.0f}ms"
+    method_padded = method.ljust(7)  # Longest method is OPTIONS (7)
+
+    if use_color:
+        status_str = f"{status_color(status)}{status}{_RESET}"
+        timing_str = f"{_TIMING}{timing}{_RESET}"
+        method_str = f"{method_color(method)}{method_padded}{_RESET}"
+        host_str = f"{_HOST}{host}{_RESET}"
+        path_str = f"{_PATH}{path}{_RESET}"
+    else:
+        status_str = str(status)
+        timing_str = timing
+        method_str = method_padded
+        host_str = host
+        path_str = path
+
+    # Format: "IP STATUS METHOD host path TIMING"
+    return f"{ip} {status_str} {method_str} {host_str}{path_str} {timing_str}"
+
+
+# WebSocket connection counter (mod 100)
+_ws_counter = 0
+
+
+def _next_ws_id() -> int:
+    """Get next WebSocket connection ID (0-99)."""
+    global _ws_counter
+    ws_id = _ws_counter
+    _ws_counter = (_ws_counter + 1) % 100
+    return ws_id
+
+
+def log_ws_open(client: str, host: str, path: str) -> int:
+    """Log WebSocket connection open. Returns connection ID for use in close."""
+    use_color = sys.stderr.isatty()
+    ws_id = _next_ws_id()
+
+    ip = format_client_ip(client).ljust(15)
+    id_str = f"{ws_id:02d}".ljust(7)  # Align with method field (7 chars)
+
+    if use_color:
+        # 🔌 aligned with status (takes ~2 char width), ID aligned with method
+        prefix = f"🔌  {_WS_OPEN}{id_str}{_RESET}"
+        host_str = f"{_HOST}{host}{_RESET}"
+        path_str = f"{_PATH}{path}{_RESET}"
+    else:
+        prefix = f"WS+ {id_str}"
+        host_str = host
+        path_str = path
+
+    logger.info(f"{ip} {prefix} {host_str}{path_str}")
+    return ws_id
+
+
+# WebSocket close codes to human-readable status
+WS_CLOSE_CODES = {
+    1000: "ok",
+    1001: "going away",
+    1002: "protocol error",
+    1003: "unsupported",
+    1005: "no status",
+    1006: "abnormal",
+    1007: "invalid data",
+    1008: "policy violation",
+    1009: "too large",
+    1010: "extension required",
+    1011: "server error",
+    1012: "restarting",
+    1013: "try again",
+    1014: "bad gateway",
+    1015: "tls error",
+}
+
+
+def log_ws_close(
+    client: str, ws_id: int, close_code: int | None, duration_ms: float
+) -> None:
+    """Log WebSocket connection close with duration and status."""
+    use_color = sys.stderr.isatty()
+
+    ip = format_client_ip(client).ljust(15)
+    id_str = f"{ws_id:02d}".ljust(7)  # Align with method field (7 chars)
+    timing = f"{duration_ms:.0f}ms"
+
+    # Convert close code to status text
+    if close_code is None:
+        status = "closed"
+    else:
+        status = WS_CLOSE_CODES.get(close_code, f"code {close_code}")
+
+    if use_color:
+        # 🔌 aligned with status, ID aligned with method
+        prefix = f"🔌  {_WS_CLOSE}{id_str}{_RESET}"
+        status_str = f"{_WS_STATUS}{status}{_RESET}"
+        timing_str = f"{_TIMING}{timing}{_RESET}"
+    else:
+        prefix = f"WS- {id_str}"
+        status_str = status
+        timing_str = timing
+
+    logger.info(f"{ip} {prefix} {status_str} {timing_str}")
+
+
+class AccessLogMiddleware(BaseHTTPMiddleware):
+    """Middleware that logs HTTP requests with custom format."""
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        start = time.perf_counter()
+        response = await call_next(request)
+        duration_ms = (time.perf_counter() - start) * 1000
+
+        client = request.client.host if request.client else "-"
+        host = request.headers.get("host", "-")
+        method = request.method
+        path = request.url.path
+        if request.url.query:
+            path = f"{path}?{request.url.query}"
+        status = response.status_code
+
+        line = format_access_log(client, status, method, host, path, duration_ms)
+        logger.info(line)
+
+        return response
+
+
+def configure_access_logging():
+    """Configure the access logger to output to stderr."""
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setFormatter(logging.Formatter("%(message)s"))
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)
+    logger.propagate = False

paskia/fastapi/mainapp.py

@@ -10,10 +10,16 @@ from fastapi_vue import Frontend
 
 from paskia import globals
 from paskia.db import start_background, stop_background
+from paskia.db.logging import configure_db_logging
 from paskia.fastapi import admin, api, auth_host, ws
+from paskia.fastapi.logging import AccessLogMiddleware, configure_access_logging
 from paskia.fastapi.session import AUTH_COOKIE
 from paskia.util import hostutil, passphrase, vitedev
 
+# Configure custom logging
+configure_access_logging()
+configure_db_logging()
+
 # Vue Frontend static files
 frontend = Frontend(
     Path(__file__).parent.parent / "frontend-build",
@@ -48,10 +54,11 @@ async def lifespan(app: FastAPI):  # pragma: no cover - startup path
         # Re-raise to fail fast
         raise
 
-    # Restore info level logging after startup (suppressed during uvicorn init in dev mode)
+    # Restore uvicorn info logging (suppressed during startup in dev mode)
+    # Keep uvicorn.error at WARNING to suppress WebSocket "connection open/closed" messages
     if frontend.devmode:
         logging.getLogger("uvicorn").setLevel(logging.INFO)
-        logging.getLogger("uvicorn.access").setLevel(logging.INFO)
+    logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
 
     await frontend.load()
     await start_background()
@@ -67,6 +74,9 @@ app = FastAPI(
     openapi_url=None,
 )
 
+# Custom access logging (uvicorn's access_log is disabled)
+app.add_middleware(AccessLogMiddleware)
+
 # Apply redirections to auth-host if configured (deny access to restricted endpoints, remove /auth/)
 app.middleware("http")(auth_host.redirect_middleware)
 

paskia/fastapi/remote.py

@@ -324,7 +324,7 @@ async def websocket_remote_auth_permit(ws: WebSocket):
                     token_str = passphrase.generate()
                     expiry = expires()
                     db.create_reset_token(
-                        user_uuid=cred.user,
+                        user_uuid=cred.user_uuid,
                         passphrase=token_str,
                         expiry=expiry,
                         token_type="device addition",
@@ -333,7 +333,7 @@ async def websocket_remote_auth_permit(ws: WebSocket):
                     # Also create a session so the device is logged in
                     normalized_host = hostutil.normalize_host(request.host)
                     session_token = db.login(
-                        user_uuid=cred.user,
+                        user_uuid=cred.user_uuid,
                         credential_uuid=cred.uuid,
                         sign_count=new_sign_count,
                         host=normalized_host,
@@ -346,7 +346,7 @@ async def websocket_remote_auth_permit(ws: WebSocket):
 
                     normalized_host = hostutil.normalize_host(request.host)
                     session_token = db.login(
-                        user_uuid=cred.user,
+                        user_uuid=cred.user_uuid,
                         credential_uuid=cred.uuid,
                         sign_count=new_sign_count,
                         host=normalized_host,
@@ -359,7 +359,7 @@ async def websocket_remote_auth_permit(ws: WebSocket):
                 completed = await remoteauth.instance.complete_request(
                     token=request.key,
                     session_token=session_token,
-                    user_uuid=cred.user,
+                    user_uuid=cred.user_uuid,
                     credential_uuid=cred.uuid,
                     reset_token=reset_token,
                 )

paskia/fastapi/reset.py

@@ -10,8 +10,6 @@ display name. If multiple users match, they are listed and the command
 aborts. A new one-time reset link is always created.
 """
 
-from __future__ import annotations
-
 import asyncio
 from uuid import UUID
 

paskia/fastapi/response.py

@@ -0,0 +1,22 @@
+"""FastAPI response utilities for msgspec.Struct serialization."""
+
+import msgspec
+from fastapi import Response
+
+
+class MsgspecResponse(Response):
+    """Response that uses msgspec for JSON encoding.
+
+    Use this for returning msgspec.Struct, dict, or list with proper serialization.
+    """
+
+    media_type = "application/json"
+
+    def __init__(
+        self,
+        content: msgspec.Struct | dict | list,
+        status_code: int = 200,
+        headers: dict | None = None,
+    ):
+        body = msgspec.json.encode(content)
+        super().__init__(content=body, status_code=status_code, headers=headers)

paskia/fastapi/user.py

@@ -1,4 +1,4 @@
-from datetime import timezone
+from datetime import UTC
 from uuid import UUID
 
 from fastapi import (
@@ -43,7 +43,7 @@ async def user_update_display_name(
             status_code=401, detail="Authentication Required", mode="login"
         )
     host = request.headers.get("host")
-    ctx = db.get_session_context(auth, host)
+    ctx = db.data().session_ctx(auth, host)
     if not ctx:
         raise authz.AuthException(
             status_code=401, detail="Session expired", mode="login"
@@ -62,7 +62,7 @@ async def api_logout_all(request: Request, response: Response, auth=AUTH_COOKIE)
     if not auth:
         return {"message": "Already logged out"}
     host = request.headers.get("host")
-    ctx = db.get_session_context(auth, host)
+    ctx = db.data().session_ctx(auth, host)
     if not ctx:
         raise authz.AuthException(
             status_code=401, detail="Session expired", mode="login"
@@ -84,14 +84,14 @@ async def api_delete_session(
             status_code=401, detail="Authentication Required", mode="login"
         )
     host = request.headers.get("host")
-    ctx = db.get_session_context(auth, host)
+    ctx = db.data().session_ctx(auth, host)
     if not ctx:
         raise authz.AuthException(
             status_code=401, detail="Session expired", mode="login"
         )
 
     target_session = db.data().sessions.get(session_id)
-    if not target_session or target_session.user != ctx.user.uuid:
+    if not target_session or target_session.user_uuid != ctx.user.uuid:
         raise HTTPException(status_code=404, detail="Session not found")
 
     db.delete_session(session_id, ctx=ctx)
@@ -141,8 +141,8 @@ async def api_create_link(
         "message": "Registration link generated successfully",
         "url": url,
         "expires": (
-            expiry.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
+            expiry.astimezone(UTC).isoformat().replace("+00:00", "Z")
             if expiry.tzinfo
-            else expiry.replace(tzinfo=timezone.utc).isoformat().replace("+00:00", "Z")
+            else expiry.replace(tzinfo=UTC).isoformat().replace("+00:00", "Z")
         ),
     }

paskia/fastapi/ws.py

@@ -38,11 +38,11 @@ async def websocket_register_add(
                 f"The reset link for {passkey.instance.rp_name} is invalid or has expired"
             )
         s = get_reset(reset)
-        user_uuid = s.user
+        user_uuid = s.user_uuid
     else:
         # Require recent authentication for adding a new passkey
         ctx = await authz.verify(auth, perm=[], host=host, max_age="5m")
-        user_uuid = ctx.session.user
+        user_uuid = ctx.session.user_uuid
         s = ctx.session
 
     # Get user information and determine effective user_name for this registration
@@ -91,7 +91,7 @@ async def websocket_authenticate(ws: WebSocket, auth=AUTH_COOKIE):
     session_user_uuid = None
     credential_ids = None
     if auth:
-        ctx = db.get_session_context(auth, host)
+        ctx = db.data().session_ctx(auth, host)
         if ctx:
             session_user_uuid = ctx.user.uuid
             credential_ids = db.get_user_credential_ids(session_user_uuid) or None
@@ -99,7 +99,7 @@ async def websocket_authenticate(ws: WebSocket, auth=AUTH_COOKIE):
     cred, new_sign_count = await authenticate_chat(ws, origin, credential_ids)
 
     # If reauth mode, verify the credential belongs to the session's user
-    if session_user_uuid and cred.user != session_user_uuid:
+    if session_user_uuid and cred.user_uuid != session_user_uuid:
         raise ValueError("This passkey belongs to a different account")
 
     # Create session and update user/credential in a single transaction
@@ -114,7 +114,7 @@ async def websocket_authenticate(ws: WebSocket, auth=AUTH_COOKIE):
         raise ValueError(f"Host must be the same as or a subdomain of {rp_id}")
 
     token = db.login(
-        user_uuid=cred.user,
+        user_uuid=cred.user_uuid,
         credential_uuid=cred.uuid,
         sign_count=new_sign_count,
         host=normalized_host,
@@ -125,7 +125,7 @@ async def websocket_authenticate(ws: WebSocket, auth=AUTH_COOKIE):
 
     await ws.send_json(
         {
-            "user": str(cred.user),
+            "user": str(cred.user_uuid),
             "session_token": token,
         }
     )

paskia/fastapi/wsutil.py

@@ -3,6 +3,7 @@ Shared WebSocket utilities for FastAPI endpoints.
 """
 
 import logging
+import time
 from functools import wraps
 
 import base64url
@@ -10,6 +11,7 @@ from fastapi import WebSocket, WebSocketDisconnect
 from webauthn.helpers.exceptions import InvalidAuthenticationResponse
 
 from paskia.fastapi import authz
+from paskia.fastapi.logging import log_ws_close, log_ws_open
 from paskia.globals import passkey
 from paskia.util import pow
 
@@ -19,11 +21,19 @@ def websocket_error_handler(func):
 
     @wraps(func)
     async def wrapper(ws: WebSocket, *args, **kwargs):
+        client = ws.client.host if ws.client else "-"
+        host = ws.headers.get("host", "-")
+        path = ws.url.path
+
+        start = time.perf_counter()
+        ws_id = log_ws_open(client, host, path)
+        close_code = None
+
         try:
             await ws.accept()
             return await func(ws, *args, **kwargs)
-        except WebSocketDisconnect:
-            pass
+        except WebSocketDisconnect as e:
+            close_code = e.code
         except authz.AuthException as e:
             await ws.send_json(
                 {
@@ -36,6 +46,9 @@ def websocket_error_handler(func):
         except Exception:
             logging.exception("Internal Server Error")
             await ws.send_json({"status": 500, "detail": "Internal Server Error"})
+        finally:
+            duration_ms = (time.perf_counter() - start) * 1000
+            log_ws_close(client, ws_id, close_code, duration_ms)
 
     return wrapper
 

paskia/migrate/__init__.py

@@ -14,7 +14,7 @@ Or via the CLI entry point (if installed):
 import argparse
 import asyncio
 import re
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from uuid import UUID
 
 import base64url
@@ -154,7 +154,7 @@ async def migrate_from_sql(
                 if perm_uuid:
                     new_permissions[perm_uuid] = True
             new_role = Role(
-                org=role.org_uuid,
+                org_uuid=role.org_uuid,
                 display_name=role.display_name,
                 permissions=new_permissions,
             )
@@ -172,8 +172,8 @@ async def migrate_from_sql(
             user_key: UUID = legacy_user.uuid
             new_user = User(
                 display_name=legacy_user.display_name,
-                role=legacy_user.role_uuid,
-                created_at=legacy_user.created_at or datetime.now(timezone.utc),
+                role_uuid=legacy_user.role_uuid,
+                created_at=legacy_user.created_at or datetime.now(UTC),
                 last_seen=legacy_user.last_seen,
                 visits=legacy_user.visits,
             )
@@ -190,7 +190,7 @@ async def migrate_from_sql(
             cred_key: UUID = legacy_cred.uuid
             new_cred = Credential(
                 credential_id=legacy_cred.credential_id,
-                user=legacy_cred.user_uuid,
+                user_uuid=legacy_cred.user_uuid,
                 aaguid=legacy_cred.aaguid,
                 public_key=legacy_cred.public_key,
                 sign_count=legacy_cred.sign_count,
@@ -217,8 +217,8 @@ async def migrate_from_sql(
                 # Already in new format or unknown - try to use as-is
                 session_key = base64url.enc(old_key[:12])
             db.sessions[session_key] = Session(
-                user=sess.user_uuid,
-                credential=sess.credential_uuid,
+                user_uuid=sess.user_uuid,
+                credential_uuid=sess.credential_uuid,
                 host=sess.host,
                 ip=sess.ip,
                 user_agent=sess.user_agent,
@@ -241,14 +241,14 @@ async def migrate_from_sql(
                 # Already in new format or unknown - truncate to 9 bytes
                 token_key = old_key[:9]
             db.reset_tokens[token_key] = ResetToken(
-                user=token.user_uuid,
+                user_uuid=token.user_uuid,
                 expiry=token.expiry,
                 token_type=token.token_type,
             )
         print(f"  Migrated {len(token_models)} reset tokens")
 
     # Queue and flush all changes using the transaction mechanism
-    with db.transaction("migrate"):
+    with db.transaction("migrate:sql"):
         pass  # All data already added to _data, transaction commits on exit
 
     await store.flush()

paskia/migrate/sql.py

@@ -9,7 +9,7 @@ DO NOT use this module for new code. Use paskia.db instead.
 
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from uuid import UUID
 
 from sqlalchemy import (
@@ -25,11 +25,6 @@ from sqlalchemy.dialects.sqlite import BLOB
 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
-from paskia.db import (
-    Org,
-    Role,
-)
-
 
 # Legacy User class for SQL schema (uses 'role_uuid' not 'role')
 @dataclass
@@ -71,6 +66,17 @@ class _LegacyRole:
     permissions: list[str] | None = None
 
 
+# Legacy Org class for SQL schema (has mutable permissions/roles lists)
+@dataclass
+class _LegacyOrg:
+    """Org as stored in the old SQL schema with mutable permissions/roles."""
+
+    uuid: UUID
+    display_name: str
+    permissions: list[str] | None = None
+    roles: list[_LegacyRole] | None = None
+
+
 # Legacy Session class for SQL schema (uses 'key' as field, 'user_uuid', 'credential_uuid')
 @dataclass
 class _LegacySession:
@@ -112,8 +118,8 @@ def _normalize_dt(value: datetime | None) -> datetime | None:
     if value is None:
         return None
     if value.tzinfo is None:
-        return value.replace(tzinfo=timezone.utc)
-    return value.astimezone(timezone.utc)
+        return value.replace(tzinfo=UTC)
+    return value.astimezone(UTC)
 
 
 class Base(DeclarativeBase):
@@ -128,12 +134,13 @@ class OrgModel(Base):
 
     def as_dataclass(self):
         # Base Org without permissions/roles (filled by data accessors)
-        org = Org(display_name=self.display_name)
-        org.uuid = UUID(bytes=self.uuid)
-        return org
+        return _LegacyOrg(
+            uuid=UUID(bytes=self.uuid),
+            display_name=self.display_name,
+        )
 
     @staticmethod
-    def from_dataclass(org: Org):
+    def from_dataclass(org: _LegacyOrg):
         return OrgModel(uuid=org.uuid.bytes, display_name=org.display_name)
 
 
@@ -172,7 +179,7 @@ class UserModel(Base):
         LargeBinary(16), ForeignKey("roles.uuid", ondelete="CASCADE"), nullable=False
     )
     created_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+        DateTime(timezone=True), default=lambda: datetime.now(UTC)
     )
     last_seen: Mapped[datetime | None] = mapped_column(
         DateTime(timezone=True), nullable=True
@@ -195,7 +202,7 @@ class UserModel(Base):
             uuid=user.uuid.bytes,
             display_name=user.display_name,
             role_uuid=user.role_uuid.bytes,
-            created_at=user.created_at or datetime.now(timezone.utc),
+            created_at=user.created_at or datetime.now(UTC),
             last_seen=user.last_seen,
             visits=user.visits,
         )
@@ -215,7 +222,7 @@ class CredentialModel(Base):
     public_key: Mapped[bytes] = mapped_column(BLOB, nullable=False)
     sign_count: Mapped[int] = mapped_column(Integer, nullable=False)
     created_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+        DateTime(timezone=True), default=lambda: datetime.now(UTC)
     )
     last_used: Mapped[datetime | None] = mapped_column(
         DateTime(timezone=True), nullable=True
@@ -255,7 +262,7 @@ class SessionModel(Base):
     user_agent: Mapped[str] = mapped_column(String(512), nullable=False)
     renewed: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
-        default=lambda: datetime.now(timezone.utc),
+        default=lambda: datetime.now(UTC),
         nullable=False,
     )
 
@@ -388,7 +395,7 @@ class DB:
             result = await session.execute(select(PermissionModel))
             return [p.as_dataclass() for p in result.scalars().all()]
 
-    async def list_organizations(self) -> list[Org]:
+    async def list_organizations(self) -> list[_LegacyOrg]:
         async with self.session() as session:
             # Load all orgs
             orgs_result = await session.execute(select(OrgModel))
@@ -415,13 +422,13 @@ class DB:
                 perms_by_role.setdefault(rp.role_uuid, []).append(rp.permission_id)
 
             # Build org dataclasses with roles and permission IDs
-            roles_by_org: dict[bytes, list[Role]] = {}
+            roles_by_org: dict[bytes, list[_LegacyRole]] = {}
             for rm in role_models:
                 r_dc = rm.as_dataclass()
                 r_dc.permissions = perms_by_role.get(rm.uuid, [])
                 roles_by_org.setdefault(rm.org_uuid, []).append(r_dc)
 
-            orgs: list[Org] = []
+            orgs: list[_LegacyOrg] = []
             for om in org_models:
                 o_dc = om.as_dataclass()
                 o_dc.permissions = perms_by_org.get(om.uuid, [])