diff options
| -rw-r--r-- | generalresearch/gr_django/migrations/0003_user_email_alter_membership_owner_and_more.py | 32 | ||||
| -rw-r--r-- | generalresearch/gr_django/models.py | 15 | ||||
| -rw-r--r-- | generalresearch/managers/gr/authentication.py | 98 | ||||
| -rw-r--r-- | generalresearch/managers/gr/team.py | 180 | ||||
| -rw-r--r-- | generalresearch/models/gr/__init__.py | 7 | ||||
| -rw-r--r-- | generalresearch/models/gr/authentication.py | 15 | ||||
| -rw-r--r-- | generalresearch/models/gr/team.py | 25 | ||||
| -rw-r--r-- | pyproject.toml | 2 | ||||
| -rw-r--r-- | test_utils/models/gr/conftest.py | 2 |
9 files changed, 250 insertions, 126 deletions
diff --git a/generalresearch/gr_django/migrations/0003_user_email_alter_membership_owner_and_more.py b/generalresearch/gr_django/migrations/0003_user_email_alter_membership_owner_and_more.py new file mode 100644 index 0000000..a36979f --- /dev/null +++ b/generalresearch/gr_django/migrations/0003_user_email_alter_membership_owner_and_more.py @@ -0,0 +1,32 @@ +# Generated by Django 6.1.1 on 2026-10-16 19:20 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('common', '0002_business_balance_business_balance_net_7d_and_more'), + ] + + operations = [ + migrations.AddField( + model_name='user', + name='email', + field=models.EmailField(max_length=254, null=True, unique=True), + ), + migrations.AlterField( + model_name='membership', + name='owner', + field=models.BooleanField(default=False, null=True), + ), + migrations.AlterField( + model_name='membership', + name='privilege', + field=models.PositiveSmallIntegerField(choices=[(0, 'Admin'), (1, 'Member')], default=1), + ), + migrations.AddConstraint( + model_name='membership', + constraint=models.UniqueConstraint(fields=('user_id', 'team_id'), name='uniq_membership_user_team'), + ), + ] diff --git a/generalresearch/gr_django/models.py b/generalresearch/gr_django/models.py index 01bf749..7f26eaf 100644 --- a/generalresearch/gr_django/models.py +++ b/generalresearch/gr_django/models.py @@ -6,6 +6,7 @@ from django.db import models class User(models.Model): id = models.BigAutoField(primary_key=True) sub = models.CharField(max_length=200, blank=False, null=False) + email = models.EmailField(null=True, unique=True) is_superuser = models.BooleanField(default=False) date_joined = models.DateTimeField(auto_now_add=True) @@ -37,14 +38,22 @@ class Membership(models.Model): PRIVILEGES = ( (0, "Admin"), - (1, "Maintain"), - (2, "Read"), + (1, "Member"), ) privilege = models.PositiveSmallIntegerField(default=1, choices=PRIVILEGES) - owner = models.BooleanField(default=False) + + # todo: delete this field once managers are not querying for it + owner = models.BooleanField(default=False, null=True) created = models.DateTimeField(auto_now_add=True, null=False) + class Meta: + constraints = [ + models.UniqueConstraint( + fields=['user_id', 'team_id'], + name='uniq_membership_user_team' + ) + ] class Team(models.Model): """ diff --git a/generalresearch/managers/gr/authentication.py b/generalresearch/managers/gr/authentication.py index f1ac2de..18c8ba4 100644 --- a/generalresearch/managers/gr/authentication.py +++ b/generalresearch/managers/gr/authentication.py @@ -26,6 +26,7 @@ class GRUserManager(PostgresManagerWithRedis): def create( self, sub: str, + email: str | None = None, is_superuser: bool = False, ) -> GRUser: from generalresearch.models.gr.authentication import GRUser @@ -35,6 +36,7 @@ class GRUserManager(PostgresManagerWithRedis): instance = GRUser.model_validate( { "sub": sub, + "email": email, "is_superuser": is_superuser, "date_joined": now, } @@ -45,8 +47,8 @@ class GRUserManager(PostgresManagerWithRedis): with conn.cursor() as c: query = sql.SQL(""" INSERT INTO gr_user - (sub, is_superuser, date_joined) - VALUES (%(sub)s, %(is_superuser)s, %(date_joined)s) + (sub, email, is_superuser, date_joined) + VALUES (%(sub)s, %(email)s, %(is_superuser)s, %(date_joined)s) RETURNING id """) c.execute(query=query, params=data) @@ -56,63 +58,61 @@ class GRUserManager(PostgresManagerWithRedis): instance.id = gr_user_id return instance - def get_by_id(self, gr_user_id: int) -> GRUser | None: + def filter( + self, + *, + gr_user_id: int | None = None, + sub: str | None = None, + email: str | None = None, + ) -> list[GRUser]: from generalresearch.models.gr.authentication import GRUser + conditions = [] + params = {} + + if gr_user_id is not None: + conditions.append("u.id = %(gr_user_id)s") + params['gr_user_id'] = gr_user_id + if sub is not None: + conditions.append("u.sub = %(sub)s") + params['sub'] = sub + if email is not None: + conditions.append("u.email = %(email)s") + params['email'] = email + + filter_str = " AND ".join(conditions) + filter_str = "WHERE " + filter_str if filter_str else "" + + query = f""" + SELECT u.* + FROM gr_user AS u + {filter_str} + """ with self.pg_config.make_connection() as conn, conn.cursor() as c: - c.execute( - query=""" - SELECT u.* - FROM gr_user AS u - WHERE u.id = %s - LIMIT 1; - """, - params=(gr_user_id,), - ) - res = c.fetchone() + c.execute(query=query, params=params) + res = c.fetchall() - if res is None: - raise ValueError("GRUser not found") - assert isinstance( - res, dict - ), "GRUserManager.get_by_id query returned invalid results" + return [GRUser.from_postgresql(item) for item in res] - # We can return None if no MySQL results were found... but raise an - # error if returning failed for a different reason - gr_user = GRUser.from_postgresql(res) - assert isinstance(gr_user, GRUser), "GRUser not serialized correctly" - return gr_user + def _get_one(self, *, raises: bool = True, **filters: int | str) -> GRUser | None: + users = self.filter(**filters) - def get_by_sub(self, sub: str, raises=True) -> GRUser | None: - from generalresearch.models.gr.authentication import GRUser - - with self.pg_config.make_connection() as conn, conn.cursor() as c: - c.execute( - query=""" - SELECT u.* - FROM gr_user AS u - WHERE u.sub = %s - LIMIT 1; - """, - params=(sub,), - ) - res = c.fetchone() + if not users: + if raises: + raise ValueError("GRUser not found") + return None - if raises and res is None: - raise ValueError("GRUser not found") + assert len(users) == 1, "GRUserManager filter returned multiple users" + return users[0] - if res is None: - return None + def get_by_id(self, gr_user_id: int) -> GRUser | None: + return self._get_one(gr_user_id=gr_user_id) - assert isinstance( - res, dict - ), "GRUserManager.get_by_id query returned invalid results" + def get_by_sub(self, sub: str, raises=True) -> GRUser | None: + return self._get_one(sub=sub, raises=raises) - # We can return None if no MySQL results were found... but raise an - # error if returning failed for a different reason - gr_user = GRUser.from_postgresql(res) - assert isinstance(gr_user, GRUser), "GRUser not serialized correctly" - return gr_user + def get_by_email(self, email: str, raises=True) -> GRUser | None: + return self._get_one(email=email, raises=raises) def get_by_sub_or_create(self, sub: str) -> GRUser: return self.get_by_sub(sub=sub, raises=False) or self.create(sub=sub) diff --git a/generalresearch/managers/gr/team.py b/generalresearch/managers/gr/team.py index 5d9c6bb..89a3a19 100644 --- a/generalresearch/managers/gr/team.py +++ b/generalresearch/managers/gr/team.py @@ -11,7 +11,6 @@ from generalresearch.managers.base import ( PostgresManager, PostgresManagerWithRedis, ) -from generalresearch.managers.gr.authentication import GRUserManager from generalresearch.models.custom_types import UUIDStr from generalresearch.models.gr.team import ( Membership, @@ -37,20 +36,21 @@ class MembershipManager(PostgresManager): self, team: Team, gr_user: GRUser, - privilege: MembershipPrivilege = MembershipPrivilege.READ, + privilege: MembershipPrivilege = MembershipPrivilege.MEMBER, ) -> Membership: membership = Membership( uuid=uuid4().hex, - privilege=MembershipPrivilege.READ, - owner=False, + privilege=privilege, team_id=team.id, user_id=gr_user.id, + user=gr_user.to_user_ref(), created=datetime.now(tz=UTC), ) data = membership.model_dump(by_alias=True) data["team_id"] = team.id data["user_id"] = gr_user.id + data["privilege"] = membership.privilege.value # 'user_id' = {int} 5774 # 'team_id' = {int} 20736 @@ -70,8 +70,8 @@ class MembershipManager(PostgresManager): c.execute( query=sql.SQL(""" INSERT INTO common_membership - (uuid, privilege, owner, team_id, user_id, created) - VALUES (%(uuid)s, %(privilege)s, %(owner)s, %(team_id)s, + (uuid, privilege, team_id, user_id, created) + VALUES (%(uuid)s, %(privilege)s, %(team_id)s, %(user_id)s, %(created)s) RETURNING id """), @@ -86,55 +86,109 @@ class MembershipManager(PostgresManager): def exists( self, gr_user_id: PositiveInt, team_id: PositiveInt ) -> Membership | None: - with self.pg_config.make_connection() as conn, conn.cursor() as c: - c.execute( - query=sql.SQL(""" - SELECT id, uuid, privilege, owner, created, - user_id, team_id - FROM common_membership - WHERE team_id = %s AND user_id = %s - LIMIT 1 - """), - params=(team_id, gr_user_id), - ) - res = c.fetchone() + memberships = self.filter(gr_user_id=gr_user_id, team_id=team_id) + if len(memberships) > 1: + raise ValueError("User has multiple memberships in the same Team") + return memberships[0] if memberships else None - if not res: - return None + def get_by_team_id(self, team_id: PositiveInt) -> list[Membership]: + return self.filter(team_id=team_id) - return Membership.model_validate(res) + def get_by_gr_user_id(self, gr_user_id: PositiveInt) -> list[Membership]: + return self.filter(gr_user_id=gr_user_id) + + def get_by_uuid(self, membership_uuid: UUIDStr) -> Membership | None: + memberships = self.filter(membership_uuid=membership_uuid) + if len(memberships) > 1: + raise ValueError("Multiple memberships have the same UUID") + return memberships[0] if memberships else None + + def filter( + self, + *, + gr_user_id: PositiveInt | None = None, + team_id: PositiveInt | None = None, + membership_uuid: UUIDStr | None = None, + ) -> list[Membership]: + """Return memberships matching all the supplied identifiers.""" + conditions = [] + params = {} + + if gr_user_id is not None: + conditions.append("m.user_id = %(user_id)s") + params["user_id"] = gr_user_id + if team_id is not None: + conditions.append("m.team_id = %(team_id)s") + params["team_id"] = team_id + if membership_uuid is not None: + conditions.append("m.uuid = %(membership_uuid)s") + params["membership_uuid"] = membership_uuid + + filter_str = " AND ".join(conditions) + filter_str = "WHERE " + filter_str if filter_str else "" + + query = f""" + SELECT + m.id, + m.uuid, + m.privilege, + m.created, + m.user_id, + u.email AS user_email, + u.date_joined AS user_date_joined, + m.team_id + FROM common_membership AS m + JOIN gr_user AS u + ON u.id = m.user_id + {filter_str} + LIMIT 250 + """ - def get_by_team_id(self, team_id: PositiveInt) -> list[Membership]: with self.pg_config.make_connection() as conn, conn.cursor() as c: - c.execute( - query=sql.SQL(""" - SELECT id, uuid, privilege, owner, created, - user_id, team_id - FROM common_membership - WHERE team_id = %s - LIMIT 250 - """), - params=(team_id,), - ) + c.execute(query=query, params=params) res = c.fetchall() - return [Membership.model_validate(i) for i in res] + for item in res: + item["user"] = { + "email": item.pop("user_email"), + "date_joined": item.pop("user_date_joined"), + } - def get_by_gr_user_id(self, gr_user_id: PositiveInt) -> list[Membership]: - with self.pg_config.make_connection() as conn, conn.cursor() as c: - c.execute( - query=sql.SQL(""" - SELECT id, uuid, privilege, owner, created, - user_id, team_id - FROM common_membership - WHERE user_id = %s - LIMIT 250 - """), - params=(gr_user_id,), + return [Membership.model_validate(item) for item in res] + + def update_privilege( + self, + membership_id: PositiveInt, + privilege: MembershipPrivilege, + ) -> Membership: + with self.pg_config.make_connection() as conn, conn.cursor() as cursor: + cursor.execute( + query=""" + UPDATE common_membership + SET privilege = %s + WHERE id = %s + RETURNING uuid + """, + params=(privilege.value, membership_id), ) - res = c.fetchall() + row = cursor.fetchone() + if row is None: + raise LookupError(f"Membership {membership_id} does not exist") + conn.commit() - return [Membership.model_validate(i) for i in res] + membership = self.get_by_uuid(membership_uuid=row["uuid"]) + assert membership is not None + return membership + + def delete(self, membership_id: PositiveInt) -> None: + with self.pg_config.make_connection() as conn, conn.cursor() as cursor: + cursor.execute( + query="DELETE FROM common_membership WHERE id = %s", + params=(membership_id,), + ) + if cursor.rowcount != 1: + raise LookupError(f"Membership {membership_id} does not exist") + conn.commit() class TeamManager(PostgresManagerWithRedis): @@ -217,20 +271,6 @@ class TeamManager(PostgresManagerWithRedis): self.redis_client.delete(team.cache_key) return team - def add_user( - self, team: Team, gr_user: GRUser, gr_user_manager: GRUserManager - ) -> Membership: - """Create a Membership between a GRUser and a Team""" - - team.prefetch_gr_users(gr_user_manager=gr_user_manager) - - assert gr_user not in team.gr_users, ( - "Can't create multiple Memberships for the same User to the same Team" - ) - mm = MembershipManager(pg_config=self.pg_config) - - return mm.create(team=team, gr_user=gr_user) - def add_business( self, team_id: PositiveInt, @@ -251,6 +291,26 @@ class TeamManager(PostgresManagerWithRedis): ) conn.commit() + def remove_business( + self, + team_id: PositiveInt, + business_id: PositiveInt, + conn: Connection | None = None, + ) -> None: + with self.connection(conn) as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" + DELETE FROM common_team_businesses + WHERE team_id = %s and business_id = %s + """), + params=( + team_id, + business_id, + ), + ) + assert c.rowcount == 1, "Team Business Association not found" + conn.commit() + def get_by_uuid(self, team_uuid: UUIDStr) -> Team | None: from generalresearch.models.gr.team import Team diff --git a/generalresearch/models/gr/__init__.py b/generalresearch/models/gr/__init__.py index 132f72b..fb27718 100644 --- a/generalresearch/models/gr/__init__.py +++ b/generalresearch/models/gr/__init__.py @@ -1,15 +1,16 @@ -from generalresearch.models.gr.authentication import GRToken, GRUser +from generalresearch.models.gr.authentication import GRToken, GRUser, GRUserRef from generalresearch.models.gr.business import Business, BusinessCreate -from generalresearch.models.gr.team import Team +from generalresearch.models.gr.team import Team, TeamCreate from generalresearch.models.thl.finance import BusinessBalances from generalresearch.models.thl.payout import BrokerageProductPayoutEvent from generalresearch.models.thl.product import Product # To prevent type checkers from removing these imports from this file -_ = Product, BrokerageProductPayoutEvent, BusinessBalances +_ = Product, BrokerageProductPayoutEvent, BusinessBalances, GRUserRef GRUser.model_rebuild() GRToken.model_rebuild() BusinessCreate.model_rebuild() Business.model_rebuild() +TeamCreate.model_rebuild() Team.model_rebuild() diff --git a/generalresearch/models/gr/authentication.py b/generalresearch/models/gr/authentication.py index fa0bcae..5366db8 100644 --- a/generalresearch/models/gr/authentication.py +++ b/generalresearch/models/gr/authentication.py @@ -10,6 +10,7 @@ from pydantic import ( AnyHttpUrl, BaseModel, ConfigDict, + EmailStr, Field, NonNegativeInt, PositiveInt, @@ -97,6 +98,13 @@ class Claims(BaseModel): return self.sub +class GRUserRef(BaseModel): + email: EmailStr | None = Field(default=None) + date_joined: AwareDatetimeISO = Field( + description="When the GR User account signed up." + ) + + class GRUser(BaseModel): model_config = ConfigDict( # extra="forbid", @@ -106,6 +114,7 @@ class GRUser(BaseModel): id: PositiveInt | None = Field(default=None) sub: str | None = Field(max_length=200) + email: EmailStr | None = Field(default=None) is_superuser: bool = Field(default=False) date_joined: AwareDatetimeISO = Field( @@ -119,6 +128,12 @@ class GRUser(BaseModel): token: GRToken | None = Field(default=None) claims: Claims | None = Field(default=None) + def to_user_ref(self) -> GRUserRef: + return GRUserRef( + email=self.email, + date_joined=self.date_joined, + ) + def prefetch_claims( self, token: str, key: dict[str, Any], audience: str, issuer: AnyHttpUrl ) -> None: diff --git a/generalresearch/models/gr/team.py b/generalresearch/models/gr/team.py index b53b913..ee379eb 100644 --- a/generalresearch/models/gr/team.py +++ b/generalresearch/models/gr/team.py @@ -29,6 +29,7 @@ from generalresearch.models.custom_types import ( UUIDStr, UUIDStrCoerce, ) +from generalresearch.models.gr.authentication import GRUserRef from generalresearch.utils.enum import ReprEnumMeta if TYPE_CHECKING: @@ -53,9 +54,14 @@ if TYPE_CHECKING: class MembershipPrivilege(Enum, metaclass=ReprEnumMeta): + """ + This determines permissions regarding team membership management + """ + + # An admin can invite other users to a team ADMIN = 0 - MAINTAIN = 1 - READ = 2 + # A member is not + MEMBER = 1 class Membership(BaseModel): @@ -66,30 +72,29 @@ class Membership(BaseModel): a Team can have various levels of permissions and rights. """ - model_config = ConfigDict(use_enum_values=True) + model_config = ConfigDict() - id: SkipJsonSchema[PositiveInt | None] = Field( + id: PositiveInt | None = Field( default=None, ) uuid: UUIDStrCoerce = Field(examples=[uuid4().hex]) privilege: MembershipPrivilege = Field( - default=MembershipPrivilege.MAINTAIN, - examples=[MembershipPrivilege.READ.value], + default=MembershipPrivilege.ADMIN, + examples=[MembershipPrivilege.MEMBER.value], description=MembershipPrivilege.as_openapi(), ) - owner: bool = Field(default=False, examples=[True]) - created: AwareDatetimeISO = Field( description="This is when the User was added to the Team, it's when" "the Membership was created and not when the GR User " "account was created." ) - user_id: SkipJsonSchema[PositiveInt] = Field() + user_id: PositiveInt = Field() + user: GRUserRef = Field() - team_id: SkipJsonSchema[PositiveInt] = Field() + team_id: PositiveInt = Field() # prefetch attributes team: SkipJsonSchema[Team | None] = Field(default=None) diff --git a/pyproject.toml b/pyproject.toml index 15d9118..c165cd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "generalresearch" -version = "3.8.6" +version = "3.8.7" description = "Python Utilities for General Research" readme = "README.md" requires-python = ">=3.14" diff --git a/test_utils/models/gr/conftest.py b/test_utils/models/gr/conftest.py index f5dcaa1..d99bd89 100644 --- a/test_utils/models/gr/conftest.py +++ b/test_utils/models/gr/conftest.py @@ -217,6 +217,7 @@ def gr_user_factory(gr_user_manager: GRUserManager) -> Callable[..., GRUser]: def _inner( save: bool = True, sub: str | None = None, + email: str | None = None, is_superuser: bool = False, ) -> GRUser: sub = sub or f"{uuid4().hex}-{uuid4().hex}" @@ -224,6 +225,7 @@ def gr_user_factory(gr_user_manager: GRUserManager) -> Callable[..., GRUser]: if save: return gr_user_manager.create( sub=sub, + email=email, is_superuser=is_superuser, ) else: |