gr: a lot of changes to managers (membership priveleges, checks for team membershipness, add/remove business)

5 files changed, 128 insertions, 118 deletions

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..405435e 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,12 +36,11 @@ 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,
             created=datetime.now(tz=UTC),
@@ -70,8 +68,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 +84,55 @@ 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()
-
-        if not res:
-            return None
-
-        return Membership.model_validate(res)
+        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
 
     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,),
-            )
-            res = c.fetchall()
-
-        return [Membership.model_validate(i) for i in res]
+        return self.filter(team_id=team_id)
 
     def get_by_gr_user_id(self, gr_user_id: PositiveInt) -> list[Membership]:
+        return self.filter(gr_user_id=gr_user_id)
+
+    def filter(
+        self,
+        *,
+        gr_user_id: PositiveInt | None = None,
+        team_id: PositiveInt | 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
+
+        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,
+                m.team_id
+            FROM common_membership AS m
+            {filter_str}
+            LIMIT 250
+        """
+
         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,),
-            )
+            c.execute(query=query, params=params)
             res = c.fetchall()
 
-        return [Membership.model_validate(i) for i in res]
+        return [Membership.model_validate(item) for item in res]
 
 
 class TeamManager(PostgresManagerWithRedis):
@@ -217,20 +215,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 +235,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..c6d9b2a 100644
--- a/generalresearch/models/gr/__init__.py
+++ b/generalresearch/models/gr/__init__.py
@@ -1,6 +1,6 @@
 from generalresearch.models.gr.authentication import GRToken, GRUser
 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
@@ -12,4 +12,5 @@ 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..123e1dc 100644
--- a/generalresearch/models/gr/authentication.py
+++ b/generalresearch/models/gr/authentication.py
@@ -14,6 +14,7 @@ from pydantic import (
     NonNegativeInt,
     PositiveInt,
     field_validator,
+    EmailStr,
 )
 
 from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr
@@ -106,6 +107,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(
diff --git a/generalresearch/models/gr/team.py b/generalresearch/models/gr/team.py
index b53b913..d340c68 100644
--- a/generalresearch/models/gr/team.py
+++ b/generalresearch/models/gr/team.py
@@ -53,9 +53,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):
@@ -74,13 +79,11 @@ class Membership(BaseModel):
     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 "