From aaa2009579cf1192ab87e7609806f88c5246a3a7 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 13 Nov 2024 01:49:52 +0530 Subject: [PATCH 01/57] apis and migrations added --- ddpui/api/org_preferences_api.py | 281 ++++++++++++++++++ ddpui/api/user_preferences_api.py | 19 +- ...serpreferences_discord_webhook_and_more.py | 20 ++ ddpui/migrations/0105_orgpreferences.py | 51 ++++ ...er_orgpreferences_llm_optin_approved_by.py | 24 ++ ddpui/migrations/0107_orgsupersets.py | 37 +++ .../0108_userpreferences_llm_optin.py | 17 ++ ...gpreferences_org_alter_orgsupersets_org.py | 29 ++ ddpui/models/__init__.py | 2 + ddpui/models/org_preferences.py | 32 ++ ddpui/models/org_supersets.py | 13 + ddpui/models/userpreferences.py | 3 +- ddpui/routes.py | 3 + ddpui/schemas/org_preferences_schema.py | 42 +++ ddpui/schemas/userpreferences_schema.py | 8 +- ddpui/settings.py | 20 +- ddpui/utils/constants.py | 5 + seed/002_permissions.json | 16 + seed/003_role_permissions.json | 32 ++ 19 files changed, 624 insertions(+), 30 deletions(-) create mode 100644 ddpui/api/org_preferences_api.py create mode 100644 ddpui/migrations/0104_remove_userpreferences_discord_webhook_and_more.py create mode 100644 ddpui/migrations/0105_orgpreferences.py create mode 100644 ddpui/migrations/0106_alter_orgpreferences_llm_optin_approved_by.py create mode 100644 ddpui/migrations/0107_orgsupersets.py create mode 100644 ddpui/migrations/0108_userpreferences_llm_optin.py create mode 100644 ddpui/migrations/0109_alter_orgpreferences_org_alter_orgsupersets_org.py create mode 100644 ddpui/models/org_preferences.py create mode 100644 ddpui/models/org_supersets.py create mode 100644 ddpui/schemas/org_preferences_schema.py diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py new file mode 100644 index 00000000..48d33482 --- /dev/null +++ b/ddpui/api/org_preferences_api.py @@ -0,0 +1,281 @@ +from ninja import Router +from ninja.errors import HttpError +from ddpui import auth +from ddpui.models.org_preferences import OrgPreferences +from ddpui.models.org_supersets import OrgSupersets +from django.utils import timezone +from ddpui.schemas.org_preferences_schema import ( + CreateOrgPreferencesSchema, + UpdateOrgPreferencesSchema, + UpdateLLMOptinSchema, + UpdateDiscordNotificationsSchema, + CreateOrgSupersetDetailsSchema +) +from ddpui.utils.constants import ( + AIRBYTE_URL_TO_GET_VERSION, + PREFECT_URL_TO_GET_VERSION, + DBT_VERSION_COMMAND, + ELEMENTARY_VERSION_COMMAND +) +from ddpui.auth import has_permission +from ddpui.models.org import Org +from ddpui.models.org_user import OrgUser +import requests +import subprocess +orgpreference_router = Router() + + + +@orgpreference_router.post("/", auth=auth.CustomAuthMiddleware()) +def create_org_preferences(request, payload: CreateOrgPreferencesSchema): + print(payload, "payload") + """Creates preferences for an organization""" + if OrgPreferences.objects.filter(org_id=payload.org_id).exists(): + raise HttpError(400, "Organization preferences already exist") + + org_preferences = OrgPreferences.objects.create( + **payload.dict() + ) + + preferences = { + "trial_start_date": org_preferences.trial_start_date, + "trial_end_date": org_preferences.trial_end_date, + "llm_optin": org_preferences.llm_optin, + "llm_optin_approved_by": org_preferences.llm_optin_approved_by, + "llm_optin_date": org_preferences.llm_optin_date, + "enable_discord_notifications": org_preferences.enable_discord_notifications, + "discord_webhook": org_preferences.discord_webhook, + } + + return {"success": True, "res": preferences} + + +# @orgpreference_router.put("/{org_id}/", auth=auth.CustomAuthMiddleware()) +# def update_org_preferences(request, payload: UpdateOrgPreferencesSchema, org_id: int): +# """Updates preferences for an organization""" +# payload.llm_optin_approved_by = request.orguser +# try: +# org_preferences = OrgPreferences.objects.get(org_id=org_id) +# except OrgPreferences.DoesNotExist: +# raise HttpError(404, "Preferences for this organization not found") + +# #check this line please. +# # if org.type=="demo": +# # if payload.trial_start_date is None or payload.trial_end_date is None: +# # raise HttpError(400, "Trial accounts must have trial_start_date and trial_end_date.") + + +# if payload.llm_optin is True: +# # if not request.orguser.has_permission("can_edit_llm_settings"): +# # raise HttpError(403, "You do not have permission to edit LLM settings.") + +# payload.llm_optin_date = timezone.now() +# if payload.llm_optin_approved_by is None or payload.llm_optin_date is None: +# raise HttpError(400, "If LLM opt-in is true, both llm_optin_approved_by and llm_optin_date must be provided.") + +# # Update fields only if they are present in the payload + +# org_preferences.save() + +# return {"success": True, "res": 1} + +@orgpreference_router.put("/{org_id}/llm_approval", auth=auth.CustomAuthMiddleware()) +@has_permission(['can_edit_llm_settings']) +def update_org_preferences(request, payload: UpdateLLMOptinSchema, org_id: int): + """Updates llm preferences for an organization""" + + try: + org_preferences = OrgPreferences.objects.get(org_id=org_id) + except OrgPreferences.DoesNotExist: + raise HttpError(404, "Preferences for this organization not found") + + # Determine fields based on the llm_optin value + if payload.llm_optin is True: + llm_optin_approved_by = request.orguser + llm_optin_date = timezone.now() + else: + # Set them to None for False + llm_optin_approved_by = None + llm_optin_date = None + + # Update the OrgPreferences instance + org_preferences.llm_optin = payload.llm_optin + org_preferences.llm_optin_approved_by = llm_optin_approved_by + org_preferences.llm_optin_date = llm_optin_date + org_preferences.save() + + return {"success": True, "res": 1} + +@orgpreference_router.put("/{org_id}/enable-discord-notifications", auth=auth.CustomAuthMiddleware()) +@has_permission(['can_edit_discord_notifications_settings']) +def update_org_preferences(request, payload: UpdateDiscordNotificationsSchema, org_id: int): + """Updates preferences for an organization""" + + try: + org_preferences = OrgPreferences.objects.get(org_id=org_id) + except OrgPreferences.DoesNotExist: + raise HttpError(404, "Preferences for this organization not found") + + # Check for discord_link in the database when enabling notifications + if payload.enable_discord_notifications: + if not org_preferences.discord_webhook and not payload.discord_webhook: + raise HttpError(400, "Discord link is missing. Please add a Discord link to enable notifications.") + + # Update the OrgPreferences instance + org_preferences.enable_discord_notifications = payload.enable_discord_notifications + if payload.discord_webhook: + org_preferences.discord_webhook =payload.discord_webhook + org_preferences.save() + + return {"success": True, "res": 1} + + + + + +@orgpreference_router.get("/{org_id}/", auth=auth.CustomAuthMiddleware()) +def get_org_preferences(request, org_id: int): + """Gets preferences for an organization by org_id""" + try: + org_preferences = OrgPreferences.objects.select_related( + "org", + "llm_optin_approved_by" + ).get(org_id=org_id) + except OrgPreferences.DoesNotExist: + raise HttpError(404, "Organization preferences not found") + + approved_by_data = None + if org_preferences.llm_optin_approved_by and org_preferences.llm_optin_approved_by.user: + user = org_preferences.llm_optin_approved_by.user + approved_by_data = { + "user_id": user.pk, + "email": user.email + } + + preferences = { + "org": { + "id": org_preferences.org.id, + "name": org_preferences.org.name, + "slug": org_preferences.org.slug, + "type": org_preferences.org.type, + }, + "trial_start_date": org_preferences.trial_start_date, + "trial_end_date": org_preferences.trial_end_date, + "llm_optin": org_preferences.llm_optin, + "llm_optin_approved_by": approved_by_data, + "llm_optin_date": org_preferences.llm_optin_date, + "enable_discord_notifications": org_preferences.enable_discord_notifications, + "discord_webhook": org_preferences.discord_webhook, + } + + return {"success": True, "res": preferences} + +# api to get more superset related information. +@orgpreference_router.get("/{org_id}/org-superset", auth=auth.CustomAuthMiddleware()) +def get_org_superset_details(request, org_id: int): + """Gets preferences for an organization by org_id""" + try: + org_superset = OrgSupersets.objects.select_related( + "org", + ).get(org_id=org_id) + + except OrgSupersets.DoesNotExist: + raise HttpError(404, "Organizations superset details not found") + + response_data = { + "org": { + "id": org_superset.org.id, + "name": org_superset.org.name, + "slug": org_superset.org.slug, + "type": org_superset.org.type, + }, + "superset_version": org_superset.superset_version + } + return {"success": True, "res": response_data} + + + +@orgpreference_router.post("/{org_id}/org-superset", auth=auth.CustomAuthMiddleware()) +def create_org_superset_details(request, payload:CreateOrgSupersetDetailsSchema, org_id: int): + """Creates an entry in db with org superset details""" + try: + org = Org.objects.get(id=org_id) + except Org.DoesNotExist: + raise HttpError(404, "Organization not found") + + org_superset, created = OrgSupersets.objects.update_or_create( + org=org, + **payload.dict() + ) + + response_data = { + "org": { + "id": org.id, + "name": org.name, + "slug": org.slug, + "type": org.type, + }, + "id": org_superset.id, + "superset_version": org_superset.superset_version + } + + return {"success": True, "res": response_data} + +@orgpreference_router.get("/{org_id}/versions", auth=auth.CustomAuthMiddleware()) +def get_tools_versions( request, org_id: int): + versions = {} + + # Airbyte Version + try: + airbyte_url = AIRBYTE_URL_TO_GET_VERSION + airbyte_response = requests.get(airbyte_url) + if airbyte_response.status_code == 200: + airbyte_data = airbyte_response.json() + versions["Airbyte"] = airbyte_data.get("version") + else: + versions["Airbyte"] = "Error: Unable to retrieve version" + except Exception as e: + versions["Airbyte"] = f"Error: {e}" + + # Prefect Version + try: + prefect_url =PREFECT_URL_TO_GET_VERSION + prefect_response = requests.get(prefect_url) + if prefect_response.status_code == 200: + versions["Prefect"] = prefect_response.text.strip() + else: + versions["Prefect"] = "Error: Unable to retrieve version" + except Exception as e: + versions["Prefect"] = f"Error: {e}" + + # dbt Version + try: + dbt_version_command = DBT_VERSION_COMMAND + dbt_output = subprocess.check_output(dbt_version_command, text=True) + for line in dbt_output.splitlines(): + if "installed:" in line: + versions["dbt"] = line.split(":")[1].strip() + break + except Exception as e: + versions["dbt"] = f"Error: {e}" + + # Elementary Version + try: + elementary_version_command = ELEMENTARY_VERSION_COMMAND + elementary_output = subprocess.check_output(elementary_version_command, text=True) + for line in elementary_output.splitlines(): + if "Elementary version" in line: + versions["Elementary"] = line.split()[-1].strip() + break + except Exception as e: + versions["Elementary"] = f"Error: {e}" + + # Superset_version + try: + org_superset = OrgSupersets.objects.get(org_id=org_id) + except OrgSupersets.DoesNotExist: + raise HttpError(404, "Organizations superset details not found") + + versions["Superset"] = org_superset.superset_version + + return {"success": True, "res": versions} \ No newline at end of file diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index ba428231..de19ca59 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -21,15 +21,14 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): userpreferences = UserPreferences.objects.create( orguser=orguser, - enable_discord_notifications=payload.enable_discord_notifications, + llm_optin = payload.llm_optin, enable_email_notifications=payload.enable_email_notifications, - discord_webhook=payload.discord_webhook, + ) preferences = { - "discord_webhook": userpreferences.discord_webhook, "enable_email_notifications": userpreferences.enable_email_notifications, - "enable_discord_notifications": userpreferences.enable_discord_notifications, + "llm_optin": userpreferences.llm_optin } return {"success": True, "res": preferences} @@ -42,19 +41,16 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) - if payload.enable_discord_notifications is not None: - user_preferences.enable_discord_notifications = payload.enable_discord_notifications + if payload.llm_optin is not None: + user_preferences.llm_optin = payload.llm_optin if payload.enable_email_notifications is not None: user_preferences.enable_email_notifications = payload.enable_email_notifications - if payload.discord_webhook is not None: - user_preferences.discord_webhook = payload.discord_webhook user_preferences.save() preferences = { - "discord_webhook": user_preferences.discord_webhook, "enable_email_notifications": user_preferences.enable_email_notifications, - "enable_discord_notifications": user_preferences.enable_discord_notifications, + "llm_optin": user_preferences.llm_optin, } return {"success": True, "res": preferences} @@ -68,8 +64,7 @@ def get_user_preferences(request): user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) preferences = { - "discord_webhook": user_preferences.discord_webhook, "enable_email_notifications": user_preferences.enable_email_notifications, - "enable_discord_notifications": user_preferences.enable_discord_notifications, + "llm_optin": user_preferences.llm_optin, } return {"success": True, "res": preferences} diff --git a/ddpui/migrations/0104_remove_userpreferences_discord_webhook_and_more.py b/ddpui/migrations/0104_remove_userpreferences_discord_webhook_and_more.py new file mode 100644 index 00000000..913fec20 --- /dev/null +++ b/ddpui/migrations/0104_remove_userpreferences_discord_webhook_and_more.py @@ -0,0 +1,20 @@ +# Generated by Django 4.2 on 2024-11-06 15:34 + +from django.db import migrations + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0103_connectionjob_connectionmeta_and_more"), + ] + + operations = [ + migrations.RemoveField( + model_name="userpreferences", + name="discord_webhook", + ), + migrations.RemoveField( + model_name="userpreferences", + name="enable_discord_notifications", + ), + ] diff --git a/ddpui/migrations/0105_orgpreferences.py b/ddpui/migrations/0105_orgpreferences.py new file mode 100644 index 00000000..85761f5f --- /dev/null +++ b/ddpui/migrations/0105_orgpreferences.py @@ -0,0 +1,51 @@ +# Generated by Django 4.2 on 2024-11-06 15:41 + +from django.db import migrations, models +import django.db.models.deletion +import django.utils.timezone + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0104_remove_userpreferences_discord_webhook_and_more"), + ] + + operations = [ + migrations.CreateModel( + name="OrgPreferences", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, primary_key=True, serialize=False, verbose_name="ID" + ), + ), + ("trial_start_date", models.DateTimeField(blank=True, default=None, null=True)), + ("trial_end_date", models.DateTimeField(blank=True, default=None, null=True)), + ("llm_optin", models.BooleanField(default=False)), + ("llm_optin_date", models.DateTimeField(blank=True, null=True)), + ("enable_discord_notifications", models.BooleanField(default=False)), + ("discord_webhook", models.URLField(blank=True, null=True)), + ("created_at", models.DateTimeField(default=django.utils.timezone.now)), + ("updated_at", models.DateTimeField(default=django.utils.timezone.now)), + ( + "llm_optin_approved_by", + models.OneToOneField( + blank=True, + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="approvedby", + to="ddpui.orguser", + ), + ), + ( + "org", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + related_name="preferences", + to="ddpui.org", + ), + ), + ], + ), + ] diff --git a/ddpui/migrations/0106_alter_orgpreferences_llm_optin_approved_by.py b/ddpui/migrations/0106_alter_orgpreferences_llm_optin_approved_by.py new file mode 100644 index 00000000..03b0847d --- /dev/null +++ b/ddpui/migrations/0106_alter_orgpreferences_llm_optin_approved_by.py @@ -0,0 +1,24 @@ +# Generated by Django 4.2 on 2024-11-12 08:27 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0105_orgpreferences"), + ] + + operations = [ + migrations.AlterField( + model_name="orgpreferences", + name="llm_optin_approved_by", + field=models.ForeignKey( + blank=True, + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="approvedby", + to="ddpui.orguser", + ), + ), + ] diff --git a/ddpui/migrations/0107_orgsupersets.py b/ddpui/migrations/0107_orgsupersets.py new file mode 100644 index 00000000..845aa262 --- /dev/null +++ b/ddpui/migrations/0107_orgsupersets.py @@ -0,0 +1,37 @@ +# Generated by Django 4.2 on 2024-11-12 17:21 + +from django.db import migrations, models +import django.db.models.deletion +import django.utils.timezone + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0106_alter_orgpreferences_llm_optin_approved_by"), + ] + + operations = [ + migrations.CreateModel( + name="OrgSupersets", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, primary_key=True, serialize=False, verbose_name="ID" + ), + ), + ("container_name", models.CharField(blank=True, max_length=255, null=True)), + ("superset_version", models.CharField(blank=True, max_length=255, null=True)), + ("created_at", models.DateTimeField(default=django.utils.timezone.now)), + ("updated_at", models.DateTimeField(default=django.utils.timezone.now)), + ( + "org", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + related_name="orgInfo", + to="ddpui.org", + ), + ), + ], + ), + ] diff --git a/ddpui/migrations/0108_userpreferences_llm_optin.py b/ddpui/migrations/0108_userpreferences_llm_optin.py new file mode 100644 index 00000000..f7950b8c --- /dev/null +++ b/ddpui/migrations/0108_userpreferences_llm_optin.py @@ -0,0 +1,17 @@ +# Generated by Django 4.2 on 2024-11-12 19:51 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0107_orgsupersets"), + ] + + operations = [ + migrations.AddField( + model_name="userpreferences", + name="llm_optin", + field=models.BooleanField(default=False), + ), + ] diff --git a/ddpui/migrations/0109_alter_orgpreferences_org_alter_orgsupersets_org.py b/ddpui/migrations/0109_alter_orgpreferences_org_alter_orgsupersets_org.py new file mode 100644 index 00000000..7edb6434 --- /dev/null +++ b/ddpui/migrations/0109_alter_orgpreferences_org_alter_orgsupersets_org.py @@ -0,0 +1,29 @@ +# Generated by Django 4.2 on 2024-11-12 20:03 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0108_userpreferences_llm_optin"), + ] + + operations = [ + migrations.AlterField( + model_name="orgpreferences", + name="org", + field=models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, + related_name="preferences", + to="ddpui.org", + ), + ), + migrations.AlterField( + model_name="orgsupersets", + name="org", + field=models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, related_name="orgInfo", to="ddpui.org" + ), + ), + ] diff --git a/ddpui/models/__init__.py b/ddpui/models/__init__.py index c0757144..7608b502 100644 --- a/ddpui/models/__init__.py +++ b/ddpui/models/__init__.py @@ -1,3 +1,5 @@ from ddpui.models.tasks import Task from ddpui.models.llm import AssistantPrompt from ddpui.models.syncstats import SyncStats +from ddpui.models.org_preferences import OrgPreferences +from ddpui.models.org_supersets import OrgSupersets \ No newline at end of file diff --git a/ddpui/models/org_preferences.py b/ddpui/models/org_preferences.py new file mode 100644 index 00000000..6b4caea9 --- /dev/null +++ b/ddpui/models/org_preferences.py @@ -0,0 +1,32 @@ +from django.db import models +from enum import Enum +from ddpui.models.org import Org +from ddpui.models.org_user import OrgUser +from django.utils import timezone +from django.core.exceptions import ValidationError +class OrgType(str, Enum): + """an enum representing the type of organization""" + + DEMO = "demo" + POC = "poc" + CLIENT = "client" + + @classmethod + def choices(cls): + """django model definition needs an iterable for `choices`""" + return [(key.value, key.name) for key in cls] +class OrgPreferences(models.Model): + """Model to store org preferences for settings panel""" + + org = models.OneToOneField(Org, on_delete=models.CASCADE, related_name="preferences") + trial_start_date = models.DateTimeField(null=True, blank=True, default=None) + trial_end_date = models.DateTimeField(null=True, blank=True, default=None) + llm_optin = models.BooleanField(default=False) + llm_optin_approved_by = models.ForeignKey( + OrgUser, on_delete=models.CASCADE, related_name="approvedby", null=True, blank=True + ) + llm_optin_date = models.DateTimeField(null=True, blank=True) + enable_discord_notifications = models.BooleanField(default=False) + discord_webhook = models.URLField(blank=True, null=True) + created_at = models.DateTimeField(default=timezone.now) + updated_at = models.DateTimeField(default=timezone.now) diff --git a/ddpui/models/org_supersets.py b/ddpui/models/org_supersets.py new file mode 100644 index 00000000..f60c3441 --- /dev/null +++ b/ddpui/models/org_supersets.py @@ -0,0 +1,13 @@ +from django.db import models +from ddpui.models.org import Org +from django.utils import timezone + + +class OrgSupersets(models.Model): + """Model to store org supereset details for settings panel""" + + org = models.OneToOneField(Org, on_delete=models.CASCADE, related_name="orgInfo") + container_name = models.CharField(max_length=255, blank=True, null=True) + superset_version = models.CharField(max_length=255, blank=True, null=True) + created_at = models.DateTimeField(default=timezone.now) + updated_at = models.DateTimeField(default=timezone.now) diff --git a/ddpui/models/userpreferences.py b/ddpui/models/userpreferences.py index 5e2b84b8..320309f1 100644 --- a/ddpui/models/userpreferences.py +++ b/ddpui/models/userpreferences.py @@ -7,8 +7,7 @@ class UserPreferences(models.Model): """Model to store user preferences for notifications""" orguser = models.OneToOneField(OrgUser, on_delete=models.CASCADE, related_name="preferences") - enable_discord_notifications = models.BooleanField(default=False) enable_email_notifications = models.BooleanField(default=False) - discord_webhook = models.URLField(blank=True, null=True) + llm_optin = models.BooleanField(default=False) created_at = models.DateTimeField(default=timezone.now) updated_at = models.DateTimeField(default=timezone.now) diff --git a/ddpui/routes.py b/ddpui/routes.py index d61d8e5e..a5bc6a19 100644 --- a/ddpui/routes.py +++ b/ddpui/routes.py @@ -15,6 +15,7 @@ from ddpui.api.transform_api import transform_router from ddpui.api.user_org_api import user_org_router from ddpui.api.user_preferences_api import userpreference_router +from ddpui.api.org_preferences_api import orgpreference_router from ddpui.api.warehouse_api import warehouse_router from ddpui.api.webhook_api import webhook_router @@ -71,6 +72,7 @@ def ninja_default_error_handler( transform_router.tags = ["Transform"] user_org_router.tags = ["UserOrg"] userpreference_router.tags = ["UserPreference"] +orgpreference_router.tags =["OrgPreference"] warehouse_router.tags = ["Warehouse"] webhook_router.tags = ["Webhook"] @@ -89,3 +91,4 @@ def ninja_default_error_handler( src_api.add_router("/api/warehouse/", warehouse_router) src_api.add_router("/api/", user_org_router) src_api.add_router("/webhooks/", webhook_router) +src_api.add_router("/api/orgpreferences/",orgpreference_router ) \ No newline at end of file diff --git a/ddpui/schemas/org_preferences_schema.py b/ddpui/schemas/org_preferences_schema.py new file mode 100644 index 00000000..e17c17e1 --- /dev/null +++ b/ddpui/schemas/org_preferences_schema.py @@ -0,0 +1,42 @@ +from typing import Optional +from datetime import datetime +from ninja import Schema + + +class CreateOrgPreferencesSchema(Schema): + """Schema for creating organization preferences.""" + + org_id: int + trial_start_date: Optional[datetime] + trial_end_date: Optional[datetime] + llm_optin: Optional[bool] = False + llm_optin_approved_by: Optional[int] + llm_optin_date: Optional[datetime] + enable_discord_notifications: Optional[bool] = False + discord_webhook: Optional[str] + + +class UpdateOrgPreferencesSchema(Schema): + """Schema for updating organization preferences.""" + + trial_start_date: Optional[datetime] = None + trial_end_date: Optional[datetime] = None + llm_optin: Optional[bool] = None + llm_optin_approved_by: Optional[int] = None + llm_optin_date: Optional[datetime] = None + enable_discord_notifications: Optional[bool] = None + discord_webhook: Optional[str] = None + +class UpdateLLMOptinSchema(Schema): + """Schema for updating organization LLM approval preference.""" + llm_optin:bool + +class UpdateDiscordNotificationsSchema(Schema): + """Schema for updating organization discord notification settings.""" + enable_discord_notifications: bool + discord_webhook: Optional[str] + +class CreateOrgSupersetDetailsSchema(Schema): + """Schema for creating organization superset details.""" + superset_version: Optional[str] + container_name: Optional[str] \ No newline at end of file diff --git a/ddpui/schemas/userpreferences_schema.py b/ddpui/schemas/userpreferences_schema.py index e9009a40..108b89c4 100644 --- a/ddpui/schemas/userpreferences_schema.py +++ b/ddpui/schemas/userpreferences_schema.py @@ -4,15 +4,11 @@ class CreateUserPreferencesSchema(Schema): """Schema for creating user preferences for the user.""" - - enable_discord_notifications: bool enable_email_notifications: bool - discord_webhook: Optional[str] = None + llm_optin: bool class UpdateUserPreferencesSchema(Schema): """Schema for updating user preferences for the user.""" - - enable_discord_notifications: Optional[bool] = None enable_email_notifications: Optional[bool] = None - discord_webhook: Optional[str] = None + llm_optin: Optional[bool] diff --git a/ddpui/settings.py b/ddpui/settings.py index c6243007..846592e5 100644 --- a/ddpui/settings.py +++ b/ddpui/settings.py @@ -22,16 +22,16 @@ load_dotenv() -sentry_sdk.init( - dsn=os.getenv("SENTRY_DSN"), - # Set traces_sample_rate to 1.0 to capture 100% - # of transactions for performance monitoring. - traces_sample_rate=float(os.getenv("SENTRY_TSR", "1.0")), - # Set profiles_sample_rate to 1.0 to profile 100% - # of sampled transactions. - # We recommend adjusting this value in production. - profiles_sample_rate=float(os.getenv("SENTRY_PSR", "1.0")), -) +# sentry_sdk.init( +# dsn=os.getenv("SENTRY_DSN"), +# # Set traces_sample_rate to 1.0 to capture 100% +# # of transactions for performance monitoring. +# traces_sample_rate=float(os.getenv("SENTRY_TSR", "1.0")), +# # Set profiles_sample_rate to 1.0 to profile 100% +# # of sampled transactions. +# # We recommend adjusting this value in production. +# profiles_sample_rate=float(os.getenv("SENTRY_PSR", "1.0")), +# ) # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = Path(__file__).resolve().parent.parent diff --git a/ddpui/utils/constants.py b/ddpui/utils/constants.py index f72ee0eb..ea6538d7 100644 --- a/ddpui/utils/constants.py +++ b/ddpui/utils/constants.py @@ -51,3 +51,8 @@ # LLM data analysis ; row limit fetched and sent to llm service LIMIT_ROWS_TO_SEND_TO_LLM = 500 + +AIRBYTE_URL_TO_GET_VERSION ='http://localhost:8000/api/v1/instance_configuration' +PREFECT_URL_TO_GET_VERSION='http://localhost:4200/api/admin/version' +DBT_VERSION_COMMAND =["/home/ddp/dbt/venv/bin/dbt", "--version"] +ELEMENTARY_VERSION_COMMAND=["/home/ddp/dbt/venv/bin/edr", "--version"] \ No newline at end of file diff --git a/seed/002_permissions.json b/seed/002_permissions.json index c5cc5086..01122ce4 100644 --- a/seed/002_permissions.json +++ b/seed/002_permissions.json @@ -398,5 +398,21 @@ "name": "Can view feature flags of an organization", "slug": "can_view_flags" } + }, + { + "model": "ddpui.Permission", + "pk": 60, + "fields": { + "name": "Can Edit LLM Settings", + "slug": "can_edit_llm_settings" + } + }, + { + "model": "ddpui.Permission", + "pk": 61, + "fields": { + "name": "Can Edit Discord Notifications Settings", + "slug": "can_edit_discord_notifications_settings" + } } ] diff --git a/seed/003_role_permissions.json b/seed/003_role_permissions.json index 961fd4a4..72ff4a03 100644 --- a/seed/003_role_permissions.json +++ b/seed/003_role_permissions.json @@ -1774,5 +1774,37 @@ "role": 2, "permission": 58 } + }, + { + "model": "ddpui.RolePermission", + "pk": 219, + "fields": { + "role": 2, + "permission": 60 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 220, + "fields": { + "role": 1, + "permission": 60 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 221, + "fields": { + "role": 2, + "permission": 61 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 222, + "fields": { + "role": 1, + "permission": 61 + } } ] \ No newline at end of file From 87967884ce038d3f292ca1551503b76617d90511 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 13 Nov 2024 01:50:27 +0530 Subject: [PATCH 02/57] changed the settings --- ddpui/settings.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/ddpui/settings.py b/ddpui/settings.py index 846592e5..c6243007 100644 --- a/ddpui/settings.py +++ b/ddpui/settings.py @@ -22,16 +22,16 @@ load_dotenv() -# sentry_sdk.init( -# dsn=os.getenv("SENTRY_DSN"), -# # Set traces_sample_rate to 1.0 to capture 100% -# # of transactions for performance monitoring. -# traces_sample_rate=float(os.getenv("SENTRY_TSR", "1.0")), -# # Set profiles_sample_rate to 1.0 to profile 100% -# # of sampled transactions. -# # We recommend adjusting this value in production. -# profiles_sample_rate=float(os.getenv("SENTRY_PSR", "1.0")), -# ) +sentry_sdk.init( + dsn=os.getenv("SENTRY_DSN"), + # Set traces_sample_rate to 1.0 to capture 100% + # of transactions for performance monitoring. + traces_sample_rate=float(os.getenv("SENTRY_TSR", "1.0")), + # Set profiles_sample_rate to 1.0 to profile 100% + # of sampled transactions. + # We recommend adjusting this value in production. + profiles_sample_rate=float(os.getenv("SENTRY_PSR", "1.0")), +) # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = Path(__file__).resolve().parent.parent From b291f815953fc534110e4d03eb701e7a61370fb8 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 13 Nov 2024 01:59:50 +0530 Subject: [PATCH 03/57] preformatted --- ddpui/api/org_preferences_api.py | 193 ++++++++++++------------ ddpui/api/user_preferences_api.py | 5 +- ddpui/models/__init__.py | 2 +- ddpui/models/org_preferences.py | 4 + ddpui/routes.py | 4 +- ddpui/schemas/org_preferences_schema.py | 18 ++- ddpui/schemas/userpreferences_schema.py | 2 + ddpui/utils/constants.py | 8 +- 8 files changed, 121 insertions(+), 115 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 48d33482..e0c50829 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -9,21 +9,21 @@ UpdateOrgPreferencesSchema, UpdateLLMOptinSchema, UpdateDiscordNotificationsSchema, - CreateOrgSupersetDetailsSchema + CreateOrgSupersetDetailsSchema, ) from ddpui.utils.constants import ( AIRBYTE_URL_TO_GET_VERSION, PREFECT_URL_TO_GET_VERSION, DBT_VERSION_COMMAND, - ELEMENTARY_VERSION_COMMAND + ELEMENTARY_VERSION_COMMAND, ) from ddpui.auth import has_permission from ddpui.models.org import Org from ddpui.models.org_user import OrgUser import requests import subprocess -orgpreference_router = Router() +orgpreference_router = Router() @orgpreference_router.post("/", auth=auth.CustomAuthMiddleware()) @@ -33,9 +33,7 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): if OrgPreferences.objects.filter(org_id=payload.org_id).exists(): raise HttpError(400, "Organization preferences already exist") - org_preferences = OrgPreferences.objects.create( - **payload.dict() - ) + org_preferences = OrgPreferences.objects.create(**payload.dict()) preferences = { "trial_start_date": org_preferences.trial_start_date, @@ -60,15 +58,15 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): # raise HttpError(404, "Preferences for this organization not found") # #check this line please. -# # if org.type=="demo": +# # if org.type=="demo": # # if payload.trial_start_date is None or payload.trial_end_date is None: # # raise HttpError(400, "Trial accounts must have trial_start_date and trial_end_date.") - + # if payload.llm_optin is True: # # if not request.orguser.has_permission("can_edit_llm_settings"): # # raise HttpError(403, "You do not have permission to edit LLM settings.") - + # payload.llm_optin_date = timezone.now() # if payload.llm_optin_approved_by is None or payload.llm_optin_date is None: # raise HttpError(400, "If LLM opt-in is true, both llm_optin_approved_by and llm_optin_date must be provided.") @@ -79,8 +77,9 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): # return {"success": True, "res": 1} + @orgpreference_router.put("/{org_id}/llm_approval", auth=auth.CustomAuthMiddleware()) -@has_permission(['can_edit_llm_settings']) +@has_permission(["can_edit_llm_settings"]) def update_org_preferences(request, payload: UpdateLLMOptinSchema, org_id: int): """Updates llm preferences for an organization""" @@ -106,8 +105,11 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema, org_id: int): return {"success": True, "res": 1} -@orgpreference_router.put("/{org_id}/enable-discord-notifications", auth=auth.CustomAuthMiddleware()) -@has_permission(['can_edit_discord_notifications_settings']) + +@orgpreference_router.put( + "/{org_id}/enable-discord-notifications", auth=auth.CustomAuthMiddleware() +) +@has_permission(["can_edit_discord_notifications_settings"]) def update_org_preferences(request, payload: UpdateDiscordNotificationsSchema, org_id: int): """Updates preferences for an organization""" @@ -115,42 +117,37 @@ def update_org_preferences(request, payload: UpdateDiscordNotificationsSchema, o org_preferences = OrgPreferences.objects.get(org_id=org_id) except OrgPreferences.DoesNotExist: raise HttpError(404, "Preferences for this organization not found") - + # Check for discord_link in the database when enabling notifications if payload.enable_discord_notifications: if not org_preferences.discord_webhook and not payload.discord_webhook: - raise HttpError(400, "Discord link is missing. Please add a Discord link to enable notifications.") + raise HttpError( + 400, "Discord link is missing. Please add a Discord link to enable notifications." + ) # Update the OrgPreferences instance org_preferences.enable_discord_notifications = payload.enable_discord_notifications if payload.discord_webhook: - org_preferences.discord_webhook =payload.discord_webhook + org_preferences.discord_webhook = payload.discord_webhook org_preferences.save() return {"success": True, "res": 1} - - - @orgpreference_router.get("/{org_id}/", auth=auth.CustomAuthMiddleware()) def get_org_preferences(request, org_id: int): """Gets preferences for an organization by org_id""" try: - org_preferences = OrgPreferences.objects.select_related( - "org", - "llm_optin_approved_by" - ).get(org_id=org_id) + org_preferences = OrgPreferences.objects.select_related("org", "llm_optin_approved_by").get( + org_id=org_id + ) except OrgPreferences.DoesNotExist: raise HttpError(404, "Organization preferences not found") - + approved_by_data = None if org_preferences.llm_optin_approved_by and org_preferences.llm_optin_approved_by.user: - user = org_preferences.llm_optin_approved_by.user - approved_by_data = { - "user_id": user.pk, - "email": user.email - } + user = org_preferences.llm_optin_approved_by.user + approved_by_data = {"user_id": user.pk, "email": user.email} preferences = { "org": { @@ -162,7 +159,7 @@ def get_org_preferences(request, org_id: int): "trial_start_date": org_preferences.trial_start_date, "trial_end_date": org_preferences.trial_end_date, "llm_optin": org_preferences.llm_optin, - "llm_optin_approved_by": approved_by_data, + "llm_optin_approved_by": approved_by_data, "llm_optin_date": org_preferences.llm_optin_date, "enable_discord_notifications": org_preferences.enable_discord_notifications, "discord_webhook": org_preferences.discord_webhook, @@ -170,18 +167,19 @@ def get_org_preferences(request, org_id: int): return {"success": True, "res": preferences} + # api to get more superset related information. @orgpreference_router.get("/{org_id}/org-superset", auth=auth.CustomAuthMiddleware()) def get_org_superset_details(request, org_id: int): """Gets preferences for an organization by org_id""" try: org_superset = OrgSupersets.objects.select_related( - "org", - ).get(org_id=org_id) - + "org", + ).get(org_id=org_id) + except OrgSupersets.DoesNotExist: raise HttpError(404, "Organizations superset details not found") - + response_data = { "org": { "id": org_superset.org.id, @@ -189,24 +187,20 @@ def get_org_superset_details(request, org_id: int): "slug": org_superset.org.slug, "type": org_superset.org.type, }, - "superset_version": org_superset.superset_version + "superset_version": org_superset.superset_version, } return {"success": True, "res": response_data} - @orgpreference_router.post("/{org_id}/org-superset", auth=auth.CustomAuthMiddleware()) -def create_org_superset_details(request, payload:CreateOrgSupersetDetailsSchema, org_id: int): +def create_org_superset_details(request, payload: CreateOrgSupersetDetailsSchema, org_id: int): """Creates an entry in db with org superset details""" try: org = Org.objects.get(id=org_id) except Org.DoesNotExist: raise HttpError(404, "Organization not found") - org_superset, created = OrgSupersets.objects.update_or_create( - org=org, - **payload.dict() - ) + org_superset, created = OrgSupersets.objects.update_or_create(org=org, **payload.dict()) response_data = { "org": { @@ -215,67 +209,68 @@ def create_org_superset_details(request, payload:CreateOrgSupersetDetailsSchema, "slug": org.slug, "type": org.type, }, - "id": org_superset.id, - "superset_version": org_superset.superset_version + "id": org_superset.id, + "superset_version": org_superset.superset_version, } return {"success": True, "res": response_data} + @orgpreference_router.get("/{org_id}/versions", auth=auth.CustomAuthMiddleware()) -def get_tools_versions( request, org_id: int): - versions = {} - - # Airbyte Version - try: - airbyte_url = AIRBYTE_URL_TO_GET_VERSION - airbyte_response = requests.get(airbyte_url) - if airbyte_response.status_code == 200: - airbyte_data = airbyte_response.json() - versions["Airbyte"] = airbyte_data.get("version") - else: - versions["Airbyte"] = "Error: Unable to retrieve version" - except Exception as e: - versions["Airbyte"] = f"Error: {e}" - - # Prefect Version - try: - prefect_url =PREFECT_URL_TO_GET_VERSION - prefect_response = requests.get(prefect_url) - if prefect_response.status_code == 200: - versions["Prefect"] = prefect_response.text.strip() - else: - versions["Prefect"] = "Error: Unable to retrieve version" - except Exception as e: - versions["Prefect"] = f"Error: {e}" - - # dbt Version - try: - dbt_version_command = DBT_VERSION_COMMAND - dbt_output = subprocess.check_output(dbt_version_command, text=True) - for line in dbt_output.splitlines(): - if "installed:" in line: - versions["dbt"] = line.split(":")[1].strip() - break - except Exception as e: - versions["dbt"] = f"Error: {e}" - - # Elementary Version - try: - elementary_version_command = ELEMENTARY_VERSION_COMMAND - elementary_output = subprocess.check_output(elementary_version_command, text=True) - for line in elementary_output.splitlines(): - if "Elementary version" in line: - versions["Elementary"] = line.split()[-1].strip() - break - except Exception as e: - versions["Elementary"] = f"Error: {e}" - - # Superset_version - try: - org_superset = OrgSupersets.objects.get(org_id=org_id) - except OrgSupersets.DoesNotExist: - raise HttpError(404, "Organizations superset details not found") - - versions["Superset"] = org_superset.superset_version - - return {"success": True, "res": versions} \ No newline at end of file +def get_tools_versions(request, org_id: int): + versions = {} + + # Airbyte Version + try: + airbyte_url = AIRBYTE_URL_TO_GET_VERSION + airbyte_response = requests.get(airbyte_url) + if airbyte_response.status_code == 200: + airbyte_data = airbyte_response.json() + versions["Airbyte"] = airbyte_data.get("version") + else: + versions["Airbyte"] = "Error: Unable to retrieve version" + except Exception as e: + versions["Airbyte"] = f"Error: {e}" + + # Prefect Version + try: + prefect_url = PREFECT_URL_TO_GET_VERSION + prefect_response = requests.get(prefect_url) + if prefect_response.status_code == 200: + versions["Prefect"] = prefect_response.text.strip() + else: + versions["Prefect"] = "Error: Unable to retrieve version" + except Exception as e: + versions["Prefect"] = f"Error: {e}" + + # dbt Version + try: + dbt_version_command = DBT_VERSION_COMMAND + dbt_output = subprocess.check_output(dbt_version_command, text=True) + for line in dbt_output.splitlines(): + if "installed:" in line: + versions["dbt"] = line.split(":")[1].strip() + break + except Exception as e: + versions["dbt"] = f"Error: {e}" + + # Elementary Version + try: + elementary_version_command = ELEMENTARY_VERSION_COMMAND + elementary_output = subprocess.check_output(elementary_version_command, text=True) + for line in elementary_output.splitlines(): + if "Elementary version" in line: + versions["Elementary"] = line.split()[-1].strip() + break + except Exception as e: + versions["Elementary"] = f"Error: {e}" + + # Superset_version + try: + org_superset = OrgSupersets.objects.get(org_id=org_id) + except OrgSupersets.DoesNotExist: + raise HttpError(404, "Organizations superset details not found") + + versions["Superset"] = org_superset.superset_version + + return {"success": True, "res": versions} diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index de19ca59..a28b93dc 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -21,14 +21,13 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): userpreferences = UserPreferences.objects.create( orguser=orguser, - llm_optin = payload.llm_optin, + llm_optin=payload.llm_optin, enable_email_notifications=payload.enable_email_notifications, - ) preferences = { "enable_email_notifications": userpreferences.enable_email_notifications, - "llm_optin": userpreferences.llm_optin + "llm_optin": userpreferences.llm_optin, } return {"success": True, "res": preferences} diff --git a/ddpui/models/__init__.py b/ddpui/models/__init__.py index 7608b502..93c260b3 100644 --- a/ddpui/models/__init__.py +++ b/ddpui/models/__init__.py @@ -2,4 +2,4 @@ from ddpui.models.llm import AssistantPrompt from ddpui.models.syncstats import SyncStats from ddpui.models.org_preferences import OrgPreferences -from ddpui.models.org_supersets import OrgSupersets \ No newline at end of file +from ddpui.models.org_supersets import OrgSupersets diff --git a/ddpui/models/org_preferences.py b/ddpui/models/org_preferences.py index 6b4caea9..879f2fb1 100644 --- a/ddpui/models/org_preferences.py +++ b/ddpui/models/org_preferences.py @@ -4,6 +4,8 @@ from ddpui.models.org_user import OrgUser from django.utils import timezone from django.core.exceptions import ValidationError + + class OrgType(str, Enum): """an enum representing the type of organization""" @@ -15,6 +17,8 @@ class OrgType(str, Enum): def choices(cls): """django model definition needs an iterable for `choices`""" return [(key.value, key.name) for key in cls] + + class OrgPreferences(models.Model): """Model to store org preferences for settings panel""" diff --git a/ddpui/routes.py b/ddpui/routes.py index a5bc6a19..4322c986 100644 --- a/ddpui/routes.py +++ b/ddpui/routes.py @@ -72,7 +72,7 @@ def ninja_default_error_handler( transform_router.tags = ["Transform"] user_org_router.tags = ["UserOrg"] userpreference_router.tags = ["UserPreference"] -orgpreference_router.tags =["OrgPreference"] +orgpreference_router.tags = ["OrgPreference"] warehouse_router.tags = ["Warehouse"] webhook_router.tags = ["Webhook"] @@ -91,4 +91,4 @@ def ninja_default_error_handler( src_api.add_router("/api/warehouse/", warehouse_router) src_api.add_router("/api/", user_org_router) src_api.add_router("/webhooks/", webhook_router) -src_api.add_router("/api/orgpreferences/",orgpreference_router ) \ No newline at end of file +src_api.add_router("/api/orgpreferences/", orgpreference_router) diff --git a/ddpui/schemas/org_preferences_schema.py b/ddpui/schemas/org_preferences_schema.py index e17c17e1..25742777 100644 --- a/ddpui/schemas/org_preferences_schema.py +++ b/ddpui/schemas/org_preferences_schema.py @@ -6,12 +6,12 @@ class CreateOrgPreferencesSchema(Schema): """Schema for creating organization preferences.""" - org_id: int + org_id: int trial_start_date: Optional[datetime] - trial_end_date: Optional[datetime] + trial_end_date: Optional[datetime] llm_optin: Optional[bool] = False - llm_optin_approved_by: Optional[int] - llm_optin_date: Optional[datetime] + llm_optin_approved_by: Optional[int] + llm_optin_date: Optional[datetime] enable_discord_notifications: Optional[bool] = False discord_webhook: Optional[str] @@ -27,16 +27,22 @@ class UpdateOrgPreferencesSchema(Schema): enable_discord_notifications: Optional[bool] = None discord_webhook: Optional[str] = None + class UpdateLLMOptinSchema(Schema): """Schema for updating organization LLM approval preference.""" - llm_optin:bool + + llm_optin: bool + class UpdateDiscordNotificationsSchema(Schema): """Schema for updating organization discord notification settings.""" + enable_discord_notifications: bool discord_webhook: Optional[str] + class CreateOrgSupersetDetailsSchema(Schema): """Schema for creating organization superset details.""" + superset_version: Optional[str] - container_name: Optional[str] \ No newline at end of file + container_name: Optional[str] diff --git a/ddpui/schemas/userpreferences_schema.py b/ddpui/schemas/userpreferences_schema.py index 108b89c4..6d5b857c 100644 --- a/ddpui/schemas/userpreferences_schema.py +++ b/ddpui/schemas/userpreferences_schema.py @@ -4,11 +4,13 @@ class CreateUserPreferencesSchema(Schema): """Schema for creating user preferences for the user.""" + enable_email_notifications: bool llm_optin: bool class UpdateUserPreferencesSchema(Schema): """Schema for updating user preferences for the user.""" + enable_email_notifications: Optional[bool] = None llm_optin: Optional[bool] diff --git a/ddpui/utils/constants.py b/ddpui/utils/constants.py index ea6538d7..3798f988 100644 --- a/ddpui/utils/constants.py +++ b/ddpui/utils/constants.py @@ -52,7 +52,7 @@ # LLM data analysis ; row limit fetched and sent to llm service LIMIT_ROWS_TO_SEND_TO_LLM = 500 -AIRBYTE_URL_TO_GET_VERSION ='http://localhost:8000/api/v1/instance_configuration' -PREFECT_URL_TO_GET_VERSION='http://localhost:4200/api/admin/version' -DBT_VERSION_COMMAND =["/home/ddp/dbt/venv/bin/dbt", "--version"] -ELEMENTARY_VERSION_COMMAND=["/home/ddp/dbt/venv/bin/edr", "--version"] \ No newline at end of file +AIRBYTE_URL_TO_GET_VERSION = "http://localhost:8000/api/v1/instance_configuration" +PREFECT_URL_TO_GET_VERSION = "http://localhost:4200/api/admin/version" +DBT_VERSION_COMMAND = ["/home/ddp/dbt/venv/bin/dbt", "--version"] +ELEMENTARY_VERSION_COMMAND = ["/home/ddp/dbt/venv/bin/edr", "--version"] From eaa0382d6c3d6e395522e4b589440ab75023cd15 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Thu, 14 Nov 2024 14:20:51 +0530 Subject: [PATCH 04/57] fixed the apis --- ddpui/api/org_preferences_api.py | 156 +++++++++++------- ddpui/api/user_preferences_api.py | 13 +- ...serpreferences_discord_webhook_and_more.py | 22 +++ ddpui/models/userpreferences.py | 2 + ddpui/schemas/org_preferences_schema.py | 2 +- ddpui/schemas/userpreferences_schema.py | 4 + 6 files changed, 134 insertions(+), 65 deletions(-) create mode 100644 ddpui/migrations/0110_userpreferences_discord_webhook_and_more.py diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index e0c50829..cb2ca6fc 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -28,12 +28,19 @@ @orgpreference_router.post("/", auth=auth.CustomAuthMiddleware()) def create_org_preferences(request, payload: CreateOrgPreferencesSchema): - print(payload, "payload") + orguser: OrgUser = request.orguser + org = orguser.org + payload.org = org """Creates preferences for an organization""" - if OrgPreferences.objects.filter(org_id=payload.org_id).exists(): + if OrgPreferences.objects.filter(org=org).exists(): raise HttpError(400, "Organization preferences already exist") - org_preferences = OrgPreferences.objects.create(**payload.dict()) + payload_data = payload.dict(exclude={"org"}) + + # Create OrgPreferences + org_preferences = OrgPreferences.objects.create( + org=org, **payload_data # Use the rest of the payload + ) preferences = { "trial_start_date": org_preferences.trial_start_date, @@ -78,13 +85,18 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): # return {"success": True, "res": 1} -@orgpreference_router.put("/{org_id}/llm_approval", auth=auth.CustomAuthMiddleware()) +@orgpreference_router.put("/llm_approval", auth=auth.CustomAuthMiddleware()) @has_permission(["can_edit_llm_settings"]) -def update_org_preferences(request, payload: UpdateLLMOptinSchema, org_id: int): - """Updates llm preferences for an organization""" +def update_org_preferences(request, payload: UpdateLLMOptinSchema): + """Updates llm preferences for the logged-in user's organization""" + + # Get the organization associated with the current user + orguser: OrgUser = request.orguser + org = orguser.org try: - org_preferences = OrgPreferences.objects.get(org_id=org_id) + # Fetch organization preferences for the user's organization + org_preferences = OrgPreferences.objects.get(org=org) except OrgPreferences.DoesNotExist: raise HttpError(404, "Preferences for this organization not found") @@ -106,49 +118,59 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema, org_id: int): return {"success": True, "res": 1} -@orgpreference_router.put( - "/{org_id}/enable-discord-notifications", auth=auth.CustomAuthMiddleware() -) +@orgpreference_router.put("/enable-discord-notifications", auth=auth.CustomAuthMiddleware()) @has_permission(["can_edit_discord_notifications_settings"]) -def update_org_preferences(request, payload: UpdateDiscordNotificationsSchema, org_id: int): - """Updates preferences for an organization""" +def update_discord_notifications(request, payload: UpdateDiscordNotificationsSchema): + """Updates Discord notifications preferences for the logged-in user's organization.""" + + orguser: OrgUser = request.orguser + org = orguser.org try: - org_preferences = OrgPreferences.objects.get(org_id=org_id) + org_preferences = OrgPreferences.objects.get(org=org) except OrgPreferences.DoesNotExist: raise HttpError(404, "Preferences for this organization not found") - # Check for discord_link in the database when enabling notifications if payload.enable_discord_notifications: if not org_preferences.discord_webhook and not payload.discord_webhook: - raise HttpError( - 400, "Discord link is missing. Please add a Discord link to enable notifications." - ) + raise HttpError(400, "Discord webhook is required to enable notifications.") + discord_webhook = payload.discord_webhook or org_preferences.discord_webhook + else: + discord_webhook = None - # Update the OrgPreferences instance org_preferences.enable_discord_notifications = payload.enable_discord_notifications - if payload.discord_webhook: - org_preferences.discord_webhook = payload.discord_webhook + org_preferences.discord_webhook = discord_webhook org_preferences.save() - return {"success": True, "res": 1} + response_data = { + "enable_discord_notifications": org_preferences.enable_discord_notifications, + "discord_webhook": org_preferences.discord_webhook, + } + return {"success": True, "res": response_data} + + +@orgpreference_router.get("/", auth=auth.CustomAuthMiddleware()) +def get_org_preferences(request): + """Gets preferences for an organization based on the logged-in user's organization""" + orguser: OrgUser = request.orguser + org = orguser.org # Get the organization associated with the current user -@orgpreference_router.get("/{org_id}/", auth=auth.CustomAuthMiddleware()) -def get_org_preferences(request, org_id: int): - """Gets preferences for an organization by org_id""" try: + # Fetch organization preferences for the user's organization org_preferences = OrgPreferences.objects.select_related("org", "llm_optin_approved_by").get( - org_id=org_id + org=org ) except OrgPreferences.DoesNotExist: raise HttpError(404, "Organization preferences not found") + # Prepare data for the user who approved LLM opt-in, if available approved_by_data = None if org_preferences.llm_optin_approved_by and org_preferences.llm_optin_approved_by.user: user = org_preferences.llm_optin_approved_by.user approved_by_data = {"user_id": user.pk, "email": user.email} + # Structure the preferences response preferences = { "org": { "id": org_preferences.org.id, @@ -168,14 +190,16 @@ def get_org_preferences(request, org_id: int): return {"success": True, "res": preferences} -# api to get more superset related information. -@orgpreference_router.get("/{org_id}/org-superset", auth=auth.CustomAuthMiddleware()) -def get_org_superset_details(request, org_id: int): +@orgpreference_router.get("/org-superset", auth=auth.CustomAuthMiddleware()) +def get_org_superset_details(request): """Gets preferences for an organization by org_id""" + orguser: OrgUser = request.orguser + org = orguser.org + try: org_superset = OrgSupersets.objects.select_related( "org", - ).get(org_id=org_id) + ).get(org=org) except OrgSupersets.DoesNotExist: raise HttpError(404, "Organizations superset details not found") @@ -192,33 +216,39 @@ def get_org_superset_details(request, org_id: int): return {"success": True, "res": response_data} -@orgpreference_router.post("/{org_id}/org-superset", auth=auth.CustomAuthMiddleware()) -def create_org_superset_details(request, payload: CreateOrgSupersetDetailsSchema, org_id: int): - """Creates an entry in db with org superset details""" - try: - org = Org.objects.get(id=org_id) - except Org.DoesNotExist: - raise HttpError(404, "Organization not found") +@orgpreference_router.post("/org-superset", auth=auth.CustomAuthMiddleware()) +def create_org_superset_details(request, payload: CreateOrgSupersetDetailsSchema): + """Creates or updates an entry in db with org superset details.""" + orguser: OrgUser = request.orguser + org = orguser.org - org_superset, created = OrgSupersets.objects.update_or_create(org=org, **payload.dict()) + org_superset, created = OrgSupersets.objects.update_or_create(org=org, defaults=payload.dict()) response_data = { "org": { - "id": org.id, "name": org.name, "slug": org.slug, "type": org.type, }, - "id": org_superset.id, "superset_version": org_superset.superset_version, + "created": created, } return {"success": True, "res": response_data} -@orgpreference_router.get("/{org_id}/versions", auth=auth.CustomAuthMiddleware()) -def get_tools_versions(request, org_id: int): - versions = {} +@orgpreference_router.get("/toolinfo", auth=auth.CustomAuthMiddleware()) +def get_tools_versions(request): + orguser: OrgUser = request.orguser + org = orguser.org + + try: + # Validate org_id + org_superset = OrgSupersets.objects.get(org=org) + except OrgSupersets.DoesNotExist: + return {"success": False, "error": "Organization not found"} + + versions = [] # Airbyte Version try: @@ -226,22 +256,23 @@ def get_tools_versions(request, org_id: int): airbyte_response = requests.get(airbyte_url) if airbyte_response.status_code == 200: airbyte_data = airbyte_response.json() - versions["Airbyte"] = airbyte_data.get("version") + versions.append({"Airbyte": {"version": airbyte_data.get("version")}}) else: - versions["Airbyte"] = "Error: Unable to retrieve version" - except Exception as e: - versions["Airbyte"] = f"Error: {e}" + versions.append({"Airbyte": {"version": "Not available"}}) + except Exception: + versions.append({"Airbyte": {"version": "Not available"}}) # Prefect Version try: prefect_url = PREFECT_URL_TO_GET_VERSION prefect_response = requests.get(prefect_url) if prefect_response.status_code == 200: - versions["Prefect"] = prefect_response.text.strip() + version = prefect_response.text.strip().strip('"') + versions.append({"Prefect": {"version": version}}) else: - versions["Prefect"] = "Error: Unable to retrieve version" - except Exception as e: - versions["Prefect"] = f"Error: {e}" + versions.append({"Prefect": {"version": "Not available"}}) + except Exception: + versions.append({"Prefect": {"version": "Not available"}}) # dbt Version try: @@ -249,10 +280,12 @@ def get_tools_versions(request, org_id: int): dbt_output = subprocess.check_output(dbt_version_command, text=True) for line in dbt_output.splitlines(): if "installed:" in line: - versions["dbt"] = line.split(":")[1].strip() + versions.append({"DBT": {"version": line.split(":")[1].strip()}}) break - except Exception as e: - versions["dbt"] = f"Error: {e}" + else: + versions.append({"DBT": {"version": "Not available"}}) + except Exception: + versions.append({"DBT": {"version": "Not available"}}) # Elementary Version try: @@ -260,17 +293,14 @@ def get_tools_versions(request, org_id: int): elementary_output = subprocess.check_output(elementary_version_command, text=True) for line in elementary_output.splitlines(): if "Elementary version" in line: - versions["Elementary"] = line.split()[-1].strip() + versions.append({"Elementary": {"version": line.split()[-1].strip()}}) break - except Exception as e: - versions["Elementary"] = f"Error: {e}" - - # Superset_version - try: - org_superset = OrgSupersets.objects.get(org_id=org_id) - except OrgSupersets.DoesNotExist: - raise HttpError(404, "Organizations superset details not found") + else: + versions.append({"Elementary": {"version": "Not available"}}) + except Exception: + versions.append({"Elementary": {"version": "Not available"}}) - versions["Superset"] = org_superset.superset_version + # Superset Version + versions.append({"Superset": {"version": org_superset.superset_version}}) return {"success": True, "res": versions} diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index a28b93dc..f1796792 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -22,12 +22,16 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): userpreferences = UserPreferences.objects.create( orguser=orguser, llm_optin=payload.llm_optin, + enable_discord_notifications=payload.enable_discord_notifications, + discord_webhook=payload.discord_webhook, enable_email_notifications=payload.enable_email_notifications, ) preferences = { "enable_email_notifications": userpreferences.enable_email_notifications, "llm_optin": userpreferences.llm_optin, + "discord_webhook": userpreferences.discord_webhook, + "enable_discord_notifications": userpreferences.enable_discord_notifications, } return {"success": True, "res": preferences} @@ -40,14 +44,19 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) + if payload.enable_discord_notifications is not None: + user_preferences.enable_discord_notifications = payload.enable_discord_notifications if payload.llm_optin is not None: user_preferences.llm_optin = payload.llm_optin if payload.enable_email_notifications is not None: user_preferences.enable_email_notifications = payload.enable_email_notifications - + if payload.discord_webhook is not None: + user_preferences.discord_webhook = payload.discord_webhook user_preferences.save() preferences = { + "discord_webhook": user_preferences.discord_webhook, + "enable_discord_notifications": user_preferences.enable_discord_notifications, "enable_email_notifications": user_preferences.enable_email_notifications, "llm_optin": user_preferences.llm_optin, } @@ -63,6 +72,8 @@ def get_user_preferences(request): user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) preferences = { + "discord_webhook": user_preferences.discord_webhook, + "enable_discord_notifications": user_preferences.enable_discord_notifications, "enable_email_notifications": user_preferences.enable_email_notifications, "llm_optin": user_preferences.llm_optin, } diff --git a/ddpui/migrations/0110_userpreferences_discord_webhook_and_more.py b/ddpui/migrations/0110_userpreferences_discord_webhook_and_more.py new file mode 100644 index 00000000..27d5d0c1 --- /dev/null +++ b/ddpui/migrations/0110_userpreferences_discord_webhook_and_more.py @@ -0,0 +1,22 @@ +# Generated by Django 4.2 on 2024-11-14 08:49 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0109_alter_orgpreferences_org_alter_orgsupersets_org"), + ] + + operations = [ + migrations.AddField( + model_name="userpreferences", + name="discord_webhook", + field=models.URLField(blank=True, null=True), + ), + migrations.AddField( + model_name="userpreferences", + name="enable_discord_notifications", + field=models.BooleanField(default=False), + ), + ] diff --git a/ddpui/models/userpreferences.py b/ddpui/models/userpreferences.py index 320309f1..162b07f9 100644 --- a/ddpui/models/userpreferences.py +++ b/ddpui/models/userpreferences.py @@ -7,6 +7,8 @@ class UserPreferences(models.Model): """Model to store user preferences for notifications""" orguser = models.OneToOneField(OrgUser, on_delete=models.CASCADE, related_name="preferences") + enable_discord_notifications = models.BooleanField(default=False) + discord_webhook = models.URLField(blank=True, null=True) enable_email_notifications = models.BooleanField(default=False) llm_optin = models.BooleanField(default=False) created_at = models.DateTimeField(default=timezone.now) diff --git a/ddpui/schemas/org_preferences_schema.py b/ddpui/schemas/org_preferences_schema.py index 25742777..c28b52bc 100644 --- a/ddpui/schemas/org_preferences_schema.py +++ b/ddpui/schemas/org_preferences_schema.py @@ -6,7 +6,7 @@ class CreateOrgPreferencesSchema(Schema): """Schema for creating organization preferences.""" - org_id: int + org: Optional[int] trial_start_date: Optional[datetime] trial_end_date: Optional[datetime] llm_optin: Optional[bool] = False diff --git a/ddpui/schemas/userpreferences_schema.py b/ddpui/schemas/userpreferences_schema.py index 6d5b857c..f5c6c482 100644 --- a/ddpui/schemas/userpreferences_schema.py +++ b/ddpui/schemas/userpreferences_schema.py @@ -5,6 +5,8 @@ class CreateUserPreferencesSchema(Schema): """Schema for creating user preferences for the user.""" + enable_discord_notifications: bool + discord_webhook: Optional[str] = None enable_email_notifications: bool llm_optin: bool @@ -12,5 +14,7 @@ class CreateUserPreferencesSchema(Schema): class UpdateUserPreferencesSchema(Schema): """Schema for updating user preferences for the user.""" + enable_discord_notifications: Optional[bool] = None + discord_webhook: Optional[str] = None enable_email_notifications: Optional[bool] = None llm_optin: Optional[bool] From 8c2403279cc93be9004c0959d8173c02ee583cbb Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Thu, 14 Nov 2024 19:10:08 +0530 Subject: [PATCH 05/57] formatting fixed --- ddpui/api/org_preferences_api.py | 131 ++++++++++++++----------------- ddpui/api/user_org_api.py | 26 +++++- ddpui/core/orguserfunctions.py | 13 +++ ddpui/models/org.py | 4 +- ddpui/models/org_plans.py | 25 ++++++ ddpui/models/org_preferences.py | 17 ---- ddpui/models/org_user.py | 7 ++ 7 files changed, 128 insertions(+), 95 deletions(-) create mode 100644 ddpui/models/org_plans.py diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index cb2ca6fc..80e76108 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -3,6 +3,7 @@ from ddpui import auth from ddpui.models.org_preferences import OrgPreferences from ddpui.models.org_supersets import OrgSupersets +from ddpui.models.org_plans import OrgPlans from django.utils import timezone from ddpui.schemas.org_preferences_schema import ( CreateOrgPreferencesSchema, @@ -20,6 +21,11 @@ from ddpui.auth import has_permission from ddpui.models.org import Org from ddpui.models.org_user import OrgUser +from ddpui.ddpairbyte import airbyte_service +from ddpui.ddpprefect import ( + prefect_service, +) +from ddpui.core.orgdbt_manager import DbtProjectManager import requests import subprocess @@ -37,14 +43,11 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): payload_data = payload.dict(exclude={"org"}) - # Create OrgPreferences org_preferences = OrgPreferences.objects.create( org=org, **payload_data # Use the rest of the payload ) preferences = { - "trial_start_date": org_preferences.trial_start_date, - "trial_end_date": org_preferences.trial_end_date, "llm_optin": org_preferences.llm_optin, "llm_optin_approved_by": org_preferences.llm_optin_approved_by, "llm_optin_date": org_preferences.llm_optin_date, @@ -55,61 +58,25 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): return {"success": True, "res": preferences} -# @orgpreference_router.put("/{org_id}/", auth=auth.CustomAuthMiddleware()) -# def update_org_preferences(request, payload: UpdateOrgPreferencesSchema, org_id: int): -# """Updates preferences for an organization""" -# payload.llm_optin_approved_by = request.orguser -# try: -# org_preferences = OrgPreferences.objects.get(org_id=org_id) -# except OrgPreferences.DoesNotExist: -# raise HttpError(404, "Preferences for this organization not found") - -# #check this line please. -# # if org.type=="demo": -# # if payload.trial_start_date is None or payload.trial_end_date is None: -# # raise HttpError(400, "Trial accounts must have trial_start_date and trial_end_date.") - - -# if payload.llm_optin is True: -# # if not request.orguser.has_permission("can_edit_llm_settings"): -# # raise HttpError(403, "You do not have permission to edit LLM settings.") - -# payload.llm_optin_date = timezone.now() -# if payload.llm_optin_approved_by is None or payload.llm_optin_date is None: -# raise HttpError(400, "If LLM opt-in is true, both llm_optin_approved_by and llm_optin_date must be provided.") - -# # Update fields only if they are present in the payload - -# org_preferences.save() - -# return {"success": True, "res": 1} - - @orgpreference_router.put("/llm_approval", auth=auth.CustomAuthMiddleware()) @has_permission(["can_edit_llm_settings"]) def update_org_preferences(request, payload: UpdateLLMOptinSchema): """Updates llm preferences for the logged-in user's organization""" - # Get the organization associated with the current user orguser: OrgUser = request.orguser org = orguser.org - try: - # Fetch organization preferences for the user's organization - org_preferences = OrgPreferences.objects.get(org=org) - except OrgPreferences.DoesNotExist: - raise HttpError(404, "Preferences for this organization not found") + org_preferences = OrgPreferences.objects.filter(org=org).first() + if org_preferences is None: + raise HttpError(400, "Preferences for this organization not found") - # Determine fields based on the llm_optin value if payload.llm_optin is True: llm_optin_approved_by = request.orguser llm_optin_date = timezone.now() else: - # Set them to None for False llm_optin_approved_by = None llm_optin_date = None - # Update the OrgPreferences instance org_preferences.llm_optin = payload.llm_optin org_preferences.llm_optin_approved_by = llm_optin_approved_by org_preferences.llm_optin_date = llm_optin_date @@ -126,10 +93,9 @@ def update_discord_notifications(request, payload: UpdateDiscordNotificationsSch orguser: OrgUser = request.orguser org = orguser.org - try: - org_preferences = OrgPreferences.objects.get(org=org) - except OrgPreferences.DoesNotExist: - raise HttpError(404, "Preferences for this organization not found") + org_preferences = OrgPreferences.objects.filter(org=org).first() + if org_preferences is None: + raise HttpError(400, "Preferences for this organization not found") if payload.enable_discord_notifications: if not org_preferences.discord_webhook and not payload.discord_webhook: @@ -154,32 +120,26 @@ def update_discord_notifications(request, payload: UpdateDiscordNotificationsSch def get_org_preferences(request): """Gets preferences for an organization based on the logged-in user's organization""" orguser: OrgUser = request.orguser - org = orguser.org # Get the organization associated with the current user + org = orguser.org try: - # Fetch organization preferences for the user's organization - org_preferences = OrgPreferences.objects.select_related("org", "llm_optin_approved_by").get( + org_preferences = OrgPreferences.objects.select_related("llm_optin_approved_by").get( org=org ) except OrgPreferences.DoesNotExist: raise HttpError(404, "Organization preferences not found") - # Prepare data for the user who approved LLM opt-in, if available approved_by_data = None if org_preferences.llm_optin_approved_by and org_preferences.llm_optin_approved_by.user: user = org_preferences.llm_optin_approved_by.user approved_by_data = {"user_id": user.pk, "email": user.email} - # Structure the preferences response preferences = { "org": { - "id": org_preferences.org.id, - "name": org_preferences.org.name, - "slug": org_preferences.org.slug, - "type": org_preferences.org.type, + "name": org.name, + "slug": org.slug, + "type": org.type, }, - "trial_start_date": org_preferences.trial_start_date, - "trial_end_date": org_preferences.trial_end_date, "llm_optin": org_preferences.llm_optin, "llm_optin_approved_by": approved_by_data, "llm_optin_date": org_preferences.llm_optin_date, @@ -196,20 +156,15 @@ def get_org_superset_details(request): orguser: OrgUser = request.orguser org = orguser.org - try: - org_superset = OrgSupersets.objects.select_related( - "org", - ).get(org=org) - - except OrgSupersets.DoesNotExist: - raise HttpError(404, "Organizations superset details not found") + org_superset = OrgSupersets.objects.filter(org=org).first() + if org_superset is None: + raise HttpError(400, "Organizations superset details not found") response_data = { "org": { - "id": org_superset.org.id, - "name": org_superset.org.name, - "slug": org_superset.org.slug, - "type": org_superset.org.type, + "name": org.name, + "slug": org.slug, + "type": org.type, }, "superset_version": org_superset.superset_version, } @@ -241,17 +196,16 @@ def create_org_superset_details(request, payload: CreateOrgSupersetDetailsSchema def get_tools_versions(request): orguser: OrgUser = request.orguser org = orguser.org - - try: - # Validate org_id - org_superset = OrgSupersets.objects.get(org=org) - except OrgSupersets.DoesNotExist: - return {"success": False, "error": "Organization not found"} + org_superset = OrgSupersets.objects.filter(org=org).first() + if org_superset is None: + raise HttpError(400, "Supserset not found.") versions = [] # Airbyte Version try: + a = airbyte_service.get_airbyte_version() + print(a, "a") airbyte_url = AIRBYTE_URL_TO_GET_VERSION airbyte_response = requests.get(airbyte_url) if airbyte_response.status_code == 200: @@ -264,6 +218,8 @@ def get_tools_versions(request): # Prefect Version try: + b = prefect_service.get_prefect_version() + prefect_url = PREFECT_URL_TO_GET_VERSION prefect_response = requests.get(prefect_url) if prefect_response.status_code == 200: @@ -304,3 +260,30 @@ def get_tools_versions(request): versions.append({"Superset": {"version": org_superset.superset_version}}) return {"success": True, "res": versions} + + +@orgpreference_router.get("/org-plan", auth=auth.CustomAuthMiddleware()) +def get_org_plans(request): + """Gets preferences for an organization based on the logged-in user's organization""" + orguser: OrgUser = request.orguser + org = orguser.org + + org_plan = OrgPlans.objects.filter(org=org).first() + if org_plan is None: + raise HttpError(400, "Org's Plan not found") + + response_data = { + "org": { + "name": org.name, + "slug": org.slug, + "type": org.type, + }, + "base_plan": org_plan.base_plan, + "superset_included": org_plan.superset_included, + "subscription_duration": org_plan.subscription_duration, + "features": org_plan.features, + "start_date": org_plan.start_date, + "end_date": org_plan.end_date, + "can_upgrade_plan": org_plan.can_upgrade_plan, + } + return {"success": True, "res": response_data} diff --git a/ddpui/api/user_org_api.py b/ddpui/api/user_org_api.py index 88dba50c..0fefd499 100644 --- a/ddpui/api/user_org_api.py +++ b/ddpui/api/user_org_api.py @@ -30,6 +30,7 @@ OrgUserUpdateNewRole, OrgUserUpdatev1, ResetPasswordSchema, + ChangePasswordSchema, UserAttributes, VerifyEmailSchema, ) @@ -38,6 +39,7 @@ from ddpui.utils.deleteorg import delete_warehouse_v1 from ddpui.models.org import OrgWarehouse, Org, OrgType from ddpui.ddpairbyte import airbytehelpers +from ddpui.models.org_preferences import OrgPreferences user_org_router = Router() load_dotenv() @@ -55,7 +57,6 @@ def get_current_user_v2(request, org_slug: str = None): orguser: OrgUser = request.orguser user: User = request.orguser.user org: Org = orguser.org - # warehouse warehouse = OrgWarehouse.objects.filter(org=org).first() curr_orgusers = OrgUser.objects.filter(user=user) @@ -63,6 +64,10 @@ def get_current_user_v2(request, org_slug: str = None): if org_slug: curr_orgusers = curr_orgusers.filter(org__slug=org_slug) + org_preferences = OrgPreferences.objects.filter(org=org).first() + if org_preferences is None: + raise HttpError(400, "Organizations preferences not found") + res = [] for curr_orguser in curr_orgusers.prefetch_related( Prefetch( @@ -85,6 +90,11 @@ def get_current_user_v2(request, org_slug: str = None): ): if curr_orguser.org.orgtncs.exists(): curr_orguser.org.tnc_accepted = curr_orguser.org.orgtncs.exists() + + derived_llm_optin = curr_orguser.llm_optin + if not org_preferences.llm_optin: + derived_llm_optin = False + res.append( OrgUserResponse( email=user.email, @@ -99,7 +109,7 @@ def get_current_user_v2(request, org_slug: str = None): for rolep in curr_orguser.new_role.rolepermissions.all() ], is_demo=(curr_orguser.org.type == OrgType.DEMO if curr_orguser.org else False), - llm_optin=curr_orguser.llm_optin, + llm_optin=derived_llm_optin, ) ) @@ -400,6 +410,18 @@ def post_reset_password(request, payload: ResetPasswordSchema): # pylint: disab return {"success": 1} +@user_org_router.post( + "/users/change_password/", auth=auth.CustomAuthMiddleware() +) # from the settings panel +def change_password(request, payload: ChangePasswordSchema): # pylint: disable=unused-argument + """step 2 of the forgot-password flow""" + orguser = request.orguser + _, error = orguserfunctions.change_password(payload, orguser) + if error: + raise HttpError(400, error) + return {"success": 1} + + @user_org_router.get("/users/verify_email/resend", auth=auth.CustomAuthMiddleware()) @has_permission(["can_resend_email_verification"]) def get_verify_email_resend(request): # pylint: disable=unused-argument diff --git a/ddpui/core/orguserfunctions.py b/ddpui/core/orguserfunctions.py index 05e69a70..75324a2a 100644 --- a/ddpui/core/orguserfunctions.py +++ b/ddpui/core/orguserfunctions.py @@ -27,6 +27,7 @@ OrgUserUpdate, OrgUserUpdatev1, ResetPasswordSchema, + ChangePasswordSchema, UserAttributes, VerifyEmailSchema, ) @@ -593,6 +594,18 @@ def confirm_reset_password(payload: ResetPasswordSchema): return None, None +def change_password(payload: ChangePasswordSchema, orguser): + """If password and confirm password are same reset the password""" + + if payload.password != payload.confirmPassword: + return None, "Password and confirm password must be same" + + orguser.user.set_password(payload.password.get_secret_value()) + orguser.user.save() + + return None, None + + def resend_verification_email(orguser: OrgUser, email: str): """send a verification email to the user""" redis = RedisClient.get_instance() diff --git a/ddpui/models/org.py b/ddpui/models/org.py index 8878f4bb..1d49c041 100644 --- a/ddpui/models/org.py +++ b/ddpui/models/org.py @@ -7,8 +7,8 @@ class OrgType(str, Enum): """an enum representing the type of organization""" - DEMO = "demo" - POC = "poc" + SUBSCRIPTION = "subscription" + TRIAL = "trial" CLIENT = "client" @classmethod diff --git a/ddpui/models/org_plans.py b/ddpui/models/org_plans.py new file mode 100644 index 00000000..f1fe9460 --- /dev/null +++ b/ddpui/models/org_plans.py @@ -0,0 +1,25 @@ +from django.db import models +from ddpui.models.org import Org +from ddpui.models.org_user import OrgUser +from django.utils import timezone + + +class OrgPlans(models.Model): + """Model to store org preferences for settings panel""" + + org = models.OneToOneField(Org, on_delete=models.CASCADE, related_name="org_plans") + base_plan = models.CharField( + null=True, max_length=255, default=None + ) # plan DALGO or FREE TRAIL + superset_included = models.BooleanField(null=True, default=False) + subscription_duration = models.CharField( + null=True, max_length=255, default=None + ) # monthly, quatery. + features = models.JSONField( + null=True, default=None + ) # put the features as json and map in the frontend. + start_date = models.DateTimeField(null=True, blank=True, default=None) + end_date = models.DateTimeField(null=True, blank=True, default=None) + can_upgrade_plan = models.BooleanField(default=False) + created_at = models.DateTimeField(default=timezone.now) + updated_at = models.DateTimeField(default=timezone.now) diff --git a/ddpui/models/org_preferences.py b/ddpui/models/org_preferences.py index 879f2fb1..fb4ef5dc 100644 --- a/ddpui/models/org_preferences.py +++ b/ddpui/models/org_preferences.py @@ -1,30 +1,13 @@ from django.db import models -from enum import Enum from ddpui.models.org import Org from ddpui.models.org_user import OrgUser from django.utils import timezone -from django.core.exceptions import ValidationError - - -class OrgType(str, Enum): - """an enum representing the type of organization""" - - DEMO = "demo" - POC = "poc" - CLIENT = "client" - - @classmethod - def choices(cls): - """django model definition needs an iterable for `choices`""" - return [(key.value, key.name) for key in cls] class OrgPreferences(models.Model): """Model to store org preferences for settings panel""" org = models.OneToOneField(Org, on_delete=models.CASCADE, related_name="preferences") - trial_start_date = models.DateTimeField(null=True, blank=True, default=None) - trial_end_date = models.DateTimeField(null=True, blank=True, default=None) llm_optin = models.BooleanField(default=False) llm_optin_approved_by = models.ForeignKey( OrgUser, on_delete=models.CASCADE, related_name="approvedby", null=True, blank=True diff --git a/ddpui/models/org_user.py b/ddpui/models/org_user.py index e3cd2bf9..8583ee1a 100644 --- a/ddpui/models/org_user.py +++ b/ddpui/models/org_user.py @@ -188,6 +188,13 @@ class ResetPasswordSchema(Schema): password: SecretStr +class ChangePasswordSchema(Schema): + """Reset password from settings pannel""" + + password: SecretStr + confirmPassword: SecretStr + + class VerifyEmailSchema(Schema): """the payload for the verify-email workflow""" From 63cc307d160ed7386f174c88a77fe8c724614260 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Thu, 14 Nov 2024 19:11:43 +0530 Subject: [PATCH 06/57] removed airbyte and prefect service --- ddpui/api/org_preferences_api.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 80e76108..55fb4ba6 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -204,8 +204,6 @@ def get_tools_versions(request): # Airbyte Version try: - a = airbyte_service.get_airbyte_version() - print(a, "a") airbyte_url = AIRBYTE_URL_TO_GET_VERSION airbyte_response = requests.get(airbyte_url) if airbyte_response.status_code == 200: @@ -218,8 +216,6 @@ def get_tools_versions(request): # Prefect Version try: - b = prefect_service.get_prefect_version() - prefect_url = PREFECT_URL_TO_GET_VERSION prefect_response = requests.get(prefect_url) if prefect_response.status_code == 200: From 5f139786ead9c1355e1edbbc70b7b68ce9a39373 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 22:51:29 +0530 Subject: [PATCH 07/57] use env vars for dbt path, prefect url --- ddpui/api/org_preferences_api.py | 44 +++++++++++++++++--------------- ddpui/utils/constants.py | 5 ---- 2 files changed, 24 insertions(+), 25 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 55fb4ba6..ba081ed5 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -1,3 +1,4 @@ +import os from ninja import Router from ninja.errors import HttpError from ddpui import auth @@ -12,12 +13,6 @@ UpdateDiscordNotificationsSchema, CreateOrgSupersetDetailsSchema, ) -from ddpui.utils.constants import ( - AIRBYTE_URL_TO_GET_VERSION, - PREFECT_URL_TO_GET_VERSION, - DBT_VERSION_COMMAND, - ELEMENTARY_VERSION_COMMAND, -) from ddpui.auth import has_permission from ddpui.models.org import Org from ddpui.models.org_user import OrgUser @@ -194,6 +189,7 @@ def create_org_superset_details(request, payload: CreateOrgSupersetDetailsSchema @orgpreference_router.get("/toolinfo", auth=auth.CustomAuthMiddleware()) def get_tools_versions(request): + """get versions of the tools used in the system""" orguser: OrgUser = request.orguser org = orguser.org org_superset = OrgSupersets.objects.filter(org=org).first() @@ -203,21 +199,28 @@ def get_tools_versions(request): versions = [] # Airbyte Version - try: - airbyte_url = AIRBYTE_URL_TO_GET_VERSION - airbyte_response = requests.get(airbyte_url) - if airbyte_response.status_code == 200: - airbyte_data = airbyte_response.json() - versions.append({"Airbyte": {"version": airbyte_data.get("version")}}) - else: - versions.append({"Airbyte": {"version": "Not available"}}) - except Exception: - versions.append({"Airbyte": {"version": "Not available"}}) + versions.append({"Airbyte": {"version": "0.58"}}) + # this api will eventually work at some Airbyte version... but not at 0.58 + # try: + # abhost = os.getenv("AIRBYTE_SERVER_HOST") + # abport = os.getenv("AIRBYTE_SERVER_PORT") + # abver = os.getenv("AIRBYTE_SERVER_APIVER") + # airbyte_url = f"http://{abhost}:{abport}/api/{abver}/instance_configuration" + # airbyte_response = requests.get(airbyte_url, timeout=5) + # if airbyte_response.status_code == 200: + # airbyte_data = airbyte_response.json() + # versions.append({"Airbyte": {"version": airbyte_data.get("version")}}) + # else: + # versions.append({"Airbyte": {"version": "Not available"}}) + # except Exception: + # versions.append({"Airbyte": {"version": "Not available"}}) # Prefect Version try: - prefect_url = PREFECT_URL_TO_GET_VERSION - prefect_response = requests.get(prefect_url) + prefect_host = os.getenv("PREFECT_SERVER_HOST") + prefect_port = os.getenv("PREFECT_SERVER_PORT") + prefect_url = f"http://{prefect_host}:{prefect_port}/api/admin/version" + prefect_response = requests.get(prefect_url, timeout=5) if prefect_response.status_code == 200: version = prefect_response.text.strip().strip('"') versions.append({"Prefect": {"version": version}}) @@ -226,9 +229,10 @@ def get_tools_versions(request): except Exception: versions.append({"Prefect": {"version": "Not available"}}) + dbt_venv = os.getenv("DBT_VENV") # dbt Version try: - dbt_version_command = DBT_VERSION_COMMAND + dbt_version_command = [os.path.join(dbt_venv, "venv", "bin", "dbt"), "--version"] dbt_output = subprocess.check_output(dbt_version_command, text=True) for line in dbt_output.splitlines(): if "installed:" in line: @@ -241,7 +245,7 @@ def get_tools_versions(request): # Elementary Version try: - elementary_version_command = ELEMENTARY_VERSION_COMMAND + elementary_version_command = [os.path.join(dbt_venv, "venv", "bin", "edr"), "--version"] elementary_output = subprocess.check_output(elementary_version_command, text=True) for line in elementary_output.splitlines(): if "Elementary version" in line: diff --git a/ddpui/utils/constants.py b/ddpui/utils/constants.py index 3798f988..f72ee0eb 100644 --- a/ddpui/utils/constants.py +++ b/ddpui/utils/constants.py @@ -51,8 +51,3 @@ # LLM data analysis ; row limit fetched and sent to llm service LIMIT_ROWS_TO_SEND_TO_LLM = 500 - -AIRBYTE_URL_TO_GET_VERSION = "http://localhost:8000/api/v1/instance_configuration" -PREFECT_URL_TO_GET_VERSION = "http://localhost:4200/api/admin/version" -DBT_VERSION_COMMAND = ["/home/ddp/dbt/venv/bin/dbt", "--version"] -ELEMENTARY_VERSION_COMMAND = ["/home/ddp/dbt/venv/bin/edr", "--version"] From df220fac2ece5f909792aa0811a70cb4cbb34035 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 23:04:13 +0530 Subject: [PATCH 08/57] version getters moved to service files --- ddpui/api/org_preferences_api.py | 87 ++++++----------------------- ddpui/ddpairbyte/airbyte_service.py | 19 +++++++ ddpui/ddpdbt/dbt_service.py | 28 ++++++++++ ddpui/ddpprefect/prefect_service.py | 17 ++++++ 4 files changed, 82 insertions(+), 69 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index ba081ed5..f569d7ae 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -1,28 +1,25 @@ import os +import subprocess from ninja import Router from ninja.errors import HttpError +from django.utils import timezone from ddpui import auth from ddpui.models.org_preferences import OrgPreferences from ddpui.models.org_supersets import OrgSupersets from ddpui.models.org_plans import OrgPlans -from django.utils import timezone from ddpui.schemas.org_preferences_schema import ( CreateOrgPreferencesSchema, - UpdateOrgPreferencesSchema, UpdateLLMOptinSchema, UpdateDiscordNotificationsSchema, CreateOrgSupersetDetailsSchema, ) from ddpui.auth import has_permission -from ddpui.models.org import Org from ddpui.models.org_user import OrgUser +from ddpui.ddpdbt import dbt_service from ddpui.ddpairbyte import airbyte_service from ddpui.ddpprefect import ( prefect_service, ) -from ddpui.core.orgdbt_manager import DbtProjectManager -import requests -import subprocess orgpreference_router = Router() @@ -192,72 +189,24 @@ def get_tools_versions(request): """get versions of the tools used in the system""" orguser: OrgUser = request.orguser org = orguser.org - org_superset = OrgSupersets.objects.filter(org=org).first() - if org_superset is None: - raise HttpError(400, "Supserset not found.") versions = [] - # Airbyte Version - versions.append({"Airbyte": {"version": "0.58"}}) - # this api will eventually work at some Airbyte version... but not at 0.58 - # try: - # abhost = os.getenv("AIRBYTE_SERVER_HOST") - # abport = os.getenv("AIRBYTE_SERVER_PORT") - # abver = os.getenv("AIRBYTE_SERVER_APIVER") - # airbyte_url = f"http://{abhost}:{abport}/api/{abver}/instance_configuration" - # airbyte_response = requests.get(airbyte_url, timeout=5) - # if airbyte_response.status_code == 200: - # airbyte_data = airbyte_response.json() - # versions.append({"Airbyte": {"version": airbyte_data.get("version")}}) - # else: - # versions.append({"Airbyte": {"version": "Not available"}}) - # except Exception: - # versions.append({"Airbyte": {"version": "Not available"}}) - - # Prefect Version - try: - prefect_host = os.getenv("PREFECT_SERVER_HOST") - prefect_port = os.getenv("PREFECT_SERVER_PORT") - prefect_url = f"http://{prefect_host}:{prefect_port}/api/admin/version" - prefect_response = requests.get(prefect_url, timeout=5) - if prefect_response.status_code == 200: - version = prefect_response.text.strip().strip('"') - versions.append({"Prefect": {"version": version}}) - else: - versions.append({"Prefect": {"version": "Not available"}}) - except Exception: - versions.append({"Prefect": {"version": "Not available"}}) - - dbt_venv = os.getenv("DBT_VENV") - # dbt Version - try: - dbt_version_command = [os.path.join(dbt_venv, "venv", "bin", "dbt"), "--version"] - dbt_output = subprocess.check_output(dbt_version_command, text=True) - for line in dbt_output.splitlines(): - if "installed:" in line: - versions.append({"DBT": {"version": line.split(":")[1].strip()}}) - break - else: - versions.append({"DBT": {"version": "Not available"}}) - except Exception: - versions.append({"DBT": {"version": "Not available"}}) - - # Elementary Version - try: - elementary_version_command = [os.path.join(dbt_venv, "venv", "bin", "edr"), "--version"] - elementary_output = subprocess.check_output(elementary_version_command, text=True) - for line in elementary_output.splitlines(): - if "Elementary version" in line: - versions.append({"Elementary": {"version": line.split()[-1].strip()}}) - break - else: - versions.append({"Elementary": {"version": "Not available"}}) - except Exception: - versions.append({"Elementary": {"version": "Not available"}}) - - # Superset Version - versions.append({"Superset": {"version": org_superset.superset_version}}) + airbyte_version = airbyte_service.get_airbyte_version() + versions.append({"Airbyte": {"version": airbyte_version}}) + + prefect_version = prefect_service.get_prefect_server_version() + versions.append({"Prefect": {"version": prefect_version}}) + + dbt_version = dbt_service.get_dbt_version() + versions.append({"DBT": {"version": dbt_version}}) + + edr_version = dbt_service.get_edr_version() + versions.append({"Elementary": {"version": edr_version}}) + + org_superset = OrgSupersets.objects.filter(org=org).first() + superset_version = org_superset.superset_version if org_superset else "Not available" + versions.append({"Superset": {"version": superset_version}}) return {"success": True, "res": versions} diff --git a/ddpui/ddpairbyte/airbyte_service.py b/ddpui/ddpairbyte/airbyte_service.py index 069c32ae..6b479f79 100644 --- a/ddpui/ddpairbyte/airbyte_service.py +++ b/ddpui/ddpairbyte/airbyte_service.py @@ -971,3 +971,22 @@ def update_schema_change( raise HttpError(500, "failed to trigger Prefect flow run") from error return res + + +def get_airbyte_version(): + """Get the version of the Airbyte server""" + return "0.58" + # the following will eventually work at some Airbyte version... but not at 0.58 + # try: + # abhost = os.getenv("AIRBYTE_SERVER_HOST") + # abport = os.getenv("AIRBYTE_SERVER_PORT") + # abver = os.getenv("AIRBYTE_SERVER_APIVER") + # airbyte_url = f"http://{abhost}:{abport}/api/{abver}/instance_configuration" + # airbyte_response = requests.get(airbyte_url, timeout=5) + # if airbyte_response.status_code == 200: + # airbyte_data = airbyte_response.json() + # return airbyte_data.get("version") + # else: + # return "Not available" + # except Exception: + # return "Not available" diff --git a/ddpui/ddpdbt/dbt_service.py b/ddpui/ddpdbt/dbt_service.py index a0ba332e..d5ee9c1f 100644 --- a/ddpui/ddpdbt/dbt_service.py +++ b/ddpui/ddpdbt/dbt_service.py @@ -331,3 +331,31 @@ def refresh_elementary_report_via_prefect(orguser: OrgUser) -> dict: raise HttpError(400, "failed to start a run") from error return res + + +def get_dbt_version(): + """get dbt version""" + dbt_venv = os.getenv("DBT_VENV") + try: + dbt_version_command = [os.path.join(dbt_venv, "venv", "bin", "dbt"), "--version"] + dbt_output = subprocess.check_output(dbt_version_command, text=True) + for line in dbt_output.splitlines(): + if "installed:" in line: + return line.split(":")[1].strip() + return "Not available" + except Exception: + return "Not available" + + +def get_edr_version(): + """get elementary report version""" + dbt_venv = os.getenv("DBT_VENV") + try: + elementary_version_command = [os.path.join(dbt_venv, "venv", "bin", "edr"), "--version"] + elementary_output = subprocess.check_output(elementary_version_command, text=True) + for line in elementary_output.splitlines(): + if "Elementary version" in line: + return line.split()[-1].strip() + return "Not available" + except Exception: + return "Not available" diff --git a/ddpui/ddpprefect/prefect_service.py b/ddpui/ddpprefect/prefect_service.py index c514afa2..2304b426 100644 --- a/ddpui/ddpprefect/prefect_service.py +++ b/ddpui/ddpprefect/prefect_service.py @@ -158,6 +158,23 @@ def prefect_delete(endpoint: str, **kwargs) -> None: raise HttpError(res.status_code, res.text) from error +# ================================================================================================ +def get_prefect_server_version(): + """fetches the prefect server version""" + try: + prefect_host = os.getenv("PREFECT_SERVER_HOST") + prefect_port = os.getenv("PREFECT_SERVER_PORT") + prefect_url = f"http://{prefect_host}:{prefect_port}/api/admin/version" + prefect_response = requests.get(prefect_url, timeout=5) + if prefect_response.status_code == 200: + version = prefect_response.text.strip().strip('"') + return version + else: + return "Not available" + except Exception: + return "Not available" + + # ================================================================================================ def get_airbyte_server_block_id(blockname) -> str | None: """get the block_id for the server block having this name""" From dc5b0d56fcbb9df65732a9d8021fa1534d15e46c Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 23:24:50 +0530 Subject: [PATCH 09/57] to_json() --- ddpui/models/org_plans.py | 17 +++++++++++++++++ ddpui/models/org_preferences.py | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/ddpui/models/org_plans.py b/ddpui/models/org_plans.py index f1fe9460..59a0d89c 100644 --- a/ddpui/models/org_plans.py +++ b/ddpui/models/org_plans.py @@ -23,3 +23,20 @@ class OrgPlans(models.Model): can_upgrade_plan = models.BooleanField(default=False) created_at = models.DateTimeField(default=timezone.now) updated_at = models.DateTimeField(default=timezone.now) + + def to_json(self) -> dict: + """Return a dict representation of the model""" + return { + "org": { + "name": self.org.name, + "slug": self.org.slug, + "type": self.org.type, + }, + "base_plan": self.base_plan, + "superset_included": self.superset_included, + "subscription_duration": self.subscription_duration, + "features": self.features, + "start_date": self.start_date, + "end_date": self.end_date, + "can_upgrade_plan": self.can_upgrade_plan, + } diff --git a/ddpui/models/org_preferences.py b/ddpui/models/org_preferences.py index fb4ef5dc..45f51940 100644 --- a/ddpui/models/org_preferences.py +++ b/ddpui/models/org_preferences.py @@ -17,3 +17,20 @@ class OrgPreferences(models.Model): discord_webhook = models.URLField(blank=True, null=True) created_at = models.DateTimeField(default=timezone.now) updated_at = models.DateTimeField(default=timezone.now) + + def to_json(self) -> dict: + """Return a dict representation of the model""" + return { + "org": { + "name": self.org.name, + "slug": self.org.slug, + "type": self.org.type, + }, + "llm_optin": str(self.llm_optin), + "llm_optin_approved_by": ( + self.llm_optin_approved_by.user.email if self.llm_optin_approved_by else None + ), + "llm_optin_date": self.llm_optin_date.isoformat() if self.llm_optin_date else None, + "enable_discord_notifications": str(self.enable_discord_notifications), + "discord_webhook": self.discord_webhook, + } From 48762fba67c9e1c03377019e59e560d1c90bee41 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 23:26:01 +0530 Subject: [PATCH 10/57] ensure OrgPreferences exists instead of raising exceptions to_json refactoring removed superset-version-specific endpoints --- ddpui/api/org_preferences_api.py | 136 +++++-------------------------- 1 file changed, 22 insertions(+), 114 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index f569d7ae..5cbc3673 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -1,5 +1,3 @@ -import os -import subprocess from ninja import Router from ninja.errors import HttpError from django.utils import timezone @@ -26,10 +24,10 @@ @orgpreference_router.post("/", auth=auth.CustomAuthMiddleware()) def create_org_preferences(request, payload: CreateOrgPreferencesSchema): + """Creates preferences for an organization""" orguser: OrgUser = request.orguser org = orguser.org payload.org = org - """Creates preferences for an organization""" if OrgPreferences.objects.filter(org=org).exists(): raise HttpError(400, "Organization preferences already exist") @@ -39,15 +37,7 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): org=org, **payload_data # Use the rest of the payload ) - preferences = { - "llm_optin": org_preferences.llm_optin, - "llm_optin_approved_by": org_preferences.llm_optin_approved_by, - "llm_optin_date": org_preferences.llm_optin_date, - "enable_discord_notifications": org_preferences.enable_discord_notifications, - "discord_webhook": org_preferences.discord_webhook, - } - - return {"success": True, "res": preferences} + return {"success": True, "res": org_preferences.to_json()} @orgpreference_router.put("/llm_approval", auth=auth.CustomAuthMiddleware()) @@ -60,21 +50,20 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): org_preferences = OrgPreferences.objects.filter(org=org).first() if org_preferences is None: - raise HttpError(400, "Preferences for this organization not found") + org_preferences = OrgPreferences.objects.create(org=org) if payload.llm_optin is True: - llm_optin_approved_by = request.orguser - llm_optin_date = timezone.now() + org_preferences.llm_optin = True + org_preferences.llm_optin_approved_by = orguser + org_preferences.llm_optin_date = timezone.now() else: - llm_optin_approved_by = None - llm_optin_date = None + org_preferences.llm_optin = False + org_preferences.llm_optin_approved_by = None + org_preferences.llm_optin_date = None - org_preferences.llm_optin = payload.llm_optin - org_preferences.llm_optin_approved_by = llm_optin_approved_by - org_preferences.llm_optin_date = llm_optin_date org_preferences.save() - return {"success": True, "res": 1} + return {"success": True, "res": org_preferences.to_json()} @orgpreference_router.put("/enable-discord-notifications", auth=auth.CustomAuthMiddleware()) @@ -87,25 +76,21 @@ def update_discord_notifications(request, payload: UpdateDiscordNotificationsSch org_preferences = OrgPreferences.objects.filter(org=org).first() if org_preferences is None: - raise HttpError(400, "Preferences for this organization not found") + org_preferences = OrgPreferences.objects.create(org=org) if payload.enable_discord_notifications: if not org_preferences.discord_webhook and not payload.discord_webhook: raise HttpError(400, "Discord webhook is required to enable notifications.") - discord_webhook = payload.discord_webhook or org_preferences.discord_webhook + if payload.discord_webhook: + org_preferences.discord_webhook = payload.discord_webhook + org_preferences.enable_discord_notifications = True else: - discord_webhook = None + org_preferences.discord_webhook = None + org_preferences.enable_discord_notifications = False - org_preferences.enable_discord_notifications = payload.enable_discord_notifications - org_preferences.discord_webhook = discord_webhook org_preferences.save() - response_data = { - "enable_discord_notifications": org_preferences.enable_discord_notifications, - "discord_webhook": org_preferences.discord_webhook, - } - - return {"success": True, "res": response_data} + return {"success": True, "res": org_preferences.to_json()} @orgpreference_router.get("/", auth=auth.CustomAuthMiddleware()) @@ -114,74 +99,11 @@ def get_org_preferences(request): orguser: OrgUser = request.orguser org = orguser.org - try: - org_preferences = OrgPreferences.objects.select_related("llm_optin_approved_by").get( - org=org - ) - except OrgPreferences.DoesNotExist: - raise HttpError(404, "Organization preferences not found") - - approved_by_data = None - if org_preferences.llm_optin_approved_by and org_preferences.llm_optin_approved_by.user: - user = org_preferences.llm_optin_approved_by.user - approved_by_data = {"user_id": user.pk, "email": user.email} - - preferences = { - "org": { - "name": org.name, - "slug": org.slug, - "type": org.type, - }, - "llm_optin": org_preferences.llm_optin, - "llm_optin_approved_by": approved_by_data, - "llm_optin_date": org_preferences.llm_optin_date, - "enable_discord_notifications": org_preferences.enable_discord_notifications, - "discord_webhook": org_preferences.discord_webhook, - } - - return {"success": True, "res": preferences} - - -@orgpreference_router.get("/org-superset", auth=auth.CustomAuthMiddleware()) -def get_org_superset_details(request): - """Gets preferences for an organization by org_id""" - orguser: OrgUser = request.orguser - org = orguser.org - - org_superset = OrgSupersets.objects.filter(org=org).first() - if org_superset is None: - raise HttpError(400, "Organizations superset details not found") - - response_data = { - "org": { - "name": org.name, - "slug": org.slug, - "type": org.type, - }, - "superset_version": org_superset.superset_version, - } - return {"success": True, "res": response_data} - - -@orgpreference_router.post("/org-superset", auth=auth.CustomAuthMiddleware()) -def create_org_superset_details(request, payload: CreateOrgSupersetDetailsSchema): - """Creates or updates an entry in db with org superset details.""" - orguser: OrgUser = request.orguser - org = orguser.org - - org_superset, created = OrgSupersets.objects.update_or_create(org=org, defaults=payload.dict()) - - response_data = { - "org": { - "name": org.name, - "slug": org.slug, - "type": org.type, - }, - "superset_version": org_superset.superset_version, - "created": created, - } + org_preferences = OrgPreferences.objects.filter(org=org).first() + if org_preferences is None: + org_preferences = OrgPreferences.objects.create(org=org) - return {"success": True, "res": response_data} + return {"success": True, "res": org_preferences.to_json()} @orgpreference_router.get("/toolinfo", auth=auth.CustomAuthMiddleware()) @@ -221,18 +143,4 @@ def get_org_plans(request): if org_plan is None: raise HttpError(400, "Org's Plan not found") - response_data = { - "org": { - "name": org.name, - "slug": org.slug, - "type": org.type, - }, - "base_plan": org_plan.base_plan, - "superset_included": org_plan.superset_included, - "subscription_duration": org_plan.subscription_duration, - "features": org_plan.features, - "start_date": org_plan.start_date, - "end_date": org_plan.end_date, - "can_upgrade_plan": org_plan.can_upgrade_plan, - } - return {"success": True, "res": response_data} + return {"success": True, "res": org_plan.to_json()} From b4d5bee7ece5d20f6f65c28f7bc43940fc8331a1 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 23:48:41 +0530 Subject: [PATCH 11/57] deprecated userpreferences-> {llm_optin, discord_webhook, enable_discord } since they are moving to orgpreferences --- ddpui/api/user_org_api.py | 8 ++----- ddpui/api/user_preferences_api.py | 37 ++++--------------------------- ddpui/core/orguserfunctions.py | 2 +- ddpui/models/userpreferences.py | 14 ++++++++---- 4 files changed, 17 insertions(+), 44 deletions(-) diff --git a/ddpui/api/user_org_api.py b/ddpui/api/user_org_api.py index 0fefd499..e74237e1 100644 --- a/ddpui/api/user_org_api.py +++ b/ddpui/api/user_org_api.py @@ -66,7 +66,7 @@ def get_current_user_v2(request, org_slug: str = None): org_preferences = OrgPreferences.objects.filter(org=org).first() if org_preferences is None: - raise HttpError(400, "Organizations preferences not found") + org_preferences = OrgPreferences.objects.create(org=org) res = [] for curr_orguser in curr_orgusers.prefetch_related( @@ -91,10 +91,6 @@ def get_current_user_v2(request, org_slug: str = None): if curr_orguser.org.orgtncs.exists(): curr_orguser.org.tnc_accepted = curr_orguser.org.orgtncs.exists() - derived_llm_optin = curr_orguser.llm_optin - if not org_preferences.llm_optin: - derived_llm_optin = False - res.append( OrgUserResponse( email=user.email, @@ -109,7 +105,7 @@ def get_current_user_v2(request, org_slug: str = None): for rolep in curr_orguser.new_role.rolepermissions.all() ], is_demo=(curr_orguser.org.type == OrgType.DEMO if curr_orguser.org else False), - llm_optin=derived_llm_optin, + llm_optin=org_preferences.llm_optin, ) ) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index f1796792..2a02a1e4 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -19,22 +19,12 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): if UserPreferences.objects.filter(orguser=orguser).exists(): raise HttpError(400, "Preferences already exist") - userpreferences = UserPreferences.objects.create( + user_preferences = UserPreferences.objects.create( orguser=orguser, - llm_optin=payload.llm_optin, - enable_discord_notifications=payload.enable_discord_notifications, - discord_webhook=payload.discord_webhook, enable_email_notifications=payload.enable_email_notifications, ) - preferences = { - "enable_email_notifications": userpreferences.enable_email_notifications, - "llm_optin": userpreferences.llm_optin, - "discord_webhook": userpreferences.discord_webhook, - "enable_discord_notifications": userpreferences.enable_discord_notifications, - } - - return {"success": True, "res": preferences} + return {"success": True, "res": user_preferences.to_json()} @userpreference_router.put("/", auth=auth.CustomAuthMiddleware()) @@ -44,24 +34,11 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) - if payload.enable_discord_notifications is not None: - user_preferences.enable_discord_notifications = payload.enable_discord_notifications - if payload.llm_optin is not None: - user_preferences.llm_optin = payload.llm_optin if payload.enable_email_notifications is not None: user_preferences.enable_email_notifications = payload.enable_email_notifications - if payload.discord_webhook is not None: - user_preferences.discord_webhook = payload.discord_webhook user_preferences.save() - preferences = { - "discord_webhook": user_preferences.discord_webhook, - "enable_discord_notifications": user_preferences.enable_discord_notifications, - "enable_email_notifications": user_preferences.enable_email_notifications, - "llm_optin": user_preferences.llm_optin, - } - - return {"success": True, "res": preferences} + return {"success": True, "res": user_preferences.to_json()} @userpreference_router.get("/", auth=auth.CustomAuthMiddleware()) @@ -71,10 +48,4 @@ def get_user_preferences(request): user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) - preferences = { - "discord_webhook": user_preferences.discord_webhook, - "enable_discord_notifications": user_preferences.enable_discord_notifications, - "enable_email_notifications": user_preferences.enable_email_notifications, - "llm_optin": user_preferences.llm_optin, - } - return {"success": True, "res": preferences} + return {"success": True, "res": user_preferences.to_json()} diff --git a/ddpui/core/orguserfunctions.py b/ddpui/core/orguserfunctions.py index 75324a2a..9d151643 100644 --- a/ddpui/core/orguserfunctions.py +++ b/ddpui/core/orguserfunctions.py @@ -594,7 +594,7 @@ def confirm_reset_password(payload: ResetPasswordSchema): return None, None -def change_password(payload: ChangePasswordSchema, orguser): +def change_password(payload: ChangePasswordSchema, orguser: OrgUser): """If password and confirm password are same reset the password""" if payload.password != payload.confirmPassword: diff --git a/ddpui/models/userpreferences.py b/ddpui/models/userpreferences.py index 162b07f9..89eb9f4a 100644 --- a/ddpui/models/userpreferences.py +++ b/ddpui/models/userpreferences.py @@ -1,15 +1,21 @@ from django.db import models -from ddpui.models.org_user import OrgUser from django.utils import timezone +from ddpui.models.org_user import OrgUser class UserPreferences(models.Model): """Model to store user preferences for notifications""" orguser = models.OneToOneField(OrgUser, on_delete=models.CASCADE, related_name="preferences") - enable_discord_notifications = models.BooleanField(default=False) - discord_webhook = models.URLField(blank=True, null=True) + enable_discord_notifications = models.BooleanField(default=False) # deprecated + discord_webhook = models.URLField(blank=True, null=True) # deprecated enable_email_notifications = models.BooleanField(default=False) - llm_optin = models.BooleanField(default=False) + llm_optin = models.BooleanField(default=False) # deprecated created_at = models.DateTimeField(default=timezone.now) updated_at = models.DateTimeField(default=timezone.now) + + def to_json(self) -> dict: + """Return a dict representation of the model""" + return { + "enable_email_notifications": self.enable_email_notifications, + } From 29234af09145c4645816046c292f8c142dc4b2fc Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 23:50:10 +0530 Subject: [PATCH 12/57] fix typo --- ddpui/models/org.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ddpui/models/org.py b/ddpui/models/org.py index 1d49c041..bb5ea174 100644 --- a/ddpui/models/org.py +++ b/ddpui/models/org.py @@ -9,7 +9,7 @@ class OrgType(str, Enum): SUBSCRIPTION = "subscription" TRIAL = "trial" - CLIENT = "client" + DEMO = "demo" @classmethod def choices(cls): From a40babcc281b67c15146130a7c9ec6c1e4426436 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Thu, 14 Nov 2024 23:51:25 +0530 Subject: [PATCH 13/57] fix --- ddpui/models/org.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ddpui/models/org.py b/ddpui/models/org.py index bb5ea174..458b7136 100644 --- a/ddpui/models/org.py +++ b/ddpui/models/org.py @@ -70,7 +70,7 @@ class Org(models.Model): ) viz_url = models.CharField(max_length=100, null=True) viz_login_type = models.CharField(choices=OrgVizLoginType.choices(), max_length=50, null=True) - type = models.CharField(choices=OrgType.choices(), max_length=50, default=OrgType.CLIENT) + type = models.CharField(choices=OrgType.choices(), max_length=50, default=OrgType.SUBSCRIPTION) ses_whitelisted_email = models.TextField(max_length=100, null=True) created_at = models.DateTimeField(auto_created=True, default=timezone.now) updated_at = models.DateTimeField(auto_now=True) From 82b0be2cb61f5d5bdfcb631670240284982b9ed1 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Fri, 15 Nov 2024 10:38:32 +0530 Subject: [PATCH 14/57] changes --- ddpui/api/user_org_api.py | 8 +++----- ddpui/models/org.py | 4 ++-- ddpui/models/org_user.py | 1 + 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/ddpui/api/user_org_api.py b/ddpui/api/user_org_api.py index 0fefd499..a05a7792 100644 --- a/ddpui/api/user_org_api.py +++ b/ddpui/api/user_org_api.py @@ -91,10 +91,7 @@ def get_current_user_v2(request, org_slug: str = None): if curr_orguser.org.orgtncs.exists(): curr_orguser.org.tnc_accepted = curr_orguser.org.orgtncs.exists() - derived_llm_optin = curr_orguser.llm_optin - if not org_preferences.llm_optin: - derived_llm_optin = False - + print("broooo") res.append( OrgUserResponse( email=user.email, @@ -109,7 +106,8 @@ def get_current_user_v2(request, org_slug: str = None): for rolep in curr_orguser.new_role.rolepermissions.all() ], is_demo=(curr_orguser.org.type == OrgType.DEMO if curr_orguser.org else False), - llm_optin=derived_llm_optin, + llm_optin=curr_orguser.llm_optin, + is_llm_active=org_preferences.llm_optin, ) ) diff --git a/ddpui/models/org.py b/ddpui/models/org.py index 1d49c041..458b7136 100644 --- a/ddpui/models/org.py +++ b/ddpui/models/org.py @@ -9,7 +9,7 @@ class OrgType(str, Enum): SUBSCRIPTION = "subscription" TRIAL = "trial" - CLIENT = "client" + DEMO = "demo" @classmethod def choices(cls): @@ -70,7 +70,7 @@ class Org(models.Model): ) viz_url = models.CharField(max_length=100, null=True) viz_login_type = models.CharField(choices=OrgVizLoginType.choices(), max_length=50, null=True) - type = models.CharField(choices=OrgType.choices(), max_length=50, default=OrgType.CLIENT) + type = models.CharField(choices=OrgType.choices(), max_length=50, default=OrgType.SUBSCRIPTION) ses_whitelisted_email = models.TextField(max_length=100, null=True) created_at = models.DateTimeField(auto_created=True, default=timezone.now) updated_at = models.DateTimeField(auto_now=True) diff --git a/ddpui/models/org_user.py b/ddpui/models/org_user.py index 8583ee1a..8f8bbe07 100644 --- a/ddpui/models/org_user.py +++ b/ddpui/models/org_user.py @@ -132,6 +132,7 @@ class OrgUserResponse(Schema): new_role_slug: str | None permissions: list[dict] llm_optin: bool = None + is_llm_active: bool = None class Invitation(models.Model): From 473fdad5863aaccf7800be726db4ad139058b99b Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Fri, 15 Nov 2024 12:02:46 +0530 Subject: [PATCH 15/57] fix the enums --- ddpui/migrations/0093_remove_org_is_demo_org_type.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ddpui/migrations/0093_remove_org_is_demo_org_type.py b/ddpui/migrations/0093_remove_org_is_demo_org_type.py index c37450b5..6b7d4afb 100644 --- a/ddpui/migrations/0093_remove_org_is_demo_org_type.py +++ b/ddpui/migrations/0093_remove_org_is_demo_org_type.py @@ -18,8 +18,8 @@ class Migration(migrations.Migration): model_name="org", name="type", field=models.CharField( - choices=[("demo", "DEMO"), ("poc", "POC"), ("client", "CLIENT")], - default=ddpui.models.org.OrgType["CLIENT"], + choices=[("demo", "DEMO"), ("trial", "TRIAL"), ("subscription", "SUBSCRIPTION")], + default=ddpui.models.org.OrgType["SUBSCRIPTION"], max_length=50, ), ), From 78c79733a8348d18964ba8fe4fe0a1ee5b075367 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Fri, 15 Nov 2024 12:10:05 +0530 Subject: [PATCH 16/57] fix tests --- ddpui/tests/api_tests/test_user_preferences_api.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/ddpui/tests/api_tests/test_user_preferences_api.py b/ddpui/tests/api_tests/test_user_preferences_api.py index 89f6b6cd..841fb6ee 100644 --- a/ddpui/tests/api_tests/test_user_preferences_api.py +++ b/ddpui/tests/api_tests/test_user_preferences_api.py @@ -155,8 +155,6 @@ def test_get_user_preferences_success(orguser, user_preferences): response = get_user_preferences(request) assert response["success"] is True assert response["res"] == { - "discord_webhook": user_preferences.discord_webhook, - "enable_email_notifications": user_preferences.enable_email_notifications, "enable_discord_notifications": user_preferences.enable_discord_notifications, } @@ -170,8 +168,6 @@ def test_get_user_preferences_success_if_not_exist(orguser): response = get_user_preferences(request) assert response["success"] is True assert response["res"] == { - "discord_webhook": None, "enable_email_notifications": False, - "enable_discord_notifications": False, } assert UserPreferences.objects.filter(orguser=orguser).exists() From 87b3e30fa2fdf68029145c25ac70575a509ef529 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Fri, 15 Nov 2024 15:38:11 +0530 Subject: [PATCH 17/57] airbyte and prefect function changed and formatted --- ddpui/api/org_preferences_api.py | 25 ++++---------------- ddpui/ddpairbyte/airbyte_service.py | 36 ++++++++++++++++++++++++----- ddpui/ddpprefect/prefect_service.py | 6 +++++ 3 files changed, 40 insertions(+), 27 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 55fb4ba6..042df352 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -202,29 +202,12 @@ def get_tools_versions(request): versions = [] - # Airbyte Version - try: - airbyte_url = AIRBYTE_URL_TO_GET_VERSION - airbyte_response = requests.get(airbyte_url) - if airbyte_response.status_code == 200: - airbyte_data = airbyte_response.json() - versions.append({"Airbyte": {"version": airbyte_data.get("version")}}) - else: - versions.append({"Airbyte": {"version": "Not available"}}) - except Exception: - versions.append({"Airbyte": {"version": "Not available"}}) + ver = airbyte_service.get_current_airbyte_version() + versions.append({"Airbyte": {"version": ver if ver else "Not available"}}) # Prefect Version - try: - prefect_url = PREFECT_URL_TO_GET_VERSION - prefect_response = requests.get(prefect_url) - if prefect_response.status_code == 200: - version = prefect_response.text.strip().strip('"') - versions.append({"Prefect": {"version": version}}) - else: - versions.append({"Prefect": {"version": "Not available"}}) - except Exception: - versions.append({"Prefect": {"version": "Not available"}}) + ver = prefect_service.get_prefect_version() + versions.append({"Prefect": {"version": ver if ver else "Not available"}}) # dbt Version try: diff --git a/ddpui/ddpairbyte/airbyte_service.py b/ddpui/ddpairbyte/airbyte_service.py index 069c32ae..bb7a2565 100644 --- a/ddpui/ddpairbyte/airbyte_service.py +++ b/ddpui/ddpairbyte/airbyte_service.py @@ -35,6 +35,10 @@ def abreq(endpoint, req=None, **kwargs): """Request to the airbyte server""" + method = kwargs.get("method", "POST") + if method not in ["GET", "POST"]: + raise HttpError(500, "method not supported") + request = thread.get_current_request() abhost = os.getenv("AIRBYTE_SERVER_HOST") @@ -68,12 +72,21 @@ def abreq(endpoint, req=None, **kwargs): logger.info("Making request to Airbyte server: %s", endpoint) try: - res = requests.post( - f"http://{abhost}:{abport}/api/{abver}/{endpoint}", - headers={"Authorization": f"Basic {token}"}, - json=req, - timeout=kwargs.get("timeout", 30), - ) + res = {} + if method == "POST": + res = requests.post( + f"http://{abhost}:{abport}/api/{abver}/{endpoint}", + headers={"Authorization": f"Basic {token}"}, + json=req, + timeout=kwargs.get("timeout", 30), + ) + elif method == "GET": + res = requests.get( + f"http://{abhost}:{abport}/api/{abver}/{endpoint}", + headers={"Authorization": f"Basic {token}"}, + json=req, + timeout=kwargs.get("timeout", 30), + ) except requests.exceptions.ConnectionError as conn_error: logger.exception(conn_error) raise HttpError(500, str(conn_error)) from conn_error @@ -971,3 +984,14 @@ def update_schema_change( raise HttpError(500, "failed to trigger Prefect flow run") from error return res + + +def get_current_airbyte_version(): + """Fetch airbyte version""" + + res = abreq("instance_configuration", method="GET") + print(res, "AIRBYTE RESPONSE") + if "version" not in res: + logger.error("No version found") + return None + return res["version"] diff --git a/ddpui/ddpprefect/prefect_service.py b/ddpui/ddpprefect/prefect_service.py index c514afa2..d82684b6 100644 --- a/ddpui/ddpprefect/prefect_service.py +++ b/ddpui/ddpprefect/prefect_service.py @@ -630,3 +630,9 @@ def recurse_flow_run_logs( break return logs + + +def get_prefect_version(): + """Fetch secret block id and block name""" + response = prefect_get("prefect/version") + return response From f3f0fbd207c1966ffca45be42484e0f08f1eec25 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Fri, 15 Nov 2024 16:14:14 +0530 Subject: [PATCH 18/57] update toolinfo api for versions --- ddpui/api/org_preferences_api.py | 40 ++++++++++++-------------------- ddpui/ddpdbt/dbt_service.py | 22 +++++++++++------- 2 files changed, 29 insertions(+), 33 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 1fa18820..f3350881 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -112,6 +112,8 @@ def get_tools_versions(request): orguser: OrgUser = request.orguser org = orguser.org + org_superset = OrgSupersets.objects.filter(org=org).first() + versions = [] ver = airbyte_service.get_current_airbyte_version() @@ -122,33 +124,21 @@ def get_tools_versions(request): versions.append({"Prefect": {"version": ver if ver else "Not available"}}) # dbt Version - try: - dbt_version_command = DBT_VERSION_COMMAND - dbt_output = subprocess.check_output(dbt_version_command, text=True) - for line in dbt_output.splitlines(): - if "installed:" in line: - versions.append({"DBT": {"version": line.split(":")[1].strip()}}) - break - else: - versions.append({"DBT": {"version": "Not available"}}) - except Exception: - versions.append({"DBT": {"version": "Not available"}}) - - # Elementary Version - try: - elementary_version_command = ELEMENTARY_VERSION_COMMAND - elementary_output = subprocess.check_output(elementary_version_command, text=True) - for line in elementary_output.splitlines(): - if "Elementary version" in line: - versions.append({"Elementary": {"version": line.split()[-1].strip()}}) - break - else: - versions.append({"Elementary": {"version": "Not available"}}) - except Exception: - versions.append({"Elementary": {"version": "Not available"}}) + ver = dbt_service.get_dbt_version(org) + versions.append({"DBT": {"version": ver if ver else "Not available"}}) + + # elementary Version + ver = dbt_service.get_edr_version(org) + versions.append({"Elementary": {"version": ver if ver else "Not available"}}) # Superset Version - versions.append({"Superset": {"version": org_superset.superset_version}}) + versions.append( + { + "Superset": { + "version": org_superset.superset_version if org_superset else "Not available" + } + } + ) return {"success": True, "res": versions} diff --git a/ddpui/ddpdbt/dbt_service.py b/ddpui/ddpdbt/dbt_service.py index d5ee9c1f..74ec74c9 100644 --- a/ddpui/ddpdbt/dbt_service.py +++ b/ddpui/ddpdbt/dbt_service.py @@ -31,6 +31,7 @@ TASK_DBTSEED, TASK_DBTDEPS, ) +from ddpui.core.orgdbt_manager import DbtProjectManager from ddpui.utils.timezone import as_ist from ddpui.utils.custom_logger import CustomLogger from ddpui.utils.singletaskprogress import SingleTaskProgress @@ -333,29 +334,34 @@ def refresh_elementary_report_via_prefect(orguser: OrgUser) -> dict: return res -def get_dbt_version(): +def get_dbt_version(org: Org): """get dbt version""" - dbt_venv = os.getenv("DBT_VENV") try: - dbt_version_command = [os.path.join(dbt_venv, "venv", "bin", "dbt"), "--version"] + dbt_project_params = DbtProjectManager.gather_dbt_project_params(org, org.dbt) + dbt_version_command = [str(dbt_project_params.dbt_binary), "--version"] dbt_output = subprocess.check_output(dbt_version_command, text=True) for line in dbt_output.splitlines(): if "installed:" in line: return line.split(":")[1].strip() return "Not available" - except Exception: + except Exception as err: + logger.error("Error getting dbt version: %s", err) return "Not available" -def get_edr_version(): +def get_edr_version(org: Org): """get elementary report version""" - dbt_venv = os.getenv("DBT_VENV") try: - elementary_version_command = [os.path.join(dbt_venv, "venv", "bin", "edr"), "--version"] + dbt_project_params = DbtProjectManager.gather_dbt_project_params(org, org.dbt) + elementary_version_command = [ + os.path.join(dbt_project_params.venv_binary, "edr"), + "--version", + ] elementary_output = subprocess.check_output(elementary_version_command, text=True) for line in elementary_output.splitlines(): if "Elementary version" in line: return line.split()[-1].strip() return "Not available" - except Exception: + except Exception as err: + logger.error("Error getting elementary version: %s", err) return "Not available" From 45a3b27b6016ba9ff41c4600ed8380347ac4d605 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Fri, 15 Nov 2024 16:26:06 +0530 Subject: [PATCH 19/57] missing migration --- ..._orgpreferences_trial_end_date_and_more.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 ddpui/migrations/0111_remove_orgpreferences_trial_end_date_and_more.py diff --git a/ddpui/migrations/0111_remove_orgpreferences_trial_end_date_and_more.py b/ddpui/migrations/0111_remove_orgpreferences_trial_end_date_and_more.py new file mode 100644 index 00000000..1e7df27b --- /dev/null +++ b/ddpui/migrations/0111_remove_orgpreferences_trial_end_date_and_more.py @@ -0,0 +1,63 @@ +# Generated by Django 4.2 on 2024-11-15 10:52 + +import ddpui.models.org +from django.db import migrations, models +import django.db.models.deletion +import django.utils.timezone + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0110_userpreferences_discord_webhook_and_more"), + ] + + operations = [ + migrations.RemoveField( + model_name="orgpreferences", + name="trial_end_date", + ), + migrations.RemoveField( + model_name="orgpreferences", + name="trial_start_date", + ), + migrations.AlterField( + model_name="org", + name="type", + field=models.CharField( + choices=[("subscription", "SUBSCRIPTION"), ("trial", "TRIAL"), ("demo", "DEMO")], + default=ddpui.models.org.OrgType["SUBSCRIPTION"], + max_length=50, + ), + ), + migrations.CreateModel( + name="OrgPlans", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, primary_key=True, serialize=False, verbose_name="ID" + ), + ), + ("base_plan", models.CharField(default=None, max_length=255, null=True)), + ("superset_included", models.BooleanField(default=False, null=True)), + ( + "subscription_duration", + models.CharField(default=None, max_length=255, null=True), + ), + ("features", models.JSONField(default=None, null=True)), + ("start_date", models.DateTimeField(blank=True, default=None, null=True)), + ("end_date", models.DateTimeField(blank=True, default=None, null=True)), + ("can_upgrade_plan", models.BooleanField(default=False)), + ("created_at", models.DateTimeField(default=django.utils.timezone.now)), + ("updated_at", models.DateTimeField(default=django.utils.timezone.now)), + ( + "org", + models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, + related_name="org_plans", + to="ddpui.org", + ), + ), + ], + ), + ] From 88cb9df8fb06e1e42d4eb83f50e93c543a1f6f07 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Fri, 15 Nov 2024 16:33:01 +0530 Subject: [PATCH 20/57] api to send email to biz dev team for org plan upgrade --- .env.template | 3 +++ ddpui/api/org_preferences_api.py | 27 +++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/.env.template b/.env.template index 6b720d29..8c68b65b 100644 --- a/.env.template +++ b/.env.template @@ -88,3 +88,6 @@ SENTRY_DSN= SENTRT_TSR= # profile sampling rate SENTRY_PSR= + + +BIZ_DEV_EMAILS="" \ No newline at end of file diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index f3350881..7d5e6093 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -1,3 +1,4 @@ +import os from ninja import Router from ninja.errors import HttpError from django.utils import timezone @@ -18,6 +19,7 @@ from ddpui.ddpprefect import ( prefect_service, ) +from ddpui.utils.awsses import send_text_message orgpreference_router = Router() @@ -154,3 +156,28 @@ def get_org_plans(request): raise HttpError(400, "Org's Plan not found") return {"success": True, "res": org_plan.to_json()} + + +@orgpreference_router.post("/org-plan/upgrade", auth=auth.CustomAuthMiddleware()) +def initiate_upgrade_dalgo_plan(request): + """User can click on the upgrade button from the settings panel + which will trigger email to biz dev team""" + orguser: OrgUser = request.orguser + org = orguser.org + + org_plan = OrgPlans.objects.filter(org=org).first() + + if not org_plan: + raise HttpError(400, "Org's Plan not found") + + biz_dev_emails = os.getenv("BIZ_DEV_EMAILS", []).split(",") + + message = "Upgrade plan request from org: {org_name} with plan: {plan_name}".format( + org_name=org.name, plan_name=org_plan.features + ) + subject = "Upgrade plan request from org: {org_name}".format(org_name=org.name) + + for email in biz_dev_emails: + send_text_message(email, subject, message) + + return {"success": True} From 765e7b73a80f273b10c846a5673c477642c4aaa9 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Fri, 15 Nov 2024 16:52:08 +0530 Subject: [PATCH 21/57] send plan expiry email to account managers - celery task --- ddpui/celeryworkers/tasks.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/ddpui/celeryworkers/tasks.py b/ddpui/celeryworkers/tasks.py index c0f4eefd..dfffb726 100644 --- a/ddpui/celeryworkers/tasks.py +++ b/ddpui/celeryworkers/tasks.py @@ -1,3 +1,4 @@ +import pytz import os import shutil from pathlib import Path @@ -16,6 +17,8 @@ from ddpui.utils import timezone from ddpui.utils.custom_logger import CustomLogger from ddpui.core.orgtaskfunctions import get_edr_send_report_task +from ddpui.utils.awsses import send_text_message +from ddpui.models.org_plans import OrgPlans from ddpui.models.org import ( Org, OrgDbt, @@ -1024,6 +1027,32 @@ def summarize_warehouse_results( return +@app.task() +def check_org_plan_expiry_notify_people(): + """detects schema changes for all the orgs and sends an email to admins if there is a change""" + roles_to_notify = [ACCOUNT_MANAGER_ROLE] + days_before_expiry = 7 + + for org in Org.objects.all(): + org_plan = OrgPlans.objects.filter(org=org).first() + if not org_plan: + continue + + # send a notification 7 days before the plan expires + if org_plan.end_date - timedelta(days=days_before_expiry) < datetime.now(pytz.utc): + try: + org_users = OrgUser.objects.filter( + org=org, + new_role__slug__in=roles_to_notify, + ) + message = f"""This email is to let you know that your Dalgo plan is about to expire. Please renew it to continue using the services.""" + subject = "Dalgo plan expiry" + for orguser in org_users: + send_text_message(orguser.user.email, subject, message) + except Exception as err: + logger.error(err) + + @app.on_after_finalize.connect def setup_periodic_tasks(sender, **kwargs): """check for old locks every minute""" @@ -1039,6 +1068,11 @@ def setup_periodic_tasks(sender, **kwargs): sync_flow_runs_of_deployments.s(), name="sync flow runs of deployments into our db", ) + sender.add_periodic_task( + crontab(minute=0, hour="*/12"), + check_org_plan_expiry_notify_people.s(), + name="check org plan expiry and notify the right people", + ) @app.task(bind=True) From bfe34d744a0a09107606f574796eb3ba065ee0ce Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Fri, 15 Nov 2024 23:19:43 +0530 Subject: [PATCH 22/57] formmated --- ddpui/api/user_preferences_api.py | 16 +++++++++++++--- ddpui/models/org_preferences.py | 4 ++-- ddpui/models/userpreferences.py | 1 + 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index 2a02a1e4..f5a0e033 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -6,6 +6,7 @@ CreateUserPreferencesSchema, UpdateUserPreferencesSchema, ) +from ddpui.models.org_preferences import OrgPreferences from ddpui.models.org_user import OrgUser userpreference_router = Router() @@ -22,6 +23,7 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): user_preferences = UserPreferences.objects.create( orguser=orguser, enable_email_notifications=payload.enable_email_notifications, + llm_optin=payload.llm_optin, ) return {"success": True, "res": user_preferences.to_json()} @@ -36,6 +38,8 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): if payload.enable_email_notifications is not None: user_preferences.enable_email_notifications = payload.enable_email_notifications + if payload.llm_optin is not None: + user_preferences.llm_optin = payload.llm_optin user_preferences.save() return {"success": True, "res": user_preferences.to_json()} @@ -45,7 +49,13 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): def get_user_preferences(request): """gets user preferences for the user""" orguser: OrgUser = request.orguser - + org = orguser.org user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) - - return {"success": True, "res": user_preferences.to_json()} + org_preferences, created = OrgPreferences.objects.get_or_create(org=org) + + res = { + "enable_email_notifications": user_preferences.enable_email_notifications, + "llm_optin": user_preferences.llm_optin, + "is_llm_active": org_preferences.llm_optin, + } + return {"success": True, "res": res} diff --git a/ddpui/models/org_preferences.py b/ddpui/models/org_preferences.py index 45f51940..8c6af246 100644 --- a/ddpui/models/org_preferences.py +++ b/ddpui/models/org_preferences.py @@ -26,11 +26,11 @@ def to_json(self) -> dict: "slug": self.org.slug, "type": self.org.type, }, - "llm_optin": str(self.llm_optin), + "llm_optin": bool(self.llm_optin), "llm_optin_approved_by": ( self.llm_optin_approved_by.user.email if self.llm_optin_approved_by else None ), "llm_optin_date": self.llm_optin_date.isoformat() if self.llm_optin_date else None, - "enable_discord_notifications": str(self.enable_discord_notifications), + "enable_discord_notifications": bool(self.enable_discord_notifications), "discord_webhook": self.discord_webhook, } diff --git a/ddpui/models/userpreferences.py b/ddpui/models/userpreferences.py index 89eb9f4a..8769a31b 100644 --- a/ddpui/models/userpreferences.py +++ b/ddpui/models/userpreferences.py @@ -18,4 +18,5 @@ def to_json(self) -> dict: """Return a dict representation of the model""" return { "enable_email_notifications": self.enable_email_notifications, + "llm_optin": self.llm_optin, } From c2aff928fdd0f3d602b3c2893af5479f47444dbe Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Sat, 16 Nov 2024 00:18:03 +0530 Subject: [PATCH 23/57] a more generic permission for org level notifications --- ddpui/api/org_preferences_api.py | 2 +- seed/002_permissions.json | 128 ++++++++++++++++++++++++------- 2 files changed, 101 insertions(+), 29 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 7d5e6093..dc8a455b 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -69,7 +69,7 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): @orgpreference_router.put("/enable-discord-notifications", auth=auth.CustomAuthMiddleware()) -@has_permission(["can_edit_discord_notifications_settings"]) +@has_permission(["can_edit_org_notification_settings"]) def update_discord_notifications(request, payload: UpdateDiscordNotificationsSchema): """Updates Discord notifications preferences for the logged-in user's organization.""" diff --git a/seed/002_permissions.json b/seed/002_permissions.json index 01122ce4..55d3cbca 100644 --- a/seed/002_permissions.json +++ b/seed/002_permissions.json @@ -2,27 +2,42 @@ { "model": "ddpui.Permission", "pk": 1, - "fields": { "name": "Can View Dashboard", "slug": "can_view_dashboard" } + "fields": { + "name": "Can View Dashboard", + "slug": "can_view_dashboard" + } }, { "model": "ddpui.Permission", "pk": 2, - "fields": { "name": "Can View Sources", "slug": "can_view_sources" } + "fields": { + "name": "Can View Sources", + "slug": "can_view_sources" + } }, { "model": "ddpui.Permission", "pk": 3, - "fields": { "name": "Can Create Source", "slug": "can_create_source" } + "fields": { + "name": "Can Create Source", + "slug": "can_create_source" + } }, { "model": "ddpui.Permission", "pk": 4, - "fields": { "name": "Can Edit Source", "slug": "can_edit_source" } + "fields": { + "name": "Can Edit Source", + "slug": "can_edit_source" + } }, { "model": "ddpui.Permission", "pk": 5, - "fields": { "name": "Can View Source", "slug": "can_view_source" } + "fields": { + "name": "Can View Source", + "slug": "can_view_source" + } }, { "model": "ddpui.Permission", @@ -35,7 +50,10 @@ { "model": "ddpui.Permission", "pk": 7, - "fields": { "name": "Can View Warehouse", "slug": "can_view_warehouse" } + "fields": { + "name": "Can View Warehouse", + "slug": "can_view_warehouse" + } }, { "model": "ddpui.Permission", @@ -48,7 +66,10 @@ { "model": "ddpui.Permission", "pk": 9, - "fields": { "name": "Can Edit Warehouse", "slug": "can_edit_warehouse" } + "fields": { + "name": "Can Edit Warehouse", + "slug": "can_edit_warehouse" + } }, { "model": "ddpui.Permission", @@ -61,7 +82,10 @@ { "model": "ddpui.Permission", "pk": 11, - "fields": { "name": "Can Create Org", "slug": "can_create_org" } + "fields": { + "name": "Can Create Org", + "slug": "can_create_org" + } }, { "model": "ddpui.Permission", @@ -106,17 +130,26 @@ { "model": "ddpui.Permission", "pk": 17, - "fields": { "name": "Can Delete Source", "slug": "can_delete_source" } + "fields": { + "name": "Can Delete Source", + "slug": "can_delete_source" + } }, { "model": "ddpui.Permission", "pk": 18, - "fields": { "name": "View Tasks", "slug": "can_view_master_tasks" } + "fields": { + "name": "View Tasks", + "slug": "can_view_master_tasks" + } }, { "model": "ddpui.Permission", "pk": 19, - "fields": { "name": "View Task Config", "slug": "can_view_master_task" } + "fields": { + "name": "View Task Config", + "slug": "can_view_master_task" + } }, { "model": "ddpui.Permission", @@ -153,7 +186,10 @@ { "model": "ddpui.Permission", "pk": 24, - "fields": { "name": "Can Run Orgtask", "slug": "can_run_orgtask" } + "fields": { + "name": "Can Run Orgtask", + "slug": "can_run_orgtask" + } }, { "model": "ddpui.Permission", @@ -166,17 +202,26 @@ { "model": "ddpui.Permission", "pk": 26, - "fields": { "name": "Can Create Orgtask", "slug": "can_create_orgtask" } + "fields": { + "name": "Can Create Orgtask", + "slug": "can_create_orgtask" + } }, { "model": "ddpui.Permission", "pk": 27, - "fields": { "name": "Can View Orgtasks", "slug": "can_view_orgtasks" } + "fields": { + "name": "Can View Orgtasks", + "slug": "can_view_orgtasks" + } }, { "model": "ddpui.Permission", "pk": 28, - "fields": { "name": "Can Delete Orgtask", "slug": "can_delete_orgtask" } + "fields": { + "name": "Can Delete Orgtask", + "slug": "can_delete_orgtask" + } }, { "model": "ddpui.Permission", @@ -189,12 +234,18 @@ { "model": "ddpui.Permission", "pk": 30, - "fields": { "name": "Can View Pipelines", "slug": "can_view_pipelines" } + "fields": { + "name": "Can View Pipelines", + "slug": "can_view_pipelines" + } }, { "model": "ddpui.Permission", "pk": 31, - "fields": { "name": "Can View Pipeline", "slug": "can_view_pipeline" } + "fields": { + "name": "Can View Pipeline", + "slug": "can_view_pipeline" + } }, { "model": "ddpui.Permission", @@ -207,12 +258,18 @@ { "model": "ddpui.Permission", "pk": 33, - "fields": { "name": "Can Edit Pipeline", "slug": "can_edit_pipeline" } + "fields": { + "name": "Can Edit Pipeline", + "slug": "can_edit_pipeline" + } }, { "model": "ddpui.Permission", "pk": 34, - "fields": { "name": "Can Run Pipeline", "slug": "can_run_pipeline" } + "fields": { + "name": "Can Run Pipeline", + "slug": "can_run_pipeline" + } }, { "model": "ddpui.Permission", @@ -233,7 +290,10 @@ { "model": "ddpui.Permission", "pk": 37, - "fields": { "name": "Can Sync Sources", "slug": "can_sync_sources" } + "fields": { + "name": "Can Sync Sources", + "slug": "can_sync_sources" + } }, { "model": "ddpui.Permission", @@ -262,7 +322,10 @@ { "model": "ddpui.Permission", "pk": 41, - "fields": { "name": "Can Edit Dbt Model", "slug": "can_edit_dbt_model" } + "fields": { + "name": "Can Edit Dbt Model", + "slug": "can_edit_dbt_model" + } }, { "model": "ddpui.Permission", @@ -299,7 +362,10 @@ { "model": "ddpui.Permission", "pk": 46, - "fields": { "name": "Can View Org Users", "slug": "can_view_orgusers" } + "fields": { + "name": "Can View Org Users", + "slug": "can_view_orgusers" + } }, { "model": "ddpui.Permission", @@ -320,7 +386,10 @@ { "model": "ddpui.Permission", "pk": 49, - "fields": { "name": "Can Edit Org User", "slug": "can_edit_orguser" } + "fields": { + "name": "Can Edit Org User", + "slug": "can_edit_orguser" + } }, { "model": "ddpui.Permission", @@ -341,7 +410,10 @@ { "model": "ddpui.Permission", "pk": 52, - "fields": { "name": "Public", "slug": "public" } + "fields": { + "name": "Public", + "slug": "public" + } }, { "model": "ddpui.Permission", @@ -403,7 +475,7 @@ "model": "ddpui.Permission", "pk": 60, "fields": { - "name": "Can Edit LLM Settings", + "name": "Can Edit LLM settings", "slug": "can_edit_llm_settings" } }, @@ -411,8 +483,8 @@ "model": "ddpui.Permission", "pk": 61, "fields": { - "name": "Can Edit Discord Notifications Settings", - "slug": "can_edit_discord_notifications_settings" + "name": "Can edit org level notifications settings like discord", + "slug": "can_edit_org_notification_settings" } } -] +] \ No newline at end of file From 38eb5639539419d9b22936f01c6cbc8e8b13dcbe Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Sat, 16 Nov 2024 01:05:29 +0530 Subject: [PATCH 24/57] test case changes - updated the api function name for creating notification --- ddpui/api/notifications_api.py | 2 +- ddpui/tests/api_tests/test_notifications_api.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/ddpui/api/notifications_api.py b/ddpui/api/notifications_api.py index b11170b5..e530282a 100644 --- a/ddpui/api/notifications_api.py +++ b/ddpui/api/notifications_api.py @@ -13,7 +13,7 @@ @notification_router.post("/") -def create_notification(request, payload: CreateNotificationPayloadSchema): +def post_create_notification(request, payload: CreateNotificationPayloadSchema): """Handle the task of creating a notification""" # Filter OrgUser data based on sent_to field diff --git a/ddpui/tests/api_tests/test_notifications_api.py b/ddpui/tests/api_tests/test_notifications_api.py index d23e2572..1e1c2ef6 100644 --- a/ddpui/tests/api_tests/test_notifications_api.py +++ b/ddpui/tests/api_tests/test_notifications_api.py @@ -9,7 +9,7 @@ from ddpui import auth from django.contrib.auth.models import User from ddpui.api.notifications_api import ( - create_notification, + post_create_notification, get_notification_history, get_notification_recipients, get_user_notifications, @@ -77,7 +77,7 @@ def test_create_notification_success(mock_create_notification, mock_get_recipien create_notification_payload = CreateNotificationPayloadSchema(**payload) mock_get_recipients.return_value = (None, [1, 2, 3]) mock_create_notification.return_value = (None, {"res": [], "errors": []}) - response = create_notification(MagicMock(), create_notification_payload) + response = post_create_notification(MagicMock(), create_notification_payload) assert isinstance(response["res"], list) assert isinstance(response["errors"], list) mock_get_recipients.assert_called_once_with( @@ -117,7 +117,7 @@ def test_create_notification_no_recipients(mock_create_notification, mock_get_re create_notification_payload = CreateNotificationPayloadSchema(**payload) mock_get_recipients.return_value = ("No users found for the given information", []) with pytest.raises(HttpError) as excinfo: - create_notification(MagicMock(), create_notification_payload) + post_create_notification(MagicMock(), create_notification_payload) assert "No users found for the given information" in str(excinfo.value) mock_get_recipients.assert_called_once_with( payload["sent_to"], @@ -151,7 +151,7 @@ def test_create_notification_no_org_slug(mock_create_notification, mock_get_reci [], ) with pytest.raises(HttpError) as excinfo: - create_notification(MagicMock(), create_notification_payload) + post_create_notification(MagicMock(), create_notification_payload) assert "org_slug is required to sent notification to all org users." in str(excinfo.value) mock_get_recipients.assert_called_once_with( payload["sent_to"], @@ -185,7 +185,7 @@ def test_create_notification_no_user_email(mock_create_notification, mock_get_re [], ) with pytest.raises(HttpError) as excinfo: - create_notification(MagicMock(), create_notification_payload) + post_create_notification(MagicMock(), create_notification_payload) assert "user email is required to sent notification to a user." in str(excinfo.value) mock_get_recipients.assert_called_once_with( payload["sent_to"], @@ -219,7 +219,7 @@ def test_create_notification_user_does_not_exist(mock_create_notification, mock_ [], ) with pytest.raises(HttpError) as excinfo: - create_notification(MagicMock(), create_notification_payload) + post_create_notification(MagicMock(), create_notification_payload) assert "User with the provided email does not exist" in str(excinfo.value) mock_get_recipients.assert_called_once_with( payload["sent_to"], From 551b2700f58a95fc95b09bb2efd57a326dc97e98 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Sat, 16 Nov 2024 01:16:50 +0530 Subject: [PATCH 25/57] permissions and api for roles below account maager to request enabling of llm analysis feature --- ddpui/api/org_preferences_api.py | 1 + ddpui/api/user_preferences_api.py | 33 +++++++++++++++++++ ddpui/core/notifications_service.py | 14 ++++---- .../commands/create_notification.py | 15 +++++---- ddpui/schemas/notifications_api_schemas.py | 10 ++++++ .../tests/core/test_notifications_service.py | 17 +++++----- seed/002_permissions.json | 16 +++++++++ seed/003_role_permissions.json | 24 ++++++++++++++ 8 files changed, 108 insertions(+), 22 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index dc8a455b..d586e173 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -159,6 +159,7 @@ def get_org_plans(request): @orgpreference_router.post("/org-plan/upgrade", auth=auth.CustomAuthMiddleware()) +@has_permission(["can_initiate_org_plan_upgrade"]) def initiate_upgrade_dalgo_plan(request): """User can click on the upgrade button from the settings panel which will trigger email to biz dev team""" diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index f5a0e033..fcbfa719 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -8,6 +8,10 @@ ) from ddpui.models.org_preferences import OrgPreferences from ddpui.models.org_user import OrgUser +from ddpui.auth import ACCOUNT_MANAGER_ROLE +from ddpui.core.notifications_service import create_notification +from ddpui.auth import has_permission +from ddpui.schemas.notifications_api_schemas import NotificationDataSchema userpreference_router = Router() @@ -59,3 +63,32 @@ def get_user_preferences(request): "is_llm_active": org_preferences.llm_optin, } return {"success": True, "res": res} + + +@userpreference_router.post("/llm_analysis/request", auth=auth.CustomAuthMiddleware()) +@has_permission(["can_request_llm_analysis_feature"]) +def post_request_llm_analysis_feature_enabled(request): + """Sends a notification to org's account manager for enabling LLM analysis feature""" + orguser: OrgUser = request.orguser + org = orguser.org + + # get the account managers of the org + acc_managers: list[OrgUser] = OrgUser.objects.filter( + org=org, new_role__slug=ACCOUNT_MANAGER_ROLE + ).all() + + if len(acc_managers) == 0: + raise HttpError(400, "No account manager found for the organization") + + # send notification to all account managers + notification_data = NotificationDataSchema( + author=orguser.user.email, + message=f"{orguser.user.first_name} is requesting to enable LLM analysis feature", + urgent=False, + scheduled_time=None, + recipients=[acc_manager.id for acc_manager in acc_managers], + ) + + create_notification(notification_data) + + return {"success": True, "res": "Notified account manager(s) to enable LLM analysis feature"} diff --git a/ddpui/core/notifications_service.py b/ddpui/core/notifications_service.py index 07e96052..fa3ba245 100644 --- a/ddpui/core/notifications_service.py +++ b/ddpui/core/notifications_service.py @@ -11,7 +11,7 @@ from ddpui.utils import timezone from ddpui.utils.discord import send_discord_notification from ddpui.utils.sendgrid import send_email_notification -from ddpui.schemas.notifications_api_schemas import SentToEnum +from ddpui.schemas.notifications_api_schemas import SentToEnum, NotificationDataSchema from ddpui.celeryworkers.tasks import schedule_notification_task @@ -102,18 +102,18 @@ def handle_recipient( # main function for sending notification def create_notification( - notification_data, + notification_data: NotificationDataSchema, ) -> Tuple[Optional[Dict[str, str]], Optional[Dict[str, Any]]]: """ main function for creating notification. Add notification to the notification table. """ - author = notification_data["author"] - message = notification_data["message"] - urgent = notification_data["urgent"] - scheduled_time = notification_data["scheduled_time"] - recipients = notification_data["recipients"] + author = notification_data.author + message = notification_data.message + urgent = notification_data.urgent + scheduled_time = notification_data.scheduled_time + recipients = notification_data.recipients errors = [] notification = Notification.objects.create( diff --git a/ddpui/management/commands/create_notification.py b/ddpui/management/commands/create_notification.py index d32239a2..62b09e5e 100644 --- a/ddpui/management/commands/create_notification.py +++ b/ddpui/management/commands/create_notification.py @@ -5,6 +5,7 @@ from ddpui.schemas.notifications_api_schemas import ( CreateNotificationPayloadSchema, SentToEnum, + NotificationDataSchema, ) @@ -83,13 +84,13 @@ def handle(self, *args, **options): sys.exit(1) # Create notification data - notification_data = { - "author": payload.author, - "message": payload.message, - "urgent": payload.urgent, - "scheduled_time": payload.scheduled_time, - "recipients": recipients, - } + notification_data = NotificationDataSchema( + author=payload.author, + message=payload.message, + urgent=payload.urgent, + scheduled_time=payload.scheduled_time, + recipients=recipients, + ) # Call the create notification service error, result = notifications_service.create_notification(notification_data) diff --git a/ddpui/schemas/notifications_api_schemas.py b/ddpui/schemas/notifications_api_schemas.py index 367cc252..8793b2b4 100644 --- a/ddpui/schemas/notifications_api_schemas.py +++ b/ddpui/schemas/notifications_api_schemas.py @@ -44,3 +44,13 @@ class UpdateReadStatusSchemav1(Schema): notification_ids: list[int] read_status: bool + + +class NotificationDataSchema(Schema): + """Schema use to call the notification service function for creating a notification""" + + author: str + message: str + urgent: Optional[bool] = False + scheduled_time: Optional[datetime] = None + recipients: List[int] # list of orguser ids diff --git a/ddpui/tests/core/test_notifications_service.py b/ddpui/tests/core/test_notifications_service.py index 546bf9cd..7d067d43 100644 --- a/ddpui/tests/core/test_notifications_service.py +++ b/ddpui/tests/core/test_notifications_service.py @@ -21,7 +21,7 @@ delete_scheduled_notification, get_unread_notifications_count, ) -from ddpui.schemas.notifications_api_schemas import SentToEnum +from ddpui.schemas.notifications_api_schemas import SentToEnum, NotificationDataSchema from ddpui.tests.api_tests.test_user_org_api import mock_request, seed_db from django.contrib.auth.models import User @@ -209,13 +209,14 @@ def test_handle_recipient_discord_error(mocker: Mock, orguser, unsent_notificati def test_create_notification_success(orguser): - notification_data = { - "author": "test_author", - "message": "test_message", - "urgent": True, - "scheduled_time": None, - "recipients": [orguser.id], - } + notification_data = NotificationDataSchema( + author="test_author", + message="test_message", + urgent=True, + scheduled_time=None, + recipients=[orguser.id], + ) + error, result = create_notification(notification_data) assert error is None assert result is not None diff --git a/seed/002_permissions.json b/seed/002_permissions.json index 55d3cbca..e1586324 100644 --- a/seed/002_permissions.json +++ b/seed/002_permissions.json @@ -486,5 +486,21 @@ "name": "Can edit org level notifications settings like discord", "slug": "can_edit_org_notification_settings" } + }, + { + "model": "ddpui.Permission", + "pk": 62, + "fields": { + "name": "Can initiate upgrade of Dalgo subscription for the org", + "slug": "can_initiate_org_plan_upgrade" + } + }, + { + "model": "ddpui.Permission", + "pk": 63, + "fields": { + "name": "Can request (to the higher role) llm analysis feature to be enabled", + "slug": "can_request_llm_analysis_feature" + } } ] \ No newline at end of file diff --git a/seed/003_role_permissions.json b/seed/003_role_permissions.json index 72ff4a03..65cee4e5 100644 --- a/seed/003_role_permissions.json +++ b/seed/003_role_permissions.json @@ -1806,5 +1806,29 @@ "role": 1, "permission": 61 } + }, + { + "model": "ddpui.RolePermission", + "pk": 223, + "fields": { + "role": 3, + "permission": 63 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 224, + "fields": { + "role": 4, + "permission": 63 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 225, + "fields": { + "role": 5, + "permission": 63 + } } ] \ No newline at end of file From 5db249762ddadc1c1b67ebd636b8a6273a04c619 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 22:36:54 +0530 Subject: [PATCH 26/57] userpermissions2orgpermissions --- .../userpermissions2orgpermissions.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 ddpui/management/commands/userpermissions2orgpermissions.py diff --git a/ddpui/management/commands/userpermissions2orgpermissions.py b/ddpui/management/commands/userpermissions2orgpermissions.py new file mode 100644 index 00000000..8b8106d8 --- /dev/null +++ b/ddpui/management/commands/userpermissions2orgpermissions.py @@ -0,0 +1,62 @@ +from datetime import datetime +from django.core.management.base import BaseCommand + +from ddpui.models.org import Org +from ddpui.models.org_user import OrgUser +from ddpui.models.org_preferences import OrgPreferences +from ddpui.models.userpreferences import UserPreferences + + +class Command(BaseCommand): + """ + This script creates OrgPermissions from UserPermissions + """ + + help = "Edit a dbt cli profile" + + def add_arguments(self, parser): + parser.add_argument("--org", type=str, help="Org slug, use 'all' to update all orgs") + + def handle(self, *args, **options): + q = Org.objects + if options["org"] != "all": + q = q.filter(slug=options["org"]) + if q.count() == 0: + print("No orgs found") + return + for org in q.all(): + print("Processing org " + org.slug) + + orgpreferences = OrgPreferences.objects.filter(org=org).first() + if orgpreferences is None: + print("creating org preferences for " + org.slug) + orgpreferences = OrgPreferences.objects.create(org=org) + + for orguser in OrgUser.objects.filter(org=org): + userpreferences = UserPreferences.objects.filter(orguser=orguser).first() + + if userpreferences is not None: + print("Found user preferences for " + orguser.user.email) + + if orgpreferences.llm_optin is False and userpreferences.llm_optin is True: + print("Approving LLM opt-in by " + orguser.user.email) + orgpreferences.llm_optin = True + orgpreferences.llm_optin_approved_by = orguser + orgpreferences.llm_optin_date = datetime.now() + + if ( + orgpreferences.enable_discord_notifications is False + and userpreferences.enable_discord_notifications is True + and userpreferences.discord_webhook is not None + ): + # use the discord webbhook from the first user we find + print( + "Discord notifications enabled by " + + orguser.user.email + + ", settings webook to " + + userpreferences.discord_webhook + ) + orgpreferences.enable_discord_notifications = True + orgpreferences.discord_webhook = userpreferences.discord_webhook + + orgpreferences.save() From d24d2c54690f92e2e4362982f67b66baa70640f3 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 22:41:56 +0530 Subject: [PATCH 27/57] pylint --- ddpui/models/org_plans.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ddpui/models/org_plans.py b/ddpui/models/org_plans.py index 59a0d89c..b7220fb6 100644 --- a/ddpui/models/org_plans.py +++ b/ddpui/models/org_plans.py @@ -1,7 +1,6 @@ from django.db import models -from ddpui.models.org import Org -from ddpui.models.org_user import OrgUser from django.utils import timezone +from ddpui.models.org import Org class OrgPlans(models.Model): From 46d6cc06569b673b926342bfecbe1cc67de4f737 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 22:42:57 +0530 Subject: [PATCH 28/57] fix help string --- ddpui/management/commands/userpermissions2orgpermissions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ddpui/management/commands/userpermissions2orgpermissions.py b/ddpui/management/commands/userpermissions2orgpermissions.py index 8b8106d8..fefdd682 100644 --- a/ddpui/management/commands/userpermissions2orgpermissions.py +++ b/ddpui/management/commands/userpermissions2orgpermissions.py @@ -12,7 +12,7 @@ class Command(BaseCommand): This script creates OrgPermissions from UserPermissions """ - help = "Edit a dbt cli profile" + help = "Create OrgPermissions from UserPermissions" def add_arguments(self, parser): parser.add_argument("--org", type=str, help="Org slug, use 'all' to update all orgs") From 78c2e5d757a6096006cc3cd22a5b9583797995ee Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 22:52:37 +0530 Subject: [PATCH 29/57] createorgplan --- ddpui/management/commands/createorgplan.py | 73 ++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 ddpui/management/commands/createorgplan.py diff --git a/ddpui/management/commands/createorgplan.py b/ddpui/management/commands/createorgplan.py new file mode 100644 index 00000000..0833a938 --- /dev/null +++ b/ddpui/management/commands/createorgplan.py @@ -0,0 +1,73 @@ +from datetime import datetime +from django.core.management.base import BaseCommand + +from ddpui.models.org import Org +from ddpui.models.org_plans import OrgPlans + + +class Command(BaseCommand): + """ + This script creates OrgPlans for Orgs + """ + + help = "Create an OrgPlan for an Org" + + def add_arguments(self, parser): + parser.add_argument("--org", type=str, help="Org slug", required=True) + parser.add_argument("--with-superset", action="store_true", help="Include superset") + parser.add_argument( + "--duration", + choices=["Monthly", "Annual"], + help="Subscription duration", + required=True, + ) + parser.add_argument("--start-date", type=str, help="Start date", required=False) + parser.add_argument("--end-date", type=str, help="Start date", required=False) + + def handle(self, *args, **options): + """create the OrgPlan for the Org""" + org = Org.objects.filter(slug=options["org"]).first() + if org is None: + self.stdout.write(self.style.ERROR(f"Org {options['org']} not found")) + return + + if OrgPlans.objects.filter(org=org).exists(): + self.stdout.write(self.style.ERROR(f"Org {options['org']} already has a plan")) + return + + start_date = ( + datetime.strptime(options["start_date"], "%Y-%m-%d") if options["start_date"] else None + ) + end_date = ( + datetime.strptime(options["end_date"], "%Y-%m-%d") if options["end_date"] else None + ) + + if options["with_superset"]: + base_plan = "DALGO + Superset" + can_upgrade_plan = False + features = { + "pipeline": ["Ingest", "Transform", "Orchestrate"], + "aiFeatures": ["AI data analysis"], + "dataQuality": ["Data quality dashboards"], + } + else: + base_plan = "DALGO" + can_upgrade_plan = True + features = { + "pipeline": ["Ingest", "Transform", "Orchestrate"], + "aiFeatures": ["AI data analysis"], + "dataQuality": ["Data quality dashboards"], + "superset": ["Superset dashboards", "Superset Usage dashboards"], + } + + org_plan = OrgPlans.objects.create( + org=org, + base_plan=base_plan, + superset_included=options["with_superset"], + subscription_duration=options["duration"], + start_date=start_date, + end_date=end_date, + can_upgrade_plan=can_upgrade_plan, + features=features, + ) + print(org_plan.to_json()) From 7adfba602be546d4a0290764ce43d27f07de8821 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 22:54:41 +0530 Subject: [PATCH 30/57] --overwrite flag --- ddpui/management/commands/createorgplan.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ddpui/management/commands/createorgplan.py b/ddpui/management/commands/createorgplan.py index 0833a938..7f1a7396 100644 --- a/ddpui/management/commands/createorgplan.py +++ b/ddpui/management/commands/createorgplan.py @@ -23,6 +23,7 @@ def add_arguments(self, parser): ) parser.add_argument("--start-date", type=str, help="Start date", required=False) parser.add_argument("--end-date", type=str, help="Start date", required=False) + parser.add_argument("--overwrite", action="store_true", help="Overwrite existing plan") def handle(self, *args, **options): """create the OrgPlan for the Org""" @@ -31,7 +32,7 @@ def handle(self, *args, **options): self.stdout.write(self.style.ERROR(f"Org {options['org']} not found")) return - if OrgPlans.objects.filter(org=org).exists(): + if OrgPlans.objects.filter(org=org).exists() and not options["overwrite"]: self.stdout.write(self.style.ERROR(f"Org {options['org']} already has a plan")) return From 0db2e084ce616f3e298e37169f72446acf8e89db Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 22:57:27 +0530 Subject: [PATCH 31/57] create or update --- ddpui/management/commands/createorgplan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ddpui/management/commands/createorgplan.py b/ddpui/management/commands/createorgplan.py index 7f1a7396..0cb2ed4c 100644 --- a/ddpui/management/commands/createorgplan.py +++ b/ddpui/management/commands/createorgplan.py @@ -61,7 +61,7 @@ def handle(self, *args, **options): "superset": ["Superset dashboards", "Superset Usage dashboards"], } - org_plan = OrgPlans.objects.create( + org_plan = OrgPlans.objects.update_or_create( org=org, base_plan=base_plan, superset_included=options["with_superset"], From edf2cfcd86226ca72fefa4fe49bc1204c083364c Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 23:04:14 +0530 Subject: [PATCH 32/57] rewrite --- ddpui/management/commands/createorgplan.py | 37 +++++++++++----------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/ddpui/management/commands/createorgplan.py b/ddpui/management/commands/createorgplan.py index 0cb2ed4c..9d15a0f3 100644 --- a/ddpui/management/commands/createorgplan.py +++ b/ddpui/management/commands/createorgplan.py @@ -32,43 +32,42 @@ def handle(self, *args, **options): self.stdout.write(self.style.ERROR(f"Org {options['org']} not found")) return - if OrgPlans.objects.filter(org=org).exists() and not options["overwrite"]: + org_plan = OrgPlans.objects.filter(org=org).first() + if org_plan and not options["overwrite"]: self.stdout.write(self.style.ERROR(f"Org {options['org']} already has a plan")) return - start_date = ( + if not org_plan: + org_plan = OrgPlans(org=org) + + org_plan.superset_included = options["with_superset"] + org_plan.subscription_duration = options["duration"] + + org_plan.start_date = ( datetime.strptime(options["start_date"], "%Y-%m-%d") if options["start_date"] else None ) - end_date = ( + org_plan.end_date = ( datetime.strptime(options["end_date"], "%Y-%m-%d") if options["end_date"] else None ) if options["with_superset"]: - base_plan = "DALGO + Superset" - can_upgrade_plan = False - features = { + org_plan.base_plan = "DALGO + Superset" + org_plan.can_upgrade_plan = False + org_plan.features = { "pipeline": ["Ingest", "Transform", "Orchestrate"], "aiFeatures": ["AI data analysis"], "dataQuality": ["Data quality dashboards"], } else: - base_plan = "DALGO" - can_upgrade_plan = True - features = { + org_plan.base_plan = "DALGO" + org_plan.can_upgrade_plan = True + org_plan.features = { "pipeline": ["Ingest", "Transform", "Orchestrate"], "aiFeatures": ["AI data analysis"], "dataQuality": ["Data quality dashboards"], "superset": ["Superset dashboards", "Superset Usage dashboards"], } - org_plan = OrgPlans.objects.update_or_create( - org=org, - base_plan=base_plan, - superset_included=options["with_superset"], - subscription_duration=options["duration"], - start_date=start_date, - end_date=end_date, - can_upgrade_plan=can_upgrade_plan, - features=features, - ) + org_plan.save() + print(org_plan.to_json()) From 5da4e0c51e9c02a7f2f41ed19ed31849c6166d10 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 23:08:03 +0530 Subject: [PATCH 33/57] drop the trailing period --- ddpui/ddpdbt/dbt_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ddpui/ddpdbt/dbt_service.py b/ddpui/ddpdbt/dbt_service.py index 74ec74c9..baf22b46 100644 --- a/ddpui/ddpdbt/dbt_service.py +++ b/ddpui/ddpdbt/dbt_service.py @@ -359,8 +359,8 @@ def get_edr_version(org: Org): ] elementary_output = subprocess.check_output(elementary_version_command, text=True) for line in elementary_output.splitlines(): - if "Elementary version" in line: - return line.split()[-1].strip() + if line.startswith("Elementary version"): + return line.split()[-1].strip()[:-1] return "Not available" except Exception as err: logger.error("Error getting elementary version: %s", err) From e1454aed5d4406e7e12a2639cdd7d0c9fdbb9de8 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 23:13:39 +0530 Subject: [PATCH 34/57] createorgsuperset --- .../management/commands/createorgsuperset.py | 35 +++++++++++++++++++ ddpui/models/org_supersets.py | 2 +- 2 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 ddpui/management/commands/createorgsuperset.py diff --git a/ddpui/management/commands/createorgsuperset.py b/ddpui/management/commands/createorgsuperset.py new file mode 100644 index 00000000..32115336 --- /dev/null +++ b/ddpui/management/commands/createorgsuperset.py @@ -0,0 +1,35 @@ +from django.core.management.base import BaseCommand + +from ddpui.models.org import Org +from ddpui.models.org_supersets import OrgSupersets + + +class Command(BaseCommand): + """ + This script creates OrgSupersets for Orgs + """ + + help = "Create an OrgSuperset for an Org" + + def add_arguments(self, parser): + parser.add_argument("--org", type=str, help="Org slug", required=True) + parser.add_argument("--container-name", type=str, help="Container name", required=True) + parser.add_argument("--superset-version", type=str, help="Superset version", required=True) + parser.add_argument("--overwrite", action="store_true", help="Overwrite existing plan") + + def handle(self, *args, **options): + org = Org.objects.get(slug=options["org"]) + + org_superset = OrgSupersets.objects.filter(org=org).first() + if org_superset and not options["overwrite"]: + self.stdout.write(self.style.ERROR(f"Org {options['org']} already has a superset")) + return + + if not org_superset: + org_superset = OrgSupersets(org=org) + + org_superset.container_name = options["container_name"] + org_superset.superset_version = options["superset_version"] + + org_superset.save() + print("OrgSuperset created successfully for " + org.slug) diff --git a/ddpui/models/org_supersets.py b/ddpui/models/org_supersets.py index f60c3441..7395b8e1 100644 --- a/ddpui/models/org_supersets.py +++ b/ddpui/models/org_supersets.py @@ -1,6 +1,6 @@ from django.db import models -from ddpui.models.org import Org from django.utils import timezone +from ddpui.models.org import Org class OrgSupersets(models.Model): From 619a7edc0742cf67e52087f41487aa32f8837900 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 23:30:27 +0530 Subject: [PATCH 35/57] remove discord from user preferences --- ddpui/api/user_preferences_api.py | 4 ---- ddpui/tests/api_tests/test_user_preferences_api.py | 4 +--- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index fcbfa719..e928c463 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -53,14 +53,10 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): def get_user_preferences(request): """gets user preferences for the user""" orguser: OrgUser = request.orguser - org = orguser.org user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) - org_preferences, created = OrgPreferences.objects.get_or_create(org=org) res = { "enable_email_notifications": user_preferences.enable_email_notifications, - "llm_optin": user_preferences.llm_optin, - "is_llm_active": org_preferences.llm_optin, } return {"success": True, "res": res} diff --git a/ddpui/tests/api_tests/test_user_preferences_api.py b/ddpui/tests/api_tests/test_user_preferences_api.py index 841fb6ee..399abf2d 100644 --- a/ddpui/tests/api_tests/test_user_preferences_api.py +++ b/ddpui/tests/api_tests/test_user_preferences_api.py @@ -62,9 +62,7 @@ def user_preferences(orguser): """a pytest fixture which creates the user preferences for the OrgUser""" return UserPreferences.objects.create( orguser=orguser, - enable_discord_notifications=True, enable_email_notifications=True, - discord_webhook="http://example.com/webhook", ) @@ -155,7 +153,7 @@ def test_get_user_preferences_success(orguser, user_preferences): response = get_user_preferences(request) assert response["success"] is True assert response["res"] == { - "enable_discord_notifications": user_preferences.enable_discord_notifications, + "enable_email_notifications": user_preferences.enable_email_notifications, } From 3633bde6ede5ab3660975d333ba5742cf84af6ec Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sat, 16 Nov 2024 23:38:10 +0530 Subject: [PATCH 36/57] removing deprecated user-preferences fields --- ddpui/api/user_preferences_api.py | 3 --- ddpui/schemas/userpreferences_schema.py | 6 ------ ddpui/tests/api_tests/test_user_preferences_api.py | 14 -------------- 3 files changed, 23 deletions(-) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index e928c463..a78a10ae 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -27,7 +27,6 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): user_preferences = UserPreferences.objects.create( orguser=orguser, enable_email_notifications=payload.enable_email_notifications, - llm_optin=payload.llm_optin, ) return {"success": True, "res": user_preferences.to_json()} @@ -42,8 +41,6 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): if payload.enable_email_notifications is not None: user_preferences.enable_email_notifications = payload.enable_email_notifications - if payload.llm_optin is not None: - user_preferences.llm_optin = payload.llm_optin user_preferences.save() return {"success": True, "res": user_preferences.to_json()} diff --git a/ddpui/schemas/userpreferences_schema.py b/ddpui/schemas/userpreferences_schema.py index f5c6c482..e16e1593 100644 --- a/ddpui/schemas/userpreferences_schema.py +++ b/ddpui/schemas/userpreferences_schema.py @@ -5,16 +5,10 @@ class CreateUserPreferencesSchema(Schema): """Schema for creating user preferences for the user.""" - enable_discord_notifications: bool - discord_webhook: Optional[str] = None enable_email_notifications: bool - llm_optin: bool class UpdateUserPreferencesSchema(Schema): """Schema for updating user preferences for the user.""" - enable_discord_notifications: Optional[bool] = None - discord_webhook: Optional[str] = None enable_email_notifications: Optional[bool] = None - llm_optin: Optional[bool] diff --git a/ddpui/tests/api_tests/test_user_preferences_api.py b/ddpui/tests/api_tests/test_user_preferences_api.py index 399abf2d..8f5ef7b3 100644 --- a/ddpui/tests/api_tests/test_user_preferences_api.py +++ b/ddpui/tests/api_tests/test_user_preferences_api.py @@ -77,9 +77,7 @@ def test_create_user_preferences_success(orguser): """tests the success of creating user preferences for the OrgUser""" request = mock_request(orguser) payload = CreateUserPreferencesSchema( - enable_discord_notifications=True, enable_email_notifications=True, - discord_webhook="http://example.com/webhook", ) response = create_user_preferences(request, payload) @@ -87,9 +85,7 @@ def test_create_user_preferences_success(orguser): # Assertions assert response["success"] is True preferences = response["res"] - assert preferences["discord_webhook"] == "http://example.com/webhook" assert preferences["enable_email_notifications"] is True - assert preferences["enable_discord_notifications"] is True def test_create_user_preferences_already_exists(user_preferences): @@ -99,9 +95,7 @@ def test_create_user_preferences_already_exists(user_preferences): """ request = mock_request(orguser=user_preferences.orguser) payload = CreateUserPreferencesSchema( - enable_discord_notifications=True, enable_email_notifications=True, - discord_webhook="http://example.com/webhook", ) with pytest.raises(HttpError) as excinfo: @@ -114,17 +108,13 @@ def test_update_user_preferences_success(orguser, user_preferences): """tests the success of updating user preferences for the OrgUser""" request = mock_request(orguser) payload = UpdateUserPreferencesSchema( - enable_discord_notifications=False, enable_email_notifications=False, - discord_webhook="http://example.org/webhook", ) response = update_user_preferences(request, payload) assert response["success"] is True updated_preferences = UserPreferences.objects.get(orguser=user_preferences.orguser) - assert updated_preferences.enable_discord_notifications is False assert updated_preferences.enable_email_notifications is False - assert updated_preferences.discord_webhook == "http://example.org/webhook" def test_update_user_preferences_create_success_if_not_exist(orguser): @@ -134,17 +124,13 @@ def test_update_user_preferences_create_success_if_not_exist(orguser): """ request = mock_request(orguser) payload = UpdateUserPreferencesSchema( - enable_discord_notifications=True, enable_email_notifications=True, - discord_webhook="http://example.com/webhook", ) response = update_user_preferences(request, payload) assert response["success"] is True user_preferences = UserPreferences.objects.get(orguser=orguser) - assert user_preferences.enable_discord_notifications is True assert user_preferences.enable_email_notifications is True - assert user_preferences.discord_webhook == "http://example.com/webhook" def test_get_user_preferences_success(orguser, user_preferences): From 0a5080f990359e49550880c94cd2551089d18d93 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Sun, 17 Nov 2024 01:51:34 +0530 Subject: [PATCH 37/57] did changes to the createorgplan.py script --- ddpui/management/commands/createorgplan.py | 42 +++++++++++++++------- 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/ddpui/management/commands/createorgplan.py b/ddpui/management/commands/createorgplan.py index 9d15a0f3..97cd479c 100644 --- a/ddpui/management/commands/createorgplan.py +++ b/ddpui/management/commands/createorgplan.py @@ -15,6 +15,11 @@ class Command(BaseCommand): def add_arguments(self, parser): parser.add_argument("--org", type=str, help="Org slug", required=True) parser.add_argument("--with-superset", action="store_true", help="Include superset") + parser.add_argument( + "--is-free-trial", + action="store_true", + help="Set the plan as Free Trial", + ) parser.add_argument( "--duration", choices=["Monthly", "Annual"], @@ -22,11 +27,11 @@ def add_arguments(self, parser): required=True, ) parser.add_argument("--start-date", type=str, help="Start date", required=False) - parser.add_argument("--end-date", type=str, help="Start date", required=False) + parser.add_argument("--end-date", type=str, help="End date", required=False) parser.add_argument("--overwrite", action="store_true", help="Overwrite existing plan") def handle(self, *args, **options): - """create the OrgPlan for the Org""" + """Create the OrgPlan for the Org""" org = Org.objects.filter(slug=options["org"]).first() if org is None: self.stdout.write(self.style.ERROR(f"Org {options['org']} not found")) @@ -40,7 +45,6 @@ def handle(self, *args, **options): if not org_plan: org_plan = OrgPlans(org=org) - org_plan.superset_included = options["with_superset"] org_plan.subscription_duration = options["duration"] org_plan.start_date = ( @@ -50,16 +54,9 @@ def handle(self, *args, **options): datetime.strptime(options["end_date"], "%Y-%m-%d") if options["end_date"] else None ) - if options["with_superset"]: - org_plan.base_plan = "DALGO + Superset" - org_plan.can_upgrade_plan = False - org_plan.features = { - "pipeline": ["Ingest", "Transform", "Orchestrate"], - "aiFeatures": ["AI data analysis"], - "dataQuality": ["Data quality dashboards"], - } - else: - org_plan.base_plan = "DALGO" + if options["is_free_trial"]: + org_plan.base_plan = "Free trial" + org_plan.superset_included = True org_plan.can_upgrade_plan = True org_plan.features = { "pipeline": ["Ingest", "Transform", "Orchestrate"], @@ -67,6 +64,25 @@ def handle(self, *args, **options): "dataQuality": ["Data quality dashboards"], "superset": ["Superset dashboards", "Superset Usage dashboards"], } + else: + org_plan.base_plan = "DALGO" + if options["with_superset"]: + org_plan.superset_included = True + org_plan.can_upgrade_plan = False + org_plan.features = { + "pipeline": ["Ingest", "Transform", "Orchestrate"], + "aiFeatures": ["AI data analysis"], + "dataQuality": ["Data quality dashboards"], + "superset": ["Superset dashboards", "Superset Usage dashboards"], + } + else: + org_plan.superset_included = False + org_plan.can_upgrade_plan = True + org_plan.features = { + "pipeline": ["Ingest", "Transform", "Orchestrate"], + "aiFeatures": ["AI data analysis"], + "dataQuality": ["Data quality dashboards"], + } org_plan.save() From 2e5ab01a63c534dc6655f6276313d19970d23d55 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sun, 17 Nov 2024 13:14:44 +0530 Subject: [PATCH 38/57] renamed llm_optin to disclaimer_shown for the UserPreferences table --- .../commands/userpermissions2orgpermissions.py | 5 ++++- ...lm_optin_userpreferences_disclaimer_shown.py | 17 +++++++++++++++++ ddpui/models/userpreferences.py | 4 ++-- 3 files changed, 23 insertions(+), 3 deletions(-) create mode 100644 ddpui/migrations/0112_rename_llm_optin_userpreferences_disclaimer_shown.py diff --git a/ddpui/management/commands/userpermissions2orgpermissions.py b/ddpui/management/commands/userpermissions2orgpermissions.py index fefdd682..6bbea0e0 100644 --- a/ddpui/management/commands/userpermissions2orgpermissions.py +++ b/ddpui/management/commands/userpermissions2orgpermissions.py @@ -38,7 +38,10 @@ def handle(self, *args, **options): if userpreferences is not None: print("Found user preferences for " + orguser.user.email) - if orgpreferences.llm_optin is False and userpreferences.llm_optin is True: + if ( + orgpreferences.llm_optin is False + and userpreferences.disclaimer_shown is True + ): print("Approving LLM opt-in by " + orguser.user.email) orgpreferences.llm_optin = True orgpreferences.llm_optin_approved_by = orguser diff --git a/ddpui/migrations/0112_rename_llm_optin_userpreferences_disclaimer_shown.py b/ddpui/migrations/0112_rename_llm_optin_userpreferences_disclaimer_shown.py new file mode 100644 index 00000000..431c4cc4 --- /dev/null +++ b/ddpui/migrations/0112_rename_llm_optin_userpreferences_disclaimer_shown.py @@ -0,0 +1,17 @@ +# Generated by Django 4.2 on 2024-11-17 07:42 + +from django.db import migrations + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0111_remove_orgpreferences_trial_end_date_and_more"), + ] + + operations = [ + migrations.RenameField( + model_name="userpreferences", + old_name="llm_optin", + new_name="disclaimer_shown", + ), + ] diff --git a/ddpui/models/userpreferences.py b/ddpui/models/userpreferences.py index 8769a31b..0c36c91a 100644 --- a/ddpui/models/userpreferences.py +++ b/ddpui/models/userpreferences.py @@ -10,7 +10,7 @@ class UserPreferences(models.Model): enable_discord_notifications = models.BooleanField(default=False) # deprecated discord_webhook = models.URLField(blank=True, null=True) # deprecated enable_email_notifications = models.BooleanField(default=False) - llm_optin = models.BooleanField(default=False) # deprecated + disclaimer_shown = models.BooleanField(default=False) created_at = models.DateTimeField(default=timezone.now) updated_at = models.DateTimeField(default=timezone.now) @@ -18,5 +18,5 @@ def to_json(self) -> dict: """Return a dict representation of the model""" return { "enable_email_notifications": self.enable_email_notifications, - "llm_optin": self.llm_optin, + "disclaimer_shown": self.disclaimer_shown, } From 3a2c8884ccafda9b1cd0b6a9f996b78964679c98 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sun, 17 Nov 2024 13:15:17 +0530 Subject: [PATCH 39/57] removed llm_optin from OrgUser APIs, keeping it in the model for now --- ddpui/api/user_org_api.py | 1 - ddpui/core/orguserfunctions.py | 2 -- ddpui/models/org_user.py | 4 +--- 3 files changed, 1 insertion(+), 6 deletions(-) diff --git a/ddpui/api/user_org_api.py b/ddpui/api/user_org_api.py index f8d4f727..c6e41cd6 100644 --- a/ddpui/api/user_org_api.py +++ b/ddpui/api/user_org_api.py @@ -105,7 +105,6 @@ def get_current_user_v2(request, org_slug: str = None): for rolep in curr_orguser.new_role.rolepermissions.all() ], is_demo=(curr_orguser.org.type == OrgType.DEMO if curr_orguser.org else False), - llm_optin=curr_orguser.llm_optin, is_llm_active=org_preferences.llm_optin, ) ) diff --git a/ddpui/core/orguserfunctions.py b/ddpui/core/orguserfunctions.py index 9d151643..b83adf41 100644 --- a/ddpui/core/orguserfunctions.py +++ b/ddpui/core/orguserfunctions.py @@ -154,8 +154,6 @@ def update_orguser_v1(orguser: OrgUser, payload: OrgUserUpdatev1): orguser.user.is_active = payload.active if payload.role_uuid: orguser.new_role = Role.objects.filter(uuid=payload.role_uuid).first() - if payload.llm_optin is not None: - orguser.llm_optin = payload.llm_optin orguser.user.save() orguser.save() diff --git a/ddpui/models/org_user.py b/ddpui/models/org_user.py index 8f8bbe07..9fc97f58 100644 --- a/ddpui/models/org_user.py +++ b/ddpui/models/org_user.py @@ -70,7 +70,7 @@ class OrgUser(models.Model): role = models.IntegerField(choices=OrgUserRole.choices(), default=OrgUserRole.REPORT_VIEWER) new_role = models.ForeignKey(Role, on_delete=models.SET_NULL, null=True) email_verified = models.BooleanField(default=False) - llm_optin = models.BooleanField(default=False) + llm_optin = models.BooleanField(default=False) # deprecated created_at = models.DateTimeField(auto_created=True, default=timezone.now) updated_at = models.DateTimeField(auto_now=True) @@ -103,7 +103,6 @@ class OrgUserUpdatev1(Schema): role_uuid: uuid.UUID = None email: str = None active: bool = None - llm_optin: bool = False class OrgUserUpdateNewRole(Schema): @@ -131,7 +130,6 @@ class OrgUserResponse(Schema): is_demo: bool = False new_role_slug: str | None permissions: list[dict] - llm_optin: bool = None is_llm_active: bool = None From efdfe205b1728419354724e6c9aca35037ad3915 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sun, 17 Nov 2024 14:24:48 +0530 Subject: [PATCH 40/57] put "disclaimer_shown" and "is_llm_active" back into the GET userpreferences response --- ddpui/api/user_preferences_api.py | 3 +++ ddpui/tests/api_tests/test_user_preferences_api.py | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index a78a10ae..0a80a091 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -51,9 +51,12 @@ def get_user_preferences(request): """gets user preferences for the user""" orguser: OrgUser = request.orguser user_preferences, created = UserPreferences.objects.get_or_create(orguser=orguser) + org_preferences, created = OrgPreferences.objects.get_or_create(org=orguser.org) res = { "enable_email_notifications": user_preferences.enable_email_notifications, + "disclaimer_shown": user_preferences.disclaimer_shown, + "is_llm_active": org_preferences.llm_optin, } return {"success": True, "res": res} diff --git a/ddpui/tests/api_tests/test_user_preferences_api.py b/ddpui/tests/api_tests/test_user_preferences_api.py index 8f5ef7b3..5f343397 100644 --- a/ddpui/tests/api_tests/test_user_preferences_api.py +++ b/ddpui/tests/api_tests/test_user_preferences_api.py @@ -140,6 +140,8 @@ def test_get_user_preferences_success(orguser, user_preferences): assert response["success"] is True assert response["res"] == { "enable_email_notifications": user_preferences.enable_email_notifications, + "disclaimer_shown": user_preferences.disclaimer_shown, + "is_llm_active": False, } @@ -153,5 +155,7 @@ def test_get_user_preferences_success_if_not_exist(orguser): assert response["success"] is True assert response["res"] == { "enable_email_notifications": False, + "disclaimer_shown": False, + "is_llm_active": False, } assert UserPreferences.objects.filter(orguser=orguser).exists() From bb7185211e24b66a6b2b0a9cfb0a1fad14237320 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Sun, 17 Nov 2024 14:45:41 +0530 Subject: [PATCH 41/57] added field disclaimer_shown to the userpreferences_schema model and the api file --- ddpui/api/user_preferences_api.py | 3 +++ ddpui/schemas/userpreferences_schema.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index 0a80a091..67f92cda 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -27,6 +27,7 @@ def create_user_preferences(request, payload: CreateUserPreferencesSchema): user_preferences = UserPreferences.objects.create( orguser=orguser, enable_email_notifications=payload.enable_email_notifications, + disclaimer_shown=payload.disclaimer_shown, ) return {"success": True, "res": user_preferences.to_json()} @@ -41,6 +42,8 @@ def update_user_preferences(request, payload: UpdateUserPreferencesSchema): if payload.enable_email_notifications is not None: user_preferences.enable_email_notifications = payload.enable_email_notifications + if payload.disclaimer_shown is not None: + user_preferences.disclaimer_shown = payload.disclaimer_shown user_preferences.save() return {"success": True, "res": user_preferences.to_json()} diff --git a/ddpui/schemas/userpreferences_schema.py b/ddpui/schemas/userpreferences_schema.py index e16e1593..7651686f 100644 --- a/ddpui/schemas/userpreferences_schema.py +++ b/ddpui/schemas/userpreferences_schema.py @@ -6,9 +6,11 @@ class CreateUserPreferencesSchema(Schema): """Schema for creating user preferences for the user.""" enable_email_notifications: bool + disclaimer_shown: Optional[bool] = None class UpdateUserPreferencesSchema(Schema): """Schema for updating user preferences for the user.""" enable_email_notifications: Optional[bool] = None + disclaimer_shown: Optional[bool] = None From dae00dbe344e5af9561ee4b790f2b5a0843b1e32 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Sun, 17 Nov 2024 14:59:58 +0530 Subject: [PATCH 42/57] fixed a test case --- ddpui/tests/api_tests/test_user_preferences_api.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ddpui/tests/api_tests/test_user_preferences_api.py b/ddpui/tests/api_tests/test_user_preferences_api.py index 5f343397..94b14110 100644 --- a/ddpui/tests/api_tests/test_user_preferences_api.py +++ b/ddpui/tests/api_tests/test_user_preferences_api.py @@ -76,9 +76,7 @@ def test_seed_data(seed_db): def test_create_user_preferences_success(orguser): """tests the success of creating user preferences for the OrgUser""" request = mock_request(orguser) - payload = CreateUserPreferencesSchema( - enable_email_notifications=True, - ) + payload = CreateUserPreferencesSchema(enable_email_notifications=True, disclaimer_shown=True) response = create_user_preferences(request, payload) From 27dd66c0c49001d2f3435bf922323f940afa6cfb Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Sun, 17 Nov 2024 16:20:54 +0530 Subject: [PATCH 43/57] make acc manager's disclaimer_shown true when he/she enables the llm_optin --- ddpui/api/org_preferences_api.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index d586e173..58101b0e 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -6,12 +6,14 @@ from ddpui.models.org_preferences import OrgPreferences from ddpui.models.org_supersets import OrgSupersets from ddpui.models.org_plans import OrgPlans +from ddpui.models.userpreferences import UserPreferences from ddpui.schemas.org_preferences_schema import ( CreateOrgPreferencesSchema, UpdateLLMOptinSchema, UpdateDiscordNotificationsSchema, CreateOrgSupersetDetailsSchema, ) +from django.db import transaction from ddpui.auth import has_permission from ddpui.models.org_user import OrgUser from ddpui.ddpdbt import dbt_service @@ -44,6 +46,7 @@ def create_org_preferences(request, payload: CreateOrgPreferencesSchema): @orgpreference_router.put("/llm_approval", auth=auth.CustomAuthMiddleware()) @has_permission(["can_edit_llm_settings"]) +@transaction.atomic def update_org_preferences(request, payload: UpdateLLMOptinSchema): """Updates llm preferences for the logged-in user's organization""" @@ -51,6 +54,8 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): org = orguser.org org_preferences = OrgPreferences.objects.filter(org=org).first() + user_preferences = UserPreferences.objects.filter(orguser=orguser).first() + if org_preferences is None: org_preferences = OrgPreferences.objects.create(org=org) @@ -58,11 +63,12 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): org_preferences.llm_optin = True org_preferences.llm_optin_approved_by = orguser org_preferences.llm_optin_date = timezone.now() + user_preferences.disclaimer_shown = True else: org_preferences.llm_optin = False org_preferences.llm_optin_approved_by = None org_preferences.llm_optin_date = None - + user_preferences.save() org_preferences.save() return {"success": True, "res": org_preferences.to_json()} From e6083d28916a21a654c0db834521ba82590f9ad6 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Mon, 18 Nov 2024 11:21:54 +0530 Subject: [PATCH 44/57] minor change; first_name is usually null --- ddpui/api/user_preferences_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index 67f92cda..4d1cde23 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -82,7 +82,7 @@ def post_request_llm_analysis_feature_enabled(request): # send notification to all account managers notification_data = NotificationDataSchema( author=orguser.user.email, - message=f"{orguser.user.first_name} is requesting to enable LLM analysis feature", + message=f"{orguser.user.email} is requesting to enable LLM analysis feature", urgent=False, scheduled_time=None, recipients=[acc_manager.id for acc_manager in acc_managers], From afc2b6604916018d2bee3d6afaa8a443146e15ff Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Tue, 19 Nov 2024 10:03:13 +0530 Subject: [PATCH 45/57] minor change --- ddpui/api/org_preferences_api.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 58101b0e..0245149c 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -59,6 +59,9 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): if org_preferences is None: org_preferences = OrgPreferences.objects.create(org=org) + if user_preferences is None: + user_preferences = UserPreferences.objects.create(orguser=orguser) + if payload.llm_optin is True: org_preferences.llm_optin = True org_preferences.llm_optin_approved_by = orguser From bd9856a919b248386afbf274d8bb20fe36f1ed18 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 20 Nov 2024 13:20:20 +0530 Subject: [PATCH 46/57] added new fields and reqeust llmfeature done --- ddpui/api/user_preferences_api.py | 14 +++++++- ...113_orgplans_upgrade_requested_and_more.py | 34 +++++++++++++++++++ ddpui/models/org_plans.py | 1 + ddpui/models/org_preferences.py | 4 +++ 4 files changed, 52 insertions(+), 1 deletion(-) create mode 100644 ddpui/migrations/0113_orgplans_upgrade_requested_and_more.py diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index 4d1cde23..7bb5afce 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -60,6 +60,7 @@ def get_user_preferences(request): "enable_email_notifications": user_preferences.enable_email_notifications, "disclaimer_shown": user_preferences.disclaimer_shown, "is_llm_active": org_preferences.llm_optin, + "enable_llm_requested": org_preferences.enable_llm_request, } return {"success": True, "res": res} @@ -88,6 +89,17 @@ def post_request_llm_analysis_feature_enabled(request): recipients=[acc_manager.id for acc_manager in acc_managers], ) - create_notification(notification_data) + res, errors = create_notification(notification_data) + + if errors: + return HttpError(400, "Issue with creating the request notification") + + rows_updated = OrgPreferences.objects.filter(org=org).update( + enable_llm_request=True, enable_llm_requested_by=orguser + ) + if rows_updated == 0: + raise HttpError( + 400, "No rows were updated. OrgPreferences may not exist for this organization." + ) return {"success": True, "res": "Notified account manager(s) to enable LLM analysis feature"} diff --git a/ddpui/migrations/0113_orgplans_upgrade_requested_and_more.py b/ddpui/migrations/0113_orgplans_upgrade_requested_and_more.py new file mode 100644 index 00000000..ad8265bf --- /dev/null +++ b/ddpui/migrations/0113_orgplans_upgrade_requested_and_more.py @@ -0,0 +1,34 @@ +# Generated by Django 4.2 on 2024-11-20 07:39 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + dependencies = [ + ("ddpui", "0112_rename_llm_optin_userpreferences_disclaimer_shown"), + ] + + operations = [ + migrations.AddField( + model_name="orgplans", + name="upgrade_requested", + field=models.BooleanField(default=False), + ), + migrations.AddField( + model_name="orgpreferences", + name="enable_llm_request", + field=models.BooleanField(default=False), + ), + migrations.AddField( + model_name="orgpreferences", + name="enable_llm_requested_by", + field=models.ForeignKey( + blank=True, + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="llm_request", + to="ddpui.orguser", + ), + ), + ] diff --git a/ddpui/models/org_plans.py b/ddpui/models/org_plans.py index b7220fb6..7c4bfc0d 100644 --- a/ddpui/models/org_plans.py +++ b/ddpui/models/org_plans.py @@ -20,6 +20,7 @@ class OrgPlans(models.Model): start_date = models.DateTimeField(null=True, blank=True, default=None) end_date = models.DateTimeField(null=True, blank=True, default=None) can_upgrade_plan = models.BooleanField(default=False) + upgrade_requested = models.BooleanField(default=False) created_at = models.DateTimeField(default=timezone.now) updated_at = models.DateTimeField(default=timezone.now) diff --git a/ddpui/models/org_preferences.py b/ddpui/models/org_preferences.py index 8c6af246..8d896e85 100644 --- a/ddpui/models/org_preferences.py +++ b/ddpui/models/org_preferences.py @@ -13,6 +13,10 @@ class OrgPreferences(models.Model): OrgUser, on_delete=models.CASCADE, related_name="approvedby", null=True, blank=True ) llm_optin_date = models.DateTimeField(null=True, blank=True) + enable_llm_request = models.BooleanField(default=False) + enable_llm_requested_by = models.ForeignKey( + OrgUser, on_delete=models.CASCADE, related_name="llm_request", null=True, blank=True + ) enable_discord_notifications = models.BooleanField(default=False) discord_webhook = models.URLField(blank=True, null=True) created_at = models.DateTimeField(default=timezone.now) From 359843ac42615ae762f49afa705e386dc456c754 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 20 Nov 2024 13:50:50 +0530 Subject: [PATCH 47/57] updated upgradePlan button ogic --- ddpui/api/org_preferences_api.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 0245149c..2d39e152 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -190,4 +190,7 @@ def initiate_upgrade_dalgo_plan(request): for email in biz_dev_emails: send_text_message(email, subject, message) + org_plan.upgrade_requested = True + org_plan.save() + return {"success": True} From d6621294a3bf7a948547c7a54f6c80e985742975 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 20 Nov 2024 14:30:50 +0530 Subject: [PATCH 48/57] upgrade_reqeuested to orgplans --- ddpui/models/org_plans.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ddpui/models/org_plans.py b/ddpui/models/org_plans.py index 7c4bfc0d..9224e320 100644 --- a/ddpui/models/org_plans.py +++ b/ddpui/models/org_plans.py @@ -39,4 +39,5 @@ def to_json(self) -> dict: "start_date": self.start_date, "end_date": self.end_date, "can_upgrade_plan": self.can_upgrade_plan, + "upgrade_requested": self.upgrade_requested, } From 180bb24957fe53709d0d98b9917bb6820c6a21d1 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 20 Nov 2024 17:10:56 +0530 Subject: [PATCH 49/57] added notiifcation, getting error --- ddpui/api/org_preferences_api.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 2d39e152..55b04b9e 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -13,6 +13,11 @@ UpdateDiscordNotificationsSchema, CreateOrgSupersetDetailsSchema, ) +from ddpui.schemas.notifications_api_schemas import SentToEnum, CreateNotificationPayloadSchema +from ddpui.api.notifications_api import post_create_notification +from datetime import timedelta +from django.utils.timezone import now + from django.db import transaction from ddpui.auth import has_permission from ddpui.models.org_user import OrgUser @@ -67,6 +72,8 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): org_preferences.llm_optin_approved_by = orguser org_preferences.llm_optin_date = timezone.now() user_preferences.disclaimer_shown = True + org_preferences.enable_llm_request = False + org_preferences.enable_llm_requested_by = None else: org_preferences.llm_optin = False org_preferences.llm_optin_approved_by = None @@ -74,6 +81,20 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): user_preferences.save() org_preferences.save() + # sending notification to all users in the org. + if payload.llm_optin is True: + notification_payload = CreateNotificationPayloadSchema( + sent_to=SentToEnum.ALL_ORG_USERS, + org_slug=org.slug, + author=orguser.user.email, + message="The AI LLM Data Analysis feature is now enabled.", + scheduled_time=now() + timedelta(minutes=1), + ) + + error, result = post_create_notification(request, notification_payload) + if error: + return HttpError(400, error) + return {"success": True, "res": org_preferences.to_json()} From 991369d2eeffacf1229d2ac074eb4145a0ebf2b1 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 20 Nov 2024 21:23:35 +0530 Subject: [PATCH 50/57] fixd the create notification if the feature is enabled --- ddpui/api/org_preferences_api.py | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 55b04b9e..d26122c9 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -13,11 +13,8 @@ UpdateDiscordNotificationsSchema, CreateOrgSupersetDetailsSchema, ) -from ddpui.schemas.notifications_api_schemas import SentToEnum, CreateNotificationPayloadSchema -from ddpui.api.notifications_api import post_create_notification -from datetime import timedelta -from django.utils.timezone import now - +from ddpui.core.notifications_service import create_notification +from ddpui.schemas.notifications_api_schemas import NotificationDataSchema from django.db import transaction from ddpui.auth import has_permission from ddpui.models.org_user import OrgUser @@ -83,17 +80,19 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): # sending notification to all users in the org. if payload.llm_optin is True: - notification_payload = CreateNotificationPayloadSchema( - sent_to=SentToEnum.ALL_ORG_USERS, - org_slug=org.slug, + recipients: list[OrgUser] = OrgUser.objects.filter(org=org).all() + print(recipients, "recipients") + notification_payload = NotificationDataSchema( author=orguser.user.email, message="The AI LLM Data Analysis feature is now enabled.", - scheduled_time=now() + timedelta(minutes=1), + urgent=False, + scheduled_time=None, + recipients=[recipient.id for recipient in recipients], ) - - error, result = post_create_notification(request, notification_payload) - if error: - return HttpError(400, error) + print(notification_payload, "notfpaylod") + res, errors = create_notification(notification_payload) + if errors: + return HttpError(400, "Issue with creating the request notification") return {"success": True, "res": org_preferences.to_json()} From f257baf71e298a0d537489df8aa9a1853f2a4207 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Wed, 20 Nov 2024 21:24:20 +0530 Subject: [PATCH 51/57] removed print --- ddpui/api/org_preferences_api.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index d26122c9..334b332c 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -81,7 +81,7 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): # sending notification to all users in the org. if payload.llm_optin is True: recipients: list[OrgUser] = OrgUser.objects.filter(org=org).all() - print(recipients, "recipients") + notification_payload = NotificationDataSchema( author=orguser.user.email, message="The AI LLM Data Analysis feature is now enabled.", @@ -89,7 +89,7 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): scheduled_time=None, recipients=[recipient.id for recipient in recipients], ) - print(notification_payload, "notfpaylod") + res, errors = create_notification(notification_payload) if errors: return HttpError(400, "Issue with creating the request notification") From 688ce13dc38c657f83570246764bb536e9985e7b Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Thu, 21 Nov 2024 10:36:23 +0530 Subject: [PATCH 52/57] account managers and admin should be able to click the upgrade button --- seed/003_role_permissions.json | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/seed/003_role_permissions.json b/seed/003_role_permissions.json index 65cee4e5..c477970b 100644 --- a/seed/003_role_permissions.json +++ b/seed/003_role_permissions.json @@ -1810,6 +1810,22 @@ { "model": "ddpui.RolePermission", "pk": 223, + "fields": { + "role": 1, + "permission": 62 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 224, + "fields": { + "role": 2, + "permission": 62 + } + }, + { + "model": "ddpui.RolePermission", + "pk": 225, "fields": { "role": 3, "permission": 63 @@ -1817,7 +1833,7 @@ }, { "model": "ddpui.RolePermission", - "pk": 224, + "pk": 226, "fields": { "role": 4, "permission": 63 @@ -1825,7 +1841,7 @@ }, { "model": "ddpui.RolePermission", - "pk": 225, + "pk": 227, "fields": { "role": 5, "permission": 63 From a3c9ef9a8c17c98a1ef1a09b64cf1e1819eb248f Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Thu, 21 Nov 2024 10:39:59 +0530 Subject: [PATCH 53/57] minor change --- ddpui/api/org_preferences_api.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 334b332c..90f922fc 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -200,6 +200,10 @@ def initiate_upgrade_dalgo_plan(request): if not org_plan: raise HttpError(400, "Org's Plan not found") + # trigger emails only once + if org_plan.upgrade_requested: + return {"success": True, "res": "Upgrade request already sent"} + biz_dev_emails = os.getenv("BIZ_DEV_EMAILS", []).split(",") message = "Upgrade plan request from org: {org_name} with plan: {plan_name}".format( From d36cc722a18fc8d40d82b9500a7b5fe891ee16f0 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Thu, 21 Nov 2024 10:55:33 +0530 Subject: [PATCH 54/57] minor fix --- ddpui/api/user_preferences_api.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index 7bb5afce..8c59d4f8 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -89,10 +89,10 @@ def post_request_llm_analysis_feature_enabled(request): recipients=[acc_manager.id for acc_manager in acc_managers], ) - res, errors = create_notification(notification_data) + error, res = create_notification(notification_data) - if errors: - return HttpError(400, "Issue with creating the request notification") + if "errors" in res and len(res["errors"]) > 0: + raise HttpError(400, "Issue with creating the request notification") rows_updated = OrgPreferences.objects.filter(org=org).update( enable_llm_request=True, enable_llm_requested_by=orguser From bcb41f11c36a45d56418b0ccc40aa3b5fa3e5ee5 Mon Sep 17 00:00:00 2001 From: Ishankoradia Date: Thu, 21 Nov 2024 11:00:41 +0530 Subject: [PATCH 55/57] minor fix --- ddpui/api/org_preferences_api.py | 6 +++--- ddpui/api/user_preferences_api.py | 3 +-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/ddpui/api/org_preferences_api.py b/ddpui/api/org_preferences_api.py index 90f922fc..d482ff7f 100644 --- a/ddpui/api/org_preferences_api.py +++ b/ddpui/api/org_preferences_api.py @@ -90,9 +90,9 @@ def update_org_preferences(request, payload: UpdateLLMOptinSchema): recipients=[recipient.id for recipient in recipients], ) - res, errors = create_notification(notification_payload) - if errors: - return HttpError(400, "Issue with creating the request notification") + error, res = create_notification(notification_payload) + if res and "errors" in res and len(res["errors"]) > 0: + raise HttpError(400, "Issue with creating the request notification") return {"success": True, "res": org_preferences.to_json()} diff --git a/ddpui/api/user_preferences_api.py b/ddpui/api/user_preferences_api.py index 8c59d4f8..3c7da63b 100644 --- a/ddpui/api/user_preferences_api.py +++ b/ddpui/api/user_preferences_api.py @@ -90,8 +90,7 @@ def post_request_llm_analysis_feature_enabled(request): ) error, res = create_notification(notification_data) - - if "errors" in res and len(res["errors"]) > 0: + if res and "errors" in res and len(res["errors"]) > 0: raise HttpError(400, "Issue with creating the request notification") rows_updated = OrgPreferences.objects.filter(org=org).update( From 47b763726d1ec69fcfc7ac85775f8fe343034e73 Mon Sep 17 00:00:00 2001 From: himanshudube97 Date: Thu, 21 Nov 2024 12:13:50 +0530 Subject: [PATCH 56/57] fixed the test case --- ddpui/tests/api_tests/test_user_preferences_api.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ddpui/tests/api_tests/test_user_preferences_api.py b/ddpui/tests/api_tests/test_user_preferences_api.py index 94b14110..f2332a35 100644 --- a/ddpui/tests/api_tests/test_user_preferences_api.py +++ b/ddpui/tests/api_tests/test_user_preferences_api.py @@ -140,6 +140,7 @@ def test_get_user_preferences_success(orguser, user_preferences): "enable_email_notifications": user_preferences.enable_email_notifications, "disclaimer_shown": user_preferences.disclaimer_shown, "is_llm_active": False, + "enable_llm_requested": False, } @@ -155,5 +156,6 @@ def test_get_user_preferences_success_if_not_exist(orguser): "enable_email_notifications": False, "disclaimer_shown": False, "is_llm_active": False, + "enable_llm_requested": False, } assert UserPreferences.objects.filter(orguser=orguser).exists() From c89e426a3f6c7c0faba0110d8b781e810bb89891 Mon Sep 17 00:00:00 2001 From: Rohit Chatterjee Date: Sun, 24 Nov 2024 10:39:19 +0530 Subject: [PATCH 57/57] --plan flag some cleanup --- ddpui/management/commands/createorgplan.py | 51 +++++++++------------- 1 file changed, 21 insertions(+), 30 deletions(-) diff --git a/ddpui/management/commands/createorgplan.py b/ddpui/management/commands/createorgplan.py index 97cd479c..9069c754 100644 --- a/ddpui/management/commands/createorgplan.py +++ b/ddpui/management/commands/createorgplan.py @@ -16,9 +16,9 @@ def add_arguments(self, parser): parser.add_argument("--org", type=str, help="Org slug", required=True) parser.add_argument("--with-superset", action="store_true", help="Include superset") parser.add_argument( - "--is-free-trial", - action="store_true", - help="Set the plan as Free Trial", + "--plan", + choices=["Free Trial", "DALGO", "Internal"], + default="DALGO", ) parser.add_argument( "--duration", @@ -54,35 +54,26 @@ def handle(self, *args, **options): datetime.strptime(options["end_date"], "%Y-%m-%d") if options["end_date"] else None ) - if options["is_free_trial"]: - org_plan.base_plan = "Free trial" - org_plan.superset_included = True + org_plan.features = { + "pipeline": ["Ingest", "Transform", "Orchestrate"], + "aiFeatures": ["AI data analysis"], + "dataQuality": ["Data quality dashboards"], + } + + org_plan.superset_included = options["with_superset"] + if options["with_superset"]: + org_plan.features["superset"] = ["Superset dashboards", "Superset Usage dashboards"] + + org_plan.base_plan = options["plan"] + + if options["plan"] == "Free Trial": org_plan.can_upgrade_plan = True - org_plan.features = { - "pipeline": ["Ingest", "Transform", "Orchestrate"], - "aiFeatures": ["AI data analysis"], - "dataQuality": ["Data quality dashboards"], - "superset": ["Superset dashboards", "Superset Usage dashboards"], - } + + elif options["plan"] == "Internal": + org_plan.can_upgrade_plan = False + else: - org_plan.base_plan = "DALGO" - if options["with_superset"]: - org_plan.superset_included = True - org_plan.can_upgrade_plan = False - org_plan.features = { - "pipeline": ["Ingest", "Transform", "Orchestrate"], - "aiFeatures": ["AI data analysis"], - "dataQuality": ["Data quality dashboards"], - "superset": ["Superset dashboards", "Superset Usage dashboards"], - } - else: - org_plan.superset_included = False - org_plan.can_upgrade_plan = True - org_plan.features = { - "pipeline": ["Ingest", "Transform", "Orchestrate"], - "aiFeatures": ["AI data analysis"], - "dataQuality": ["Data quality dashboards"], - } + org_plan.can_upgrade_plan = not options["with_superset"] org_plan.save()