Support create user and team with uuid to integrate with external account systems

alphatrion/server/graphql/resolvers.py

@@ -294,6 +294,7 @@ class GraphQLMutations:
     def create_user(input: CreateUserInput) -> User:
         metadb = runtime.graphql_runtime().metadb
         user_id = metadb.create_user(
+            uuid=uuid.UUID(input.id) if input.id else None,
             username=input.username,
             email=input.email,
             avatar_url=input.avatar_url,
@@ -337,6 +338,7 @@ def update_user(input: UpdateUserInput) -> User:
     def create_team(input: CreateTeamInput) -> Team:
         metadb = runtime.graphql_runtime().metadb
         team_id = metadb.create_team(
+            uuid=uuid.UUID(input.id) if input.id else None,
             name=input.name,
             description=input.description,
             meta=input.meta,

alphatrion/server/graphql/types.py

@@ -145,6 +145,7 @@ class Metric:
 # Input types for mutations
 @strawberry.input
 class CreateUserInput:
+    id: strawberry.ID | None = None
     username: str
     email: str
     avatar_url: str | None = None
@@ -153,6 +154,7 @@ class CreateUserInput:
 
 @strawberry.input
 class CreateTeamInput:
+    id: strawberry.ID | None = None
     name: str
     description: str | None = None
     meta: JSON | None = None

alphatrion/storage/sqlstore.py

@@ -31,15 +31,25 @@ def __init__(self, db_url: str, init_tables: bool = False):
 
     # ---------- Team APIs ----------
 
+    # If uuid is provided, we will use the provided uuid for the new team.
+    # This is useful for binding with external team management system where
+    # the team id is already determined.
     def create_team(
-        self, name: str, description: str | None = None, meta: dict | None = None
+        self,
+        name: str,
+        uuid: uuid.UUID | None = None,
+        description: str | None = None,
+        meta: dict | None = None,
     ) -> uuid.UUID:
         session = self._session()
         new_team = Team(
             name=name,
             description=description,
             meta=meta,
         )
+        if uuid is not None:
+            new_team.uuid = uuid
+
         session.add(new_team)
         session.commit()
         team_id = new_team.uuid
@@ -71,27 +81,34 @@ def list_user_teams(self, user_id: uuid.UUID) -> list[Team]:
 
     # ---------- User APIs ----------
 
+    # If uuid is provided, we will use the provided uuid for the new user.
+    # This is useful for binding with external user management system where
+    # the user id is already determined.
     def create_user(
         self,
         username: str,
         email: str,
+        uuid: uuid.UUID | None = None,
         avatar_url: str | None = None,
         team_id: uuid.UUID | None = None,
         meta: dict | None = None,
     ) -> uuid.UUID:
+        user = User(
+            username=username,
+            email=email,
+            avatar_url=avatar_url,
+            meta=meta,
+        )
+        if uuid is not None:
+            user.uuid = uuid
+
         # If team_id is not provided, we will just create the user
         # without any team association.
         if team_id is None:
             session = self._session()
-            new_user = User(
-                username=username,
-                email=email,
-                avatar_url=avatar_url,
-                meta=meta,
-            )
-            session.add(new_user)
+            session.add(user)
             session.commit()
-            user_id = new_user.uuid
+            user_id = user.uuid
             session.close()
 
             return user_id
@@ -100,23 +117,17 @@ def create_user(
             # add to the team in a transaction.
             session = self._session()
             try:
-                new_user = User(
-                    username=username,
-                    email=email,
-                    avatar_url=avatar_url,
-                    meta=meta,
-                )
-                session.add(new_user)
+                session.add(user)
                 session.flush()  # flush to get the new user's id
 
                 new_member = TeamMember(
-                    user_id=new_user.uuid,
+                    user_id=user.uuid,
                     team_id=team_id,
                 )
                 session.add(new_member)
 
                 session.commit()
-                user_id = new_user.uuid
+                user_id = user.uuid
             except Exception as e:
                 session.rollback()
                 raise e

tests/integration/server/test_graphql_mutation.py

@@ -61,6 +61,41 @@ def test_create_team_mutation():
     assert team.name == "Test Team"
 
 
+def test_create_team_mutation_with_uuid():
+    """Test creating a team via GraphQL mutation"""
+    init(init_tables=True)
+    id = uuid.uuid4()  # Generate a UUID to use for the new team
+
+    mutation = f"""
+    mutation {{
+        createTeam(input: {{
+            id: "{str(id)}"
+            name: "Test Team"
+            description: "A team created via mutation"
+            meta: {{foo: "bar", count: 42}}
+        }}) {{
+            id
+            name
+            description
+            meta
+            createdAt
+            updatedAt
+            totalProjects
+            totalExperiments
+            totalRuns
+        }}
+    }}
+    """
+    response = schema.execute_sync(
+        mutation,
+        variable_values={},
+    )
+    assert response.errors is None
+    assert response.data["createTeam"]["name"] == "Test Team"
+    # Verify team was actually created in database
+    assert response.data["createTeam"]["id"] == str(id)  # Verify the returned ID matches the provided UUID
+
+
 def test_create_user_mutation():
     """Test creating a user via GraphQL mutation"""
     init(init_tables=True)
@@ -109,6 +144,45 @@ def test_create_user_mutation():
     assert user.username == username
 
 
+def test_create_user_mutation_with_uuid():
+    """Test creating a user via GraphQL mutation"""
+    init(init_tables=True)
+    id = uuid.uuid4()  # Generate a UUID to use for the new user
+
+    username = unique_username("testuser")
+    email = unique_email("testuser")
+
+    mutation = f"""
+    mutation {{
+        createUser(input: {{
+            id: "{str(id)}"
+            username: "{username}"
+            email: "{email}"
+            meta: {{role: "engineer", level: "senior"}}
+        }}) {{
+            id
+            username
+            email
+            meta
+            createdAt
+            updatedAt
+            teams {{
+                id
+                name
+            }}
+        }}
+    }}
+    """
+    response = schema.execute_sync(
+        mutation,
+        variable_values={},
+    )
+    assert response.errors is None
+    assert response.data["createUser"]["id"] == str(
+        id
+    )  # Verify the returned ID matches the provided UUID
+
+
 def test_add_user_to_team_mutation():
     """Test adding a user to a team via mutation"""
     init(init_tables=True)