feat(DB): support OR scope

feat(chalice): support OR scope
tahayk · Aug 8, 2024 · c68ab35 · c68ab35
1 parent 7dc0c05
commit c68ab35
Show file tree

Hide file tree

Showing 17 changed files with 105 additions and 9 deletions.
diff --git a/api/chalicelib/core/scope.py b/api/chalicelib/core/scope.py
@@ -0,0 +1,27 @@
+from cachetools import cached, TTLCache
+
+import schemas
+from chalicelib.utils import helper
+from chalicelib.utils import pg_client
+
+cache = TTLCache(maxsize=1, ttl=24 * 60 * 60)
+
+
+@cached(cache)
+def get_scope(tenant_id) -> schemas.ScopeType:
+    with pg_client.PostgresClient() as cur:
+        query = cur.mogrify(f"""SELECT scope 
+                                FROM public.tenants;""")
+        cur.execute(query)
+        return helper.dict_to_camel_case(cur.fetchone())["scope"]
+
+
+def update_scope(tenant_id, scope: schemas.ScopeType):
+    with pg_client.PostgresClient() as cur:
+        query = cur.mogrify(f"""UPDATE public.tenants
+                                SET scope = %(scope)s;""",
+                            {"scope": scope})
+        cur.execute(query)
+        if tenant_id in cache:
+            cache.pop(tenant_id)
+    return scope
diff --git a/api/chalicelib/core/signup.py b/api/chalicelib/core/signup.py
@@ -84,6 +84,7 @@ async def create_tenant(data: schemas.UserSignupSchema):
         'refreshToken': r.pop('refreshToken'),
         'refreshTokenMaxAge': r.pop('refreshTokenMaxAge'),
         'data': {
+            "scope": "full",
             "user": r
         }
     }
diff --git a/api/chalicelib/core/tenants.py b/api/chalicelib/core/tenants.py
@@ -11,7 +11,8 @@ def get_by_tenant_id(tenant_id):
                                        tenants.created_at,
                                        '{license.EDITION}' AS edition,
                                        openreplay_version() AS version_number,
-                                       tenants.opt_out
+                                       tenants.opt_out,
+                                       scope
                                 FROM public.tenants
                                 LIMIT 1;""",
                             {"tenantId": tenant_id})

diff --git a/api/requirements.txt b/api/requirements.txt
@@ -7,6 +7,7 @@ psycopg2-binary==2.9.9
 psycopg[pool,binary]==3.2.1
 elasticsearch==8.14.0
 jira==3.8.0
+cachetools==5.4.0
 
 
 

diff --git a/api/routers/core.py b/api/routers/core.py
@@ -879,8 +879,6 @@ def health_check():
     return {}
 
 
-# tags
-
 @app.post('/{projectId}/tags', tags=["tags"])
 def tags_create(projectId: int, data: schemas.TagCreate = Body(),
                 context: schemas.CurrentContext = Depends(OR_context)):

diff --git a/api/routers/core_dynamic.py b/api/routers/core_dynamic.py
@@ -11,6 +11,7 @@
 from chalicelib.core import sessions_viewed
 from chalicelib.core import tenants, users, projects, license
 from chalicelib.core import webhook
+from chalicelib.core import scope
 from chalicelib.core.collaboration_slack import Slack
 from chalicelib.utils import captcha, smtp
 from chalicelib.utils import helper
@@ -72,7 +73,8 @@ def login_user(response: JSONResponse, spot: Optional[bool] = False, data: schem
     content = {
         'jwt': r.pop('jwt'),
         'data': {
-            "user": r
+            "user": r,
+            "scope": scope.get_scope(-1)
         }
     }
     response.set_cookie(key="refreshToken", value=refresh_token, path=COOKIE_PATH,
@@ -131,6 +133,13 @@ def edit_account(data: schemas.EditAccountSchema = Body(...),
     return users.edit_account(tenant_id=context.tenant_id, user_id=context.user_id, changes=data)
 
 
+@app.post('/account/scope', tags=["account"])
+def change_scope(data: schemas.ScopeSchema = Body(),
+                 context: schemas.CurrentContext = Depends(OR_context)):
+    data = scope.update_scope(tenant_id=-1, scope=data.scope)
+    return {'data': data}
+
+
 @app.post('/integrations/slack', tags=['integrations'])
 @app.put('/integrations/slack', tags=['integrations'])
 def add_slack_integration(data: schemas.AddCollaborationSchema,

diff --git a/api/schemas/schemas.py b/api/schemas/schemas.py
@@ -1651,3 +1651,12 @@ class TagCreate(TagUpdate):
     selector: str = Field(..., min_length=1, max_length=255)
     ignoreClickRage: bool = Field(default=False)
     ignoreDeadClick: bool = Field(default=False)
+
+
+class ScopeType(str, Enum):
+    FULL_OR = "full"
+    SPOT_ONLY = "spot"
+
+
+class ScopeSchema(BaseModel):
+    scope: ScopeType = Field(default=ScopeType.FULL_OR)
diff --git a/ee/api/chalicelib/core/scope.py b/ee/api/chalicelib/core/scope.py
@@ -0,0 +1,30 @@
+from cachetools import cached, TTLCache
+
+import schemas
+from chalicelib.utils import helper
+from chalicelib.utils import pg_client
+
+cache = TTLCache(maxsize=1, ttl=24 * 60 * 60)
+
+
+@cached(cache)
+def get_scope(tenant_id) -> schemas.ScopeType:
+    with pg_client.PostgresClient() as cur:
+        query = cur.mogrify(f"""SELECT scope 
+                                FROM public.tenants
+                                WHERE tenant_id=%(tenant_id)s;""",
+                            {"tenant_id": tenant_id})
+        cur.execute(query)
+        return helper.dict_to_camel_case(cur.fetchone())["scope"]
+
+
+def update_scope(tenant_id, scope: schemas.ScopeType):
+    with pg_client.PostgresClient() as cur:
+        query = cur.mogrify(f"""UPDATE public.tenants
+                                SET scope = %(scope)s
+                                WHERE tenant_id=%(tenant_id)s;""",
+                            {"scope": scope, "tenant_id": tenant_id})
+        cur.execute(query)
+        if tenant_id in cache:
+            cache.pop(tenant_id)
+    return scope
diff --git a/ee/api/chalicelib/core/signup.py b/ee/api/chalicelib/core/signup.py
@@ -94,6 +94,7 @@ async def create_tenant(data: schemas.UserSignupSchema):
         'refreshToken': r.pop('refreshToken'),
         'refreshTokenMaxAge': r.pop('refreshTokenMaxAge'),
         'data': {
+            "scope": "full",
             "user": r
         }
     }
diff --git a/ee/api/chalicelib/core/tenants.py b/ee/api/chalicelib/core/tenants.py
@@ -31,7 +31,8 @@ def get_by_tenant_id(tenant_id):
                                         '{license.EDITION}' AS edition,
                                         openreplay_version() AS version_number,
                                         tenants.opt_out,
-                                        tenants.tenant_key
+                                        tenants.tenant_key,
+                                        scope
                                 FROM public.tenants
                                 WHERE tenants.tenant_id = %(tenantId)s 
                                     AND tenants.deleted_at ISNULL

diff --git a/ee/api/requirements.txt b/ee/api/requirements.txt
@@ -7,6 +7,7 @@ psycopg2-binary==2.9.9
 psycopg[pool,binary]==3.2.1
 elasticsearch==8.14.0
 jira==3.8.0
+cachetools==5.4.0
 
 
 

diff --git a/ee/api/routers/core_dynamic.py b/ee/api/routers/core_dynamic.py
@@ -11,6 +11,7 @@
 from chalicelib.core import sessions_viewed
 from chalicelib.core import tenants, users, projects, license
 from chalicelib.core import webhook
+from chalicelib.core import scope
 from chalicelib.core.collaboration_slack import Slack
 from chalicelib.core.users import get_user_settings
 from chalicelib.utils import SAML2_helper, smtp
@@ -78,6 +79,7 @@ def login_user(response: JSONResponse, spot: Optional[bool] = False, data: schem
     content = {
         'jwt': r.pop('jwt'),
         'data': {
+            "scope":scope.get_scope(r["tenantId"]),
             "user": r
         }
     }
@@ -138,6 +140,11 @@ def get_account(context: schemas.CurrentContext = Depends(OR_context)):
 def edit_account(data: schemas.EditAccountSchema = Body(...),
                  context: schemas.CurrentContext = Depends(OR_context)):
     return users.edit_account(tenant_id=context.tenant_id, user_id=context.user_id, changes=data)
+@app.post('/account/scope', tags=["account"])
+def change_scope(data: schemas.ScopeSchema = Body(),
+                 context: schemas.CurrentContext = Depends(OR_context)):
+    data = scope.update_scope(tenant_id=-1, scope=data.scope)
+    return {'data': data}
 
 
 @app.post('/integrations/slack', tags=['integrations'])

diff --git a/ee/scripts/schema/db/init_dbs/postgresql/1.20.0/1.20.0.sql b/ee/scripts/schema/db/init_dbs/postgresql/1.20.0/1.20.0.sql
@@ -30,7 +30,8 @@ WHERE NOT permissions @> '{SPOT}'
 UPDATE public.roles
 SET permissions = (SELECT array_agg(distinct e) FROM unnest(permissions || '{SPOT_PUBLIC}') AS e)
 WHERE NOT permissions @> '{SPOT_PUBLIC}'
-  AND name ILIKE 'owner';
+  AND NOT service_role;
+--   AND name ILIKE 'owner';
 
 ALTER TABLE IF EXISTS public.users
     ADD COLUMN IF NOT EXISTS spot_jwt_iat         timestamp without time zone NULL DEFAULT NULL,
@@ -49,6 +50,9 @@ CREATE TABLE IF NOT EXISTS or_cache.autocomplete_top_values
     UNIQUE NULLS NOT DISTINCT (project_id, event_type, event_key)
 );
 
+ALTER TABLE IF EXISTS public.tenants
+    ADD COLUMN IF NOT EXISTS scope text NOT NULL DEFAULT 'full';
+
 COMMIT;
 
 \elif :is_next

diff --git a/ee/scripts/schema/db/init_dbs/postgresql/init_schema.sql b/ee/scripts/schema/db/init_dbs/postgresql/init_schema.sql
@@ -103,7 +103,8 @@ CREATE TABLE public.tenants
     t_sessions     bigint                      NOT NULL DEFAULT 0,
     t_users        integer                     NOT NULL DEFAULT 1,
     t_integrations integer                     NOT NULL DEFAULT 0,
-    last_telemetry bigint                      NOT NULL DEFAULT CAST(EXTRACT(epoch FROM date_trunc('day', now())) * 1000 AS BIGINT)
+    last_telemetry bigint                      NOT NULL DEFAULT CAST(EXTRACT(epoch FROM date_trunc('day', now())) * 1000 AS BIGINT),
+    scope          text                        NOT NULL DEFAULT 'full'
 );
 
 

diff --git a/scripts/schema/db/init_dbs/postgresql/1.20.0/1.20.0.sql b/scripts/schema/db/init_dbs/postgresql/1.20.0/1.20.0.sql
@@ -39,6 +39,9 @@ CREATE TABLE IF NOT EXISTS or_cache.autocomplete_top_values
     UNIQUE NULLS NOT DISTINCT (project_id, event_type, event_key)
 );
 
+ALTER TABLE IF EXISTS public.tenants
+    ADD COLUMN IF NOT EXISTS scope text NOT NULL DEFAULT 'full';
+
 COMMIT;
 
 \elif :is_next

diff --git a/scripts/schema/db/init_dbs/postgresql/init_schema.sql b/scripts/schema/db/init_dbs/postgresql/init_schema.sql
@@ -103,8 +103,9 @@ CREATE TABLE public.tenants
     t_sessions     bigint                      NOT NULL DEFAULT 0,
     t_users        integer                     NOT NULL DEFAULT 1,
     t_integrations integer                     NOT NULL DEFAULT 0,
-    last_telemetry bigint                      NOT NULL DEFAULT CAST(EXTRACT(epoch FROM date_trunc('day', now())) * 1000 AS BIGINT)
-        CONSTRAINT onerow_uni CHECK (tenant_id = 1)
+    last_telemetry bigint                      NOT NULL DEFAULT CAST(EXTRACT(epoch FROM date_trunc('day', now())) * 1000 AS BIGINT),
+    scope          text                        NOT NULL DEFAULT 'full',
+    CONSTRAINT onerow_uni CHECK (tenant_id = 1)
 );
 
 CREATE TYPE user_role AS ENUM ('owner', 'admin', 'member');

diff --git a/third-party.md b/third-party.md
@@ -55,6 +55,7 @@ up to date with every new library you use.
 | sqlalchemy                 | MIT                | Python         |
 | pandas-redshift            | MIT                | Python         |
 | confluent-kafka            | Apache2            | Python         |
+| cachetools                 | MIT                | Python         |
 | amplitude-js               | MIT                | JavaScript     |
 | classnames                 | MIT                | JavaScript     |
 | codemirror                 | MIT                | JavaScript     |