AAI-302 update registration code to create database records

auth0/client.py

@@ -8,7 +8,7 @@
 
 from auth.management import get_management_token
 from config import Settings, get_settings
-from schemas.biocommons import Auth0UserData
+from schemas.biocommons import Auth0UserData, BiocommonsRegisterData
 
 
 class RoleData(BaseModel):
@@ -94,6 +94,12 @@ def get_user(self, user_id: str) -> Auth0UserData:
         resp = self._client.get(url)
         return Auth0UserData(**resp.json())
 
+    def create_user(self, user: BiocommonsRegisterData) -> Auth0UserData:
+        url = f"https://{self.domain}/api/v2/users"
+        resp = self._client.post(url, json=user.model_dump(mode="json"))
+        resp.raise_for_status()
+        return Auth0UserData(**resp.json())
+
     def add_roles_to_user(self, user_id: str, role_id: str | list[str]):
         """
         Add one or more roles to a user. The role(s) must already exist.

db/admin.py

@@ -16,7 +16,7 @@
     GroupMembership,
     GroupMembershipHistory,
 )
-from db.setup import engine
+from db.setup import get_engine
 
 
 def setup_oauth():
@@ -114,7 +114,7 @@ def __init__(self, app: FastAPI, secret_key: str):
         self.auth0_client = setup_oauth()
         self.admin = Admin(
             app,
-            engine=engine,
+            engine=get_engine(),
             base_url="/db_admin",
             authentication_backend=AdminAuth(secret_key=secret_key, auth0_client=self.auth0_client),
             title="AAI Backend Admin"

db/models.py

@@ -10,6 +10,7 @@
 
 from auth0.client import Auth0Client
 from db.core import BaseModel
+from schemas.biocommons import Auth0UserData
 from schemas.user import SessionUser
 
 
@@ -44,14 +45,23 @@ class BiocommonsUser(BaseModel, table=True):
     )
 
     @classmethod
-    def create_from_auth0(cls, auth0_id: str, auth0_client: Auth0Client):
+    def create_from_auth0(cls, auth0_id: str, auth0_client: Auth0Client) -> Self:
+        """
+        Get user data from Auth0 API and create a new BiocommonsUser object.
+        """
         user_data = auth0_client.get_user(user_id=auth0_id)
-        user = cls(
-            id=auth0_id,
-            email=user_data.email,
-            username=user_data.username
+        return cls.from_auth0_data(user_data)
+
+    @classmethod
+    def from_auth0_data(cls, data: Auth0UserData) -> Self:
+        """
+        Create a new BiocommonsUser object from Auth0 user data (no API call).
+        """
+        return cls(
+            id=data.user_id,
+            email=data.email,
+            username=data.username
         )
-        return user
 
     @classmethod
     def get_or_create(cls, auth0_id: str, db_session: Session, auth0_client: Auth0Client) -> Self:
@@ -65,6 +75,17 @@ def get_or_create(cls, auth0_id: str, db_session: Session, auth0_client: Auth0Cl
             db_session.commit()
         return user
 
+    def add_platform_membership(self, platform: PlatformEnum, db_session: Session, auto_approve: bool = False) -> "PlatformMembership":
+        membership = PlatformMembership(
+            platform_id=platform,
+            user=self,
+            approval_status=ApprovalStatusEnum.APPROVED if auto_approve else ApprovalStatusEnum.PENDING,
+            updated_by=None,
+        )
+        db_session.add(membership)
+        membership.save_history(db_session)
+        return membership
+
 
 class PlatformMembership(BaseModel, table=True):
     __table_args__ = (
@@ -81,6 +102,17 @@ class PlatformMembership(BaseModel, table=True):
     updated_by_id: str | None = Field(foreign_key="biocommons_user.id", nullable=True)
     updated_by: "BiocommonsUser" = Relationship(sa_relationship_kwargs={"foreign_keys": "PlatformMembership.updated_by_id",})
 
+    def save_history(self, session: Session) -> 'PlatformMembershipHistory':
+        history = PlatformMembershipHistory(
+            platform_id=self.platform_id,
+            user=self.user,
+            approval_status=self.approval_status,
+            updated_at=self.updated_at,
+            updated_by=self.updated_by,
+        )
+        session.add(history)
+        return history
+
 
 
 

db/setup.py

@@ -9,6 +9,17 @@
 
 log = logging.getLogger('uvicorn.error')
 
+# Set engine as None initially so it's not created on import
+_engine = None
+
+
+def get_engine():
+    global _engine
+    if _engine is None:
+        db_url, db_connect_args = get_db_config()
+        _engine = create_engine(db_url, connect_args=db_connect_args)
+    return _engine
+
 
 def get_db_config() -> Tuple[str, dict]:
     """
@@ -40,19 +51,17 @@ def get_db_config() -> Tuple[str, dict]:
     return db_url, connect_args
 
 
-DB_URL, db_connect_args = get_db_config()
-engine = create_engine(DB_URL, connect_args=db_connect_args)
-
-
 def create_db_and_tables():
     # NOTE: we only do this in dev (with sqlite).
     # For production, we manage the DB schema with alembic
     db_url, connect_args = get_db_config()
     if db_url.startswith("sqlite://"):
+        engine = get_engine()
         log.info("Automatically creating DB tables for sqlite")
         BaseModel.metadata.create_all(engine)
 
 
 def get_db_session():
+    engine = get_engine()
     with Session(engine) as session:
         yield session

main.py

@@ -8,7 +8,6 @@
 # This has to be imported even if unused
 from db import models  # noqa: F401
 from db.admin import DatabaseAdmin
-from db.setup import create_db_and_tables
 from routers import admin, biocommons_groups, bpa_register, galaxy_register, user, utils
 
 # Load .env to get CORS_ALLOWED_ORIGINS.
@@ -26,6 +25,7 @@
 async def lifespan(app: FastAPI):
     # NOTE: we only create the database and tables automatically in development:
     # we assume that if the DB is an sqlite DB, we are in dev.
+    from db.setup import create_db_and_tables
     create_db_and_tables()
     DatabaseAdmin.setup(app=app, secret_key=SECRET_KEY)
     yield

routers/bpa_register.py

@@ -1,16 +1,22 @@
+import logging
 from datetime import datetime, timezone
 from typing import Any, Dict
 
 from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
-from httpx import AsyncClient
+from httpx import HTTPStatusError
+from sqlmodel import Session
 
-from auth.management import get_management_token
 from auth.ses import EmailService
+from auth0.client import Auth0Client, get_auth0_client
 from config import Settings, get_settings
-from schemas.biocommons import BiocommonsRegisterData
+from db.models import BiocommonsUser, PlatformEnum
+from db.setup import get_db_session
+from schemas.biocommons import Auth0UserData, BiocommonsRegisterData
 from schemas.bpa import BPARegistrationRequest
 from schemas.service import Resource, Service
 
+logger = logging.getLogger(__name__)
+
 router = APIRouter(prefix="/bpa", tags=["bpa", "registration"])
 
 
@@ -34,28 +40,7 @@ def send_approval_email(registration: BPARegistrationRequest, bpa_resources: lis
     email_service.send(approver_email, subject, body_html)
 
 
-@router.post(
-    "/register",
-    response_model=Dict[str, Any],
-    responses={
-        400: {"description": "Bad Request - Validation error"},
-        409: {"description": "Conflict - User already exists"},
-        500: {"description": "Internal server error"},
-    },
-)
-async def register_bpa_user(
-    registration: BPARegistrationRequest,
-    background_tasks: BackgroundTasks,
-    settings: Settings = Depends(get_settings)
-) -> Dict[str, Any]:
-    """Register a new BPA user with selected organization resources."""
-    url = f"https://{settings.auth0_domain}/api/v2/users"
-    token = get_management_token(settings=settings)
-    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
-
-    now = datetime.now(timezone.utc)
-
-    # Create BPA resources
+def _get_bpa_resources(registration: BPARegistrationRequest, settings: Settings, update_time: datetime) -> list[Resource]:
     bpa_resources = []
     for org_id, is_selected in registration.organizations.items():
         if not is_selected:
@@ -68,13 +53,35 @@ async def register_bpa_user(
             id=org_id,
             name=settings.organizations[org_id],
             status="pending",
-            last_updated=now,
-            initial_request_time=now,
+            last_updated=update_time,
+            initial_request_time=update_time,
             updated_by="system",
         ).model_dump(mode="json")
         bpa_resources.append(resource)
+    return bpa_resources
+
 
-    # Create BPA service
+
+@router.post(
+    "/register",
+    response_model=Dict[str, Any],
+    responses={
+        400: {"description": "Bad Request - Validation error"},
+        409: {"description": "Conflict - User already exists"},
+        500: {"description": "Internal server error"},
+    },
+)
+async def register_bpa_user(
+    registration: BPARegistrationRequest,
+    background_tasks: BackgroundTasks,
+    settings: Settings = Depends(get_settings),
+    db_session: Session = Depends(get_db_session),
+    auth0_client: Auth0Client = Depends(get_auth0_client)
+) -> Dict[str, Any]:
+    """Register a new BPA user with selected organization resources."""
+    now = datetime.now(timezone.utc)
+
+    bpa_resources = _get_bpa_resources(registration, settings, update_time=now)
     bpa_service = Service(
         name="Bioplatforms Australia Data Portal",
         id="bpa",
@@ -91,24 +98,33 @@ async def register_bpa_user(
     )
 
     try:
-        async with AsyncClient() as client:
-            response = await client.post(
-                url, headers=headers, json=user_data.model_dump(mode="json")
-            )
-            if response.status_code != 201:
-                raise HTTPException(
-                    status_code=400,
-                    detail=f"Registration failed: {response.json()['message']}",
-                )
+        logger.info("Registering user with Auth0")
+        auth0_user_data = auth0_client.create_user(user_data)
+
+        logger.info("Adding user to DB")
+        _create_bpa_user_record(auth0_user_data, db_session)
 
         if bpa_resources and settings.send_email:
             background_tasks.add_task(send_approval_email, registration, bpa_resources)
 
-        return {"message": "User registered successfully", "user": response.json()}
+        return {"message": "User registered successfully", "user": auth0_user_data.model_dump(mode="json")}
 
-    except HTTPException:
-        raise
+    except HTTPStatusError as e:
+        raise HTTPException(status_code=e.response.status_code, detail=e.response.text)
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"Failed to register user: {str(e)}"
         )
+
+
+def _create_bpa_user_record(auth0_user_data: Auth0UserData, session: Session) -> BiocommonsUser:
+    db_user = BiocommonsUser.from_auth0_data(data=auth0_user_data)
+    bpa_membership = db_user.add_platform_membership(
+        platform=PlatformEnum.BPA_DATA_PORTAL,
+        db_session=session,
+        auto_approve=True
+    )
+    session.add(db_user)
+    session.add(bpa_membership)
+    session.commit()
+    return db_user

routers/galaxy_register.py

@@ -4,12 +4,15 @@
 import httpx
 from fastapi import APIRouter, Header, HTTPException
 from fastapi.params import Depends
+from sqlmodel import Session
 
-from auth.management import get_management_token
+from auth0.client import Auth0Client, get_auth0_client
 from config import Settings, get_settings
+from db.models import BiocommonsUser, PlatformEnum
+from db.setup import get_db_session
 from galaxy.client import GalaxyClient, get_galaxy_client
 from register.tokens import create_registration_token, verify_registration_token
-from schemas.biocommons import BiocommonsRegisterData
+from schemas.biocommons import Auth0UserData, BiocommonsRegisterData
 from schemas.galaxy import GalaxyRegistrationData
 
 logger = logging.getLogger(__name__)
@@ -29,6 +32,8 @@ def register(
         registration_data: GalaxyRegistrationData,
         settings: Annotated[Settings, Depends(get_settings)],
         galaxy_client: Annotated[GalaxyClient, Depends(get_galaxy_client)],
+        auth0_client: Annotated[Auth0Client, Depends(get_auth0_client)],
+        db_session: Annotated[Session, Depends(get_db_session)],
         registration_token: Optional[str] = Header(None),
 ):
     if not registration_token:
@@ -47,21 +52,25 @@ def register(
     except httpx.HTTPError as e:
         logger.warning(f"Failed to check username in Galaxy: {e}")
 
-    url = f"https://{settings.auth0_domain}/api/v2/users"
-    logger.debug("Getting management token.")
-    management_token = get_management_token(settings=settings)
-    headers = {"Authorization": f"Bearer {management_token}"}
-    logger.debug("Registering with Auth0 management API")
-    resp = httpx.post(
-        url,
-        # Use exclude_none so we don't include username/name fields
-        #   when not specified, Auth0 doesn't like this
-        json=user_data.model_dump(
-            mode="json",
-            exclude_none=True
-        ),
-        headers=headers
+    try:
+        logger.info("Registering user with Auth0")
+        auth0_user_data = auth0_client.create_user(user_data)
+    except httpx.HTTPStatusError as e:
+        raise HTTPException(status_code=e.response.status_code, detail=f'Registration failed: {e}')
+    # Add to database and record Galaxy membership
+    logger.info("Adding user to DB")
+    _create_galaxy_user_record(auth0_user_data, db_session)
+    return {"message": "User registered successfully", "user": auth0_user_data.model_dump(mode="json")}
+
+
+def _create_galaxy_user_record(auth0_user_data: Auth0UserData, session: Session) -> BiocommonsUser:
+    db_user = BiocommonsUser.from_auth0_data(data=auth0_user_data)
+    galaxy_membership = db_user.add_platform_membership(
+        platform=PlatformEnum.GALAXY,
+        db_session=session,
+        auto_approve=True
     )
-    if resp.status_code != 201:
-        raise HTTPException(status_code=400, detail=f'Registration failed: {resp.json()["message"]}')
-    return {"message": "User registered successfully", "user": resp.json()}
+    session.add(db_user)
+    session.add(galaxy_membership)
+    session.commit()
+    return db_user