diff --git a/alphatrion/server/graphql/resolvers.py b/alphatrion/server/graphql/resolvers.py index 244664a..2984c78 100644 --- a/alphatrion/server/graphql/resolvers.py +++ b/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, diff --git a/alphatrion/server/graphql/types.py b/alphatrion/server/graphql/types.py index 03d35d8..60c19f0 100644 --- a/alphatrion/server/graphql/types.py +++ b/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 diff --git a/alphatrion/storage/sqlstore.py b/alphatrion/storage/sqlstore.py index 1aeefc0..481bc06 100644 --- a/alphatrion/storage/sqlstore.py +++ b/alphatrion/storage/sqlstore.py @@ -31,8 +31,15 @@ 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( @@ -40,6 +47,9 @@ def create_team( 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 diff --git a/tests/integration/server/test_graphql_mutation.py b/tests/integration/server/test_graphql_mutation.py index 3b51d97..ecf7269 100644 --- a/tests/integration/server/test_graphql_mutation.py +++ b/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)