From ef6ee1444025da8609e4fe61d407fd769af1934c Mon Sep 17 00:00:00 2001
From: Stan Girard <girard.stanislas@gmail.com>
Date: Fri, 16 Feb 2024 11:32:13 -0800
Subject: [PATCH] feat(integration): improve (#2199)

# Description

Please include a summary of the changes and the related issue. Please
also include relevant motivation and context.

## Checklist before requesting a review

Please delete options that are not relevant.

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented hard-to-understand areas
- [ ] I have ideally added tests that prove my fix is effective or that
my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged

## Screenshots (if appropriate):
---
 backend/modules/brain/dto/inputs.py                | 9 +++++++++
 backend/modules/brain/entity/integration_brain.py  | 4 ++++
 supabase/migrations/20240216192826_integration.sql | 7 +++++++
 3 files changed, 20 insertions(+)
 create mode 100644 supabase/migrations/20240216192826_integration.sql

diff --git a/backend/modules/brain/dto/inputs.py b/backend/modules/brain/dto/inputs.py
index 509654df0..efc7908fc 100644
--- a/backend/modules/brain/dto/inputs.py
+++ b/backend/modules/brain/dto/inputs.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Optional
 from uuid import UUID
 
@@ -14,6 +15,11 @@ from pydantic import BaseModel, Extra
 logger = get_logger(__name__)
 
 
+class IntegrationType(str, Enum):
+    CUSTOM = "custom"
+    SYNC = "sync"
+
+
 class CreateApiBrainDefinition(BaseModel, extra=Extra.forbid):
     method: ApiBrainAllowedMethods
     url: str
@@ -28,6 +34,9 @@ class CreateIntegrationBrain(BaseModel, extra=Extra.forbid):
     integration_name: str
     integration_logo_url: str
     connection_settings: dict
+    integration_type: IntegrationType
+    description: str
+    max_files: int
 
 
 class BrainIntegrationSettings(BaseModel, extra=Extra.forbid):
diff --git a/backend/modules/brain/entity/integration_brain.py b/backend/modules/brain/entity/integration_brain.py
index b180150b0..5b2d46f7b 100644
--- a/backend/modules/brain/entity/integration_brain.py
+++ b/backend/modules/brain/entity/integration_brain.py
@@ -1,6 +1,7 @@
 from typing import Optional
 from uuid import UUID
 
+from modules.brain.dto.inputs import IntegrationType
 from pydantic import BaseModel
 
 
@@ -9,6 +10,9 @@ class IntegrationDescriptionEntity(BaseModel):
     integration_name: str
     integration_logo_url: Optional[str] = None
     connection_settings: Optional[dict] = None
+    integration_type: IntegrationType
+    description: str
+    max_files: int
 
 
 class IntegrationEntity(BaseModel):
diff --git a/supabase/migrations/20240216192826_integration.sql b/supabase/migrations/20240216192826_integration.sql
new file mode 100644
index 000000000..3528b1d0f
--- /dev/null
+++ b/supabase/migrations/20240216192826_integration.sql
@@ -0,0 +1,7 @@
+create type "public"."integration_type" as enum ('custom', 'sync');
+
+alter table "public"."integrations" add column "description" text not null default 'Default description'::text;
+
+alter table "public"."integrations" add column "integration_type" integration_type not null default 'custom'::integration_type;
+
+alter table "public"."integrations" add column "max_files" integer not null default 0;
\ No newline at end of file