diff options
| -rw-r--r-- | generalresearch/managers/base.py | 5 | ||||
| -rw-r--r-- | generalresearch/managers/thl/user_manager/mysql_user_manager.py | 254 | ||||
| -rw-r--r-- | generalresearch/managers/thl/user_manager/user_manager.py | 8 | ||||
| -rw-r--r-- | generalresearch/models/custom_types.py | 5 | ||||
| -rw-r--r-- | generalresearch/thl_django/common/models.py | 10 | ||||
| -rw-r--r-- | generalresearch/thl_django/userprofile/models.py | 3 |
6 files changed, 268 insertions, 17 deletions
diff --git a/generalresearch/managers/base.py b/generalresearch/managers/base.py index 257032f..faffddf 100644 --- a/generalresearch/managers/base.py +++ b/generalresearch/managers/base.py @@ -21,15 +21,16 @@ class Permission(int, Enum): class Manager: DEFAULT_PAGE_SIZE = 50 + @classmethod def validate_pagination( - self, + cls, page: int | None = 1, size: int | None = None, ): page = page if page is not None else 1 assert type(page) is int assert page >= 1, "page starts at 1" - size = size if size is not None else self.DEFAULT_PAGE_SIZE + size = size if size is not None else cls.DEFAULT_PAGE_SIZE assert type(size) is int assert 1 <= size <= 100 offset = (page - 1) * size diff --git a/generalresearch/managers/thl/user_manager/mysql_user_manager.py b/generalresearch/managers/thl/user_manager/mysql_user_manager.py index d2732ed..6b72edd 100644 --- a/generalresearch/managers/thl/user_manager/mysql_user_manager.py +++ b/generalresearch/managers/thl/user_manager/mysql_user_manager.py @@ -5,14 +5,16 @@ import operator from collections.abc import Collection from datetime import UTC, datetime from threading import Lock -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Tuple from uuid import uuid4 import psycopg from cachetools import LRUCache, cachedmethod -from psycopg import sql +from psycopg import Connection, sql +from pydantic import NonNegativeInt -from generalresearch.models.custom_types import UUIDStr +from generalresearch.managers.base import Permission, PostgresManager +from generalresearch.models.custom_types import UUIDStr, try_normalize_email from generalresearch.models.thl.user import User if TYPE_CHECKING: @@ -23,25 +25,59 @@ logger = logging.getLogger() logger.setLevel(logging.INFO) -class MysqlUserManager: - def __init__(self, pg_config: PostgresConfig, is_read_replica: bool): - self.pg_config = pg_config +class MysqlUserManager(PostgresManager): + def __init__( + self, + pg_config: PostgresConfig, + is_read_replica: bool, + permissions: Collection[Permission] | None = None, + ): self.is_read_replica = is_read_replica self.product_id_exists_cache = LRUCache(maxsize=5000) self.product_id_exists_cache_lock = Lock() + super().__init__(pg_config=pg_config, permissions=permissions) - def _set_last_seen(self, user: User) -> None: + def _set_last_seen(self, user: User, conn: Connection | None = None) -> None: # Don't call this directly. Use UserManager.set_last_seen() assert not self.is_read_replica now = datetime.now(tz=UTC) - self.pg_config.execute_write( - """ + with self.connection(conn) as conn, conn.cursor() as c: + c.execute( + """ UPDATE thl_user SET last_seen = %s WHERE id = %s """, - params=[now, user.user_id], - ) + params=[now, user.user_id], + ) + + def _set_last_country_ip( + self, + user: User, + last_ip: str, + last_country_iso: str | None = None, + conn: Connection | None = None, + ) -> None: + # Don't call this directly. Use UserManager.set_last_country_ip() + assert not self.is_read_replica + params = { + "user_id": user.user_id, + "last_ip": last_ip, + "last_country_iso": last_country_iso, + } + with self.connection(conn) as conn, conn.cursor() as c: + c.execute( + """ + UPDATE thl_user + SET last_ip = %(last_ip)s, last_country_iso = %(last_country_iso)s + WHERE id = %(user_id)s + AND ( + last_ip IS DISTINCT FROM %(last_ip)s OR + last_country_iso IS DISTINCT FROM %(last_country_iso)s + ) + """, + params, + ) def _change_product_user_id(self, *, user: User, new_product_user_id: str) -> User: """Change a user's supplier-provided ID in the primary database.""" @@ -168,7 +204,6 @@ class MysqlUserManager: "last_seen": now, } - # in postgres, you do not include the auto-increment id column query = sql.SQL(""" INSERT INTO thl_user (uuid, product_id, product_user_id, created, @@ -182,6 +217,15 @@ class MysqlUserManager: with self.pg_config.make_connection() as conn, conn.cursor() as c: c.execute(query=query, params=params) user_id = c.fetchone()["id"] + c.execute( + """ + UPDATE userprofile_brokerageproduct + SET user_count = user_count + 1 + WHERE id = %(product_id)s + """, + params=params, + ) + assert c.rowcount == 1 except psycopg.IntegrityError: # Two machines/processes are trying to create this same (product_id, product_user_id) # at the same time. There's a unique index, so mysql will not let two be created. @@ -327,3 +371,189 @@ class MysqlUserManager: params={"user_uuids": user_uuids}, ) return [User.from_db(x) for x in res] + + @staticmethod + def make_filter_str( + product_id: str | None = None, + business_id: str | None = None, + team_id: str | None = None, + product_user_id: str | None = None, + email_startswith: str | None = None, + display_name_like: str | None = None, + allow_global_search: bool = False, + ) -> tuple[str, dict[str, Any]]: + filters = [] + params = {} + + if product_id is not None: + params["product_id"] = product_id + filters.append("product_id = %(product_id)s") + if business_id: + params["business_id"] = business_id + filters.append("ubp.business_id = %(business_id)s") + if team_id: + params["team_id"] = team_id + filters.append("ubp.team_id = %(team_id)s") + if product_user_id: + # This should only be allowed within the context of a product, team, or business, + # unless we are admins + if not allow_global_search: + assert any([product_id, business_id, team_id]), ( + "must scope to a product, team, or business" + ) + params["product_user_id"] = product_user_id + filters.append("u.product_user_id = %(product_user_id)s") + if email_startswith is not None: + if not allow_global_search: + assert any([product_id, business_id, team_id]), ( + "must scope to a product, team, or business" + ) + params["email_startswith"] = f"{email_startswith}%" + filters.append("email_address LIKE %(email_startswith)s") + normalized_email = try_normalize_email(email_startswith) + if normalized_email != email_startswith: + params["normalized_email_startswith"] = f"{normalized_email}%" + filters.append("canonical_email LIKE %(normalized_email_startswith)s") + if display_name_like is not None: + if not allow_global_search: + assert any([product_id, business_id, team_id]), ( + "must scope to a product, team, or business" + ) + params["display_name_like"] = f"{display_name_like}%" + filters.append("display_name ILIKE %(display_name_like)s") + + return "WHERE " + " AND ".join(filters) if filters else "", params + + def filter( + self, + product_id: str | None = None, + business_id: str | None = None, + team_id: str | None = None, + product_user_id: str | None = None, + email_startswith: str | None = None, + display_name_like: str | None = None, + page: int | None = None, + size: int | None = None, + order_field: str = "created", + descending: bool = False, + allow_global_search: bool = False, + ) -> tuple[list[User], NonNegativeInt]: + with self.pg_config.make_connection() as conn: + items = self.filter_page( + product_id=product_id, + business_id=business_id, + team_id=team_id, + product_user_id=product_user_id, + email_startswith=email_startswith, + display_name_like=display_name_like, + page=page, + size=size, + order_field=order_field, + descending=descending, + allow_global_search=allow_global_search, + conn=conn, + ) + total = self.filter_count( + product_id=product_id, + business_id=business_id, + team_id=team_id, + product_user_id=product_user_id, + email_startswith=email_startswith, + display_name_like=display_name_like, + allow_global_search=allow_global_search, + conn=conn, + ) + return items, total + + def filter_page( + self, + product_id: str | None = None, + business_id: str | None = None, + team_id: str | None = None, + product_user_id: str | None = None, + email_startswith: str | None = None, + display_name_like: str | None = None, + page: int | None = None, + size: int | None = None, + order_field: str = "created", + descending: bool = False, + allow_global_search: bool = False, + conn: Connection | None = None, + ) -> list[User]: + pagination_sql = self.validate_pagination(page=page, size=size) + filter_str, params = self.make_filter_str( + product_id=product_id, + business_id=business_id, + team_id=team_id, + product_user_id=product_user_id, + email_startswith=email_startswith, + display_name_like=display_name_like, + allow_global_search=allow_global_search, + ) + order_fields = { + "created": "u.created", + "last_seen": "u.last_seen", + "product_user_id": "u.product_user_id", + } + order_by_sql = ( + f"{order_fields[order_field]} {'DESC' if descending else 'ASC'}, u.uuid" + ) + + query = f""" + SELECT u.id as user_id, u.uuid, u.product_id, u.product_user_id, + u.blocked, u.created, u.last_seen, + tu.email_address, tu.display_name, ubp.business_id, ubp.team_id + FROM thl_user u + LEFT JOIN thl_usermetadata tu on u.id = tu.user_id + JOIN userprofile_brokerageproduct ubp ON ubp.id = u.product_id + {filter_str} + ORDER BY {order_by_sql} + {pagination_sql} + """ + with self.connection(conn) as conn, conn.cursor() as c: + c.execute(query, params) + res = c.fetchall() + + return [User.from_db(x) for x in res] + + def filter_count( + self, + product_id: str | None = None, + business_id: str | None = None, + team_id: str | None = None, + product_user_id: str | None = None, + email_startswith: str | None = None, + display_name_like: str | None = None, + allow_global_search: bool = False, + conn: Connection | None = None, + ) -> NonNegativeInt: + + filter_str, params = self.make_filter_str( + product_id=product_id, + business_id=business_id, + team_id=team_id, + product_user_id=product_user_id, + email_startswith=email_startswith, + display_name_like=display_name_like, + allow_global_search=allow_global_search + ) + if not any([product_user_id, email_startswith, display_name_like]): + # We're filtering on product, team, or business only. So use the + # precomputed counts + query = f""" + SELECT SUM(user_count) AS cnt + FROM userprofile_brokerageproduct ubp + {filter_str} + """ + else: + query = f""" + SELECT COUNT(1) AS cnt + FROM thl_user u + LEFT JOIN thl_usermetadata tu on u.id = tu.user_id + JOIN userprofile_brokerageproduct ubp ON ubp.id = u.product_id + {filter_str} + """ + with self.connection(conn) as conn, conn.cursor() as c: + c.execute(query, params) + res = c.fetchall() + return res[0]["cnt"] diff --git a/generalresearch/managers/thl/user_manager/user_manager.py b/generalresearch/managers/thl/user_manager/user_manager.py index f322d2b..6358e84 100644 --- a/generalresearch/managers/thl/user_manager/user_manager.py +++ b/generalresearch/managers/thl/user_manager/user_manager.py @@ -95,6 +95,14 @@ class UserManager: assert Permission.UPDATE in self.sql_permissions, "permission error" return self.mysql_user_manager._set_last_seen(user) + def set_last_country_ip( + self, user: User, last_ip: str, last_country_iso: str | None = None + ) -> None: + assert Permission.UPDATE in self.sql_permissions, "permission error" + return self.mysql_user_manager._set_last_country_ip( + user, last_ip, last_country_iso + ) + def change_product_user_id(self, *, user: User, new_product_user_id: str) -> User: """Change a user's supplier-provided ID and refresh user lookup caches. diff --git a/generalresearch/models/custom_types.py b/generalresearch/models/custom_types.py index 0040b92..26c8b7e 100644 --- a/generalresearch/models/custom_types.py +++ b/generalresearch/models/custom_types.py @@ -47,6 +47,11 @@ def validate_canonical_email(value: str) -> str: ) return value +def try_normalize_email(value: str) -> str: + try: + return email_normalize.normalize(value).normalized_address + except ValueError: + return value CanonicalEmailStr = Annotated[ EmailStr, diff --git a/generalresearch/thl_django/common/models.py b/generalresearch/thl_django/common/models.py index 67c5ca5..afbe722 100644 --- a/generalresearch/thl_django/common/models.py +++ b/generalresearch/thl_django/common/models.py @@ -2,7 +2,7 @@ import uuid from django.db import models from django.db.models import Q - +from django.contrib.postgres.indexes import GinIndex class THLSession(models.Model): """ @@ -379,8 +379,7 @@ class THLUserMetadata(models.Model): canonical_email = models.CharField(max_length=320, null=True) # Not unique within a BP, or anything like that. A user - # can set this to whatever they like. No index - # as we will not ever look up a user by their name. + # can set this to whatever they like. display_name = models.CharField(max_length=255, null=True) class Meta: @@ -391,6 +390,11 @@ class THLUserMetadata(models.Model): models.Index(fields=["email_sha1"]), models.Index(fields=["email_md5"]), models.Index(fields=["canonical_email"]), + GinIndex( + fields=["display_name"], + opclasses=["gin_trgm_ops"], + name="usermetadata_name_trgm", + ), ] diff --git a/generalresearch/thl_django/userprofile/models.py b/generalresearch/thl_django/userprofile/models.py index 445a53c..7b4961f 100644 --- a/generalresearch/thl_django/userprofile/models.py +++ b/generalresearch/thl_django/userprofile/models.py @@ -88,6 +88,9 @@ class BrokerageProduct(models.Model): # ProductPOPFinancials model pop_financial = models.JSONField(default=None, null=True) + # This is the total number of users created + user_count = models.IntegerField(default=0, null=False) + users_active_7d = models.IntegerField(default=None, null=True) task_completes_7d = models.IntegerField(default=None, null=True) # Net Earnings over the last 7 days (in USD Cents, this can be positive or negative) |