diff options
| -rw-r--r-- | generalresearch/managers/gr/team.py | 40 | ||||
| -rw-r--r-- | generalresearch/models/gr/team.py | 96 | ||||
| -rw-r--r-- | pyproject.toml | 2 | ||||
| -rw-r--r-- | tests/managers/gr/test_team.py | 18 |
4 files changed, 113 insertions, 43 deletions
diff --git a/generalresearch/managers/gr/team.py b/generalresearch/managers/gr/team.py index 781cf1f..5d9c6bb 100644 --- a/generalresearch/managers/gr/team.py +++ b/generalresearch/managers/gr/team.py @@ -16,6 +16,7 @@ from generalresearch.models.custom_types import UUIDStr from generalresearch.models.gr.team import ( Membership, MembershipPrivilege, + TeamUpdate, ) if TYPE_CHECKING: @@ -167,9 +168,9 @@ class TeamManager(PostgresManagerWithRedis): name: str, uuid: UUIDStr | None = None, ) -> Team: - from generalresearch.models.gr.team import Team + from generalresearch.models.gr.team import TeamCreate - team = Team.model_validate({"uuid": uuid or uuid4().hex, "name": name}) + team = TeamCreate.model_validate({"uuid": uuid or uuid4().hex, "name": name}) with self.pg_config.make_connection() as conn: with conn.cursor() as c: @@ -183,8 +184,37 @@ class TeamManager(PostgresManagerWithRedis): ) team_id = c.fetchone()["id"] # type: ignore conn.commit() - team.id = team_id + persisted_team = self.get_by_id(team_id=team_id) + assert persisted_team is not None + return persisted_team + + def update( + self, + team_id: PositiveInt, + update: TeamUpdate, + ) -> Team: + data = update.model_dump(exclude_unset=True) + + assignments = sql.SQL(", ").join( + sql.SQL("{} = %s").format(sql.Identifier(field)) for field in data + ) + + query = sql.SQL(""" + UPDATE common_team + SET {assignments} + WHERE id = %s + """).format(assignments=assignments) + + with self.pg_config.make_connection() as conn: + with conn.cursor() as cursor: + cursor.execute(query, [*data.values(), team_id]) + if cursor.rowcount != 1: + raise LookupError(f"Team {team_id} does not exist") + conn.commit() + team = self.get_by_id(team_id) + assert team is not None + self.redis_client.delete(team.cache_key) return team def add_user( @@ -210,8 +240,8 @@ class TeamManager(PostgresManagerWithRedis): with self.connection(conn) as conn, conn.cursor() as c: c.execute( query=sql.SQL(""" - INSERT INTO common_team_businesses - (team_id, business_id) + INSERT INTO common_team_businesses + (team_id, business_id) VALUES (%s, %s) """), params=( diff --git a/generalresearch/models/gr/team.py b/generalresearch/models/gr/team.py index 217f184..b53b913 100644 --- a/generalresearch/models/gr/team.py +++ b/generalresearch/models/gr/team.py @@ -5,7 +5,7 @@ import os from datetime import UTC, datetime from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from uuid import uuid4 import pandas as pd @@ -18,6 +18,7 @@ from pydantic import ( PositiveInt, ValidationError, field_validator, + model_validator, ) from pydantic.json_schema import SkipJsonSchema @@ -113,10 +114,33 @@ class Membership(BaseModel): self.team = tm.get_by_id(team_id=self.team_id) -class Team(BaseModel): - id: SkipJsonSchema[PositiveInt | None] = Field(default=None) +class TeamUpdate(BaseModel): + model_config = ConfigDict(extra="forbid") + + # This is nullable with a validator that says it can't be set to null, + # so that in the future if we had another editable field, that field + # could be updated separate from 'name' + name: str | None = Field( + min_length=3, max_length=255, examples=["Team ABC"], default=None + ) + + @field_validator("name") + @classmethod + def name_cannot_be_null(cls, value: str | None) -> str: + if value is None: + raise ValueError("name cannot be null") + return value + + @model_validator(mode="after") + def require_field(self) -> Self: + if not self.model_fields_set: + raise ValueError("Must provide a key to update") + return self + + +class TeamCreate(BaseModel): uuid: UUIDStrCoerce = Field(examples=[uuid4().hex]) - name: str = Field(max_length=255, examples=["Team ABC"]) + name: str = Field(min_length=3, max_length=255, examples=["Team ABC"]) # prefetch attributes memberships: SkipJsonSchema[list[Membership] | None] = Field(default=None) @@ -124,7 +148,38 @@ class Team(BaseModel): businesses: SkipJsonSchema[list[Business] | None] = Field(default=None) products: SkipJsonSchema[list[Product] | None] = Field(default=None) - # --- Prefetch Methods --- + @classmethod + def required_fields(cls) -> list[str]: + return [ + field_name + for field_name, field_info in cls.model_fields.items() + if field_info.is_required() + ] + + # --- Properties --- + @property + def cache_key(self) -> str: + return f"team:{self.uuid}" + + @property + def file_key(self) -> str: + return f"team-{self.uuid}" + + @property + def product_ids(self) -> list[UUIDStr] | None: + if self.products is None: + LOG.warning("prefetch not run") + return None + + return [p.uuid for p in self.products] + + @property + def product_uuids(self) -> list[UUIDStr] | None: + return self.product_ids + + +class Team(TeamCreate): + id: SkipJsonSchema[PositiveInt] = Field(exclude=True) def prefetch_memberships(self, gr_membership_manager: MembershipManager) -> None: self.memberships = gr_membership_manager.get_by_team_id(team_id=self.id) @@ -226,37 +281,6 @@ class Team(BaseModel): except (pa.ArrowException, OSError, ValueError) as e: raise OSError(f"Parquet verification failed: {e}") - @classmethod - def required_fields(cls) -> list[str]: - return [ - field_name - for field_name, field_info in cls.model_fields.items() - if field_info.is_required() - ] - - # --- Properties --- - @property - def cache_key(self) -> str: - return f"team:{self.uuid}" - - @property - def file_key(self) -> str: - return f"team-{self.uuid}" - - @property - def product_ids(self) -> list[UUIDStr] | None: - if self.products is None: - LOG.warning("prefetch not run") - return None - - return [p.uuid for p in self.products] - - @property - def product_uuids(self) -> list[UUIDStr] | None: - return self.product_ids - - # --- Methods --- - def set_cache( self, product_manager: ProductManager, diff --git a/pyproject.toml b/pyproject.toml index 523c92f..15d9118 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "generalresearch" -version = "3.8.5" +version = "3.8.6" description = "Python Utilities for General Research" readme = "README.md" requires-python = ">=3.14" diff --git a/tests/managers/gr/test_team.py b/tests/managers/gr/test_team.py index 878a9ca..689c4be 100644 --- a/tests/managers/gr/test_team.py +++ b/tests/managers/gr/test_team.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING from uuid import uuid4 from generalresearch.models.gr.authentication import GRUser -from generalresearch.models.gr.team import Membership, Team +from generalresearch.models.gr.team import Membership, Team, TeamUpdate if TYPE_CHECKING: from generalresearch.managers.gr.authentication import GRUserManager @@ -59,6 +59,22 @@ class TestTeamManager: assert isinstance(team, Team) assert isinstance(team.id, int) + def test_update_name( + self, gr_team_factory: Callable[..., Team], gr_team_manager: TeamManager + ): + team = gr_team_factory() + updated = gr_team_manager.update( + team_id=team.id, + update=TeamUpdate(name="Updated Team Name"), + ) + + assert updated.id == team.id + assert updated.uuid == team.uuid + assert updated.name == "Updated Team Name" + + persisted = gr_team_manager.get_by_id(team_id=team.id) + assert persisted == updated + def test_add_user( self, gr_team: Team, |