add\remove data-source

Author: Roey7

app/api/data_source.py

@@ -5,11 +5,13 @@
 from fastapi import APIRouter
 from pydantic import BaseModel
 from starlette.background import BackgroundTasks
+from starlette.requests import Request
 
 from data_source.api.base_data_source import ConfigField
 from data_source.api.context import DataSourceContext
 from db_engine import Session
 from schemas import DataSourceType, DataSource
+from telemetry import Posthog
 
 router = APIRouter(
     prefix='/data-sources',
@@ -64,15 +66,16 @@ class AddDataSource(BaseModel):
 
 
 @router.delete("/{data_source_id}")
-async def delete_data_source(data_source_id: int):
-    DataSourceContext.delete_data_source(data_source_id=data_source_id)
+async def delete_data_source(request: Request, data_source_id: int):
+    deleted_name = DataSourceContext.delete_data_source(data_source_id=data_source_id)
+    Posthog.removed_data_source(uuid=request.headers.get('uuid'), name=deleted_name)
     return {"success": "Data source deleted successfully"}
 
 
 @router.post("")
-async def add_integration(dto: AddDataSource, background_tasks: BackgroundTasks) -> int:
+async def add_integration(request: Request, dto: AddDataSource, background_tasks: BackgroundTasks) -> int:
     data_source = DataSourceContext.create_data_source(name=dto.name, config=dto.config)
-
+    Posthog.added_data_source(uuid=request.headers.get('uuid'), name=dto.name)
     # in main.py we have a background task that runs every 5 minutes and indexes the data source
     # but here we want to index the data source immediately
     background_tasks.add_task(data_source.index)

app/data_source/api/context.py

@@ -47,17 +47,20 @@ def create_data_source(cls, name: str, config: dict) -> BaseDataSource:
             return data_source
 
     @classmethod
-    def delete_data_source(cls, data_source_id: int):
+    def delete_data_source(cls, data_source_id: int) -> str:
         with Session() as session:
             data_source = session.query(DataSource).filter_by(id=data_source_id).first()
             if data_source is None:
                 raise KnownException(message=f"Data source {data_source_id} does not exist")
 
+            data_source_name = data_source.type.name
             session.delete(data_source)
             session.commit()
 
             del cls._data_sources[data_source_id]
 
+            return data_source_name
+
     @classmethod
     def init(cls):
         cls._add_data_sources_to_db()

app/telemetry.py

@@ -18,6 +18,8 @@ class Posthog:
 
     RUN_EVENT = "run"
     BACKEND_SEARCH_EVENT = "backend_search"
+    BACKEND_ADDED = "backend_added"
+    BACKEND_REMOVED = "backend_removed"
     _should_capture = False
     _identified_uuid: Optional[str] = None
 
@@ -57,15 +59,15 @@ def _identify(cls):
             pass
 
     @classmethod
-    def _capture(cls, event: str, uuid=None):
+    def _capture(cls, event: str, uuid=None, properties=None):
         if cls._identified_uuid is None:
             cls._identify()
 
         if not cls._should_capture:
             return
 
         try:
-            posthog.capture(uuid or cls._identified_uuid, event)
+            posthog.capture(uuid or cls._identified_uuid, event, properties)
         except Exception as e:
             pass
 
@@ -80,3 +82,11 @@ def send_startup_telemetry(cls):
     @classmethod
     def increase_search_count(cls, uuid: str):
         cls._capture(cls.BACKEND_SEARCH_EVENT, uuid=uuid)
+
+    @classmethod
+    def added_data_source(cls, uuid: str, name: str):
+        cls._capture(cls.BACKEND_ADDED, uuid=uuid, properties={"name": name})
+
+    @classmethod
+    def removed_data_source(cls, uuid: str, name: str):
+        cls._capture(cls.BACKEND_REMOVED, uuid=uuid, properties={"name": name})

ui/src/components/data-source-panel.tsx

@@ -412,7 +412,11 @@ export default class DataSourcePanel extends React.Component<DataSourcePanelProp
          config: config
       }
       this.setState({ isAddingLoading: true });
-      api.post(`/data-sources`, payload).then(response => {
+      api.post(`/data-sources`, payload, {
+         headers: {
+            uuid: localStorage.getItem('uuid')
+         }
+      }).then(response => {
          toast.success("Data source added successfully, indexing...");
 
          let selectedDataSource = this.state.selectedDataSource;
@@ -446,7 +450,11 @@ export default class DataSourcePanel extends React.Component<DataSourcePanelProp
       let connectedDataSource = this.props.connectedDataSources[index];
       this.setState({ removeInProgressIndex: index });
 
-      api.delete(`/data-sources/${connectedDataSource.id}`).then(response => {
+      api.delete(`/data-sources/${connectedDataSource.id}`, {
+         headers: {
+            uuid: localStorage.getItem('uuid')
+         }
+      }).then(response => {
          toast.success(`${this.capitilize(connectedDataSource.name)} removed.`);
          this.setState({ removeInProgressIndex: -1 });
          this.props.onRemoved(connectedDataSource);