From 5e969766afacd93c812206b94fcbc5df26bfc277 Mon Sep 17 00:00:00 2001
From: Tal Borenstein <tal@keephq.dev>
Date: Tue, 22 Oct 2024 11:23:51 +0300
Subject: [PATCH] feat: provider pulling enable/disable from UI

---
 keep-ui/app/providers/provider-form.tsx       | 107 +++++++++++++++---
 keep-ui/app/providers/providers.tsx           |   3 +
 keep-ui/app/topology/ui/map/service-node.tsx  |   2 +-
 .../workflow-execution-table.tsx              |   2 -
 .../versions/2024-10-22-10-38_8438f041ee0e.py |  32 ++++++
 keep/api/models/db/provider.py                |   1 +
 keep/api/models/provider.py                   |   1 +
 keep/api/routes/preset.py                     |  22 ++--
 keep/api/routes/providers.py                  |   2 +
 keep/providers/providers_factory.py           |   1 +
 keep/providers/providers_service.py           |   5 +
 11 files changed, 151 insertions(+), 27 deletions(-)
 create mode 100644 keep/api/models/db/migrations/versions/2024-10-22-10-38_8438f041ee0e.py

diff --git a/keep-ui/app/providers/provider-form.tsx b/keep-ui/app/providers/provider-form.tsx
index 2db75e4e7..f907e6c08 100644
--- a/keep-ui/app/providers/provider-form.tsx
+++ b/keep-ui/app/providers/provider-form.tsx
@@ -51,6 +51,8 @@ import cookieCutter from "@boiseitguru/cookie-cutter";
 import { useSearchParams } from "next/navigation";
 import "./provider-form.css";
 import { useProviders } from "@/utils/hooks/useProviders";
+import TimeAgo from "react-timeago";
+import { toast } from "react-toastify";
 
 type ProviderFormProps = {
   provider: Provider;
@@ -156,6 +158,9 @@ const ProviderForm = ({
   };
   if (provider.can_setup_webhook) {
     initialData["install_webhook"] = provider.can_setup_webhook;
+    if (provider.pulling_enabled) {
+      initialData["pulling_enabled"] = true;
+    }
   }
   const [formValues, setFormValues] = useState<{
     [key: string]: string | boolean;
@@ -178,6 +183,9 @@ const ProviderForm = ({
 
     if (provider.can_setup_webhook) {
       initialValues["install_webhook"] = provider.can_setup_webhook;
+      if (provider.pulling_enabled) {
+        initialValues["pulling_enabled"] = provider.pulling_enabled;
+      }
     }
     return initialValues;
   });
@@ -342,6 +350,16 @@ const ProviderForm = ({
     }));
   };
 
+  const handlePullingEnabledChange = (
+    event: React.ChangeEvent<HTMLInputElement>
+  ) => {
+    const checked = event.target.checked;
+    setFormValues((prevValues) => ({
+      ...prevValues,
+      pulling_enabled: checked,
+    }));
+  };
+
   const validate = () => {
     const errors = validateForm(formValues);
     if (Object.keys(errors).length === 0) {
@@ -424,9 +442,13 @@ const ProviderForm = ({
       submit(`${getApiURL()}/providers/${provider.id}`, "PUT")
         .then((data) => {
           setIsLoading(false);
+          toast.success("Updated provider successfully", {
+            position: "top-left",
+          });
           mutate();
         })
         .catch((error) => {
+          toast.error("Failed to update provider", { position: "top-left" });
           const updatedFormErrors = error.toString();
           setFormErrors(updatedFormErrors);
           onFormChange(formValues, updatedFormErrors);
@@ -741,7 +763,12 @@ const ProviderForm = ({
             />
           </Link>
         </div>
-
+        {installedProvidersMode && provider.last_pull_time && (
+          <Subtitle>
+            Provider last pull time:{" "}
+            <TimeAgo date={provider.last_pull_time + "Z"} />
+          </Subtitle>
+        )}
         {provider.provisioned && (
           <div className="w-full mt-4">
             <Callout
@@ -902,6 +929,30 @@ const ProviderForm = ({
                       tooltip={`Whether to install Keep as a webhook integration in ${provider.type}. This allows Keep to asynchronously receive alerts from ${provider.type}. Please note that this will install a new integration in ${provider.type} and slightly modify your monitors/notification policy to include Keep.`}
                     />
                   </label>
+                  {
+                    // This is here because pulling is only enabled for providers we can get alerts from (e.g., support webhook)
+                  }
+                  <input
+                    type="checkbox"
+                    id="pulling_enabled"
+                    name="pulling_enabled"
+                    className="mr-2.5"
+                    onChange={handlePullingEnabledChange}
+                    checked={formValues["pulling_enabled"] || false}
+                  />
+                  <label
+                    htmlFor="pulling_enabled"
+                    className="flex items-center"
+                  >
+                    <Text className="capitalize">Pulling Enabled</Text>
+                    <Icon
+                      icon={QuestionMarkCircleIcon}
+                      variant="simple"
+                      color="gray"
+                      size="sm"
+                      tooltip={`Whether Keep should try to pull alerts automatically from the provider once in a while`}
+                    />
+                  </label>
                 </div>
                 {isLocalhost && (
                   <span className="text-sm">
@@ -929,21 +980,45 @@ const ProviderForm = ({
           </div>
 
           {provider.can_setup_webhook && installedProvidersMode && (
-            <Button
-              icon={GlobeAltIcon}
-              onClick={callInstallWebhook}
-              variant="secondary"
-              color="orange"
-              className="mt-2.5"
-              disabled={!installOrUpdateWebhookEnabled || provider.provisioned}
-              tooltip={
-                !installOrUpdateWebhookEnabled
-                  ? "Fix required webhook scopes and refresh scopes to enable"
-                  : "This uses server saved credentials. If needed, please use the `Update` button first"
-              }
-            >
-              Install/Update Webhook
-            </Button>
+            <>
+              <div className="flex">
+                <input
+                  type="checkbox"
+                  id="pulling_enabled"
+                  name="pulling_enabled"
+                  className="mr-2.5"
+                  onChange={handlePullingEnabledChange}
+                  checked={formValues["pulling_enabled"] || false}
+                />
+                <label htmlFor="pulling_enabled" className="flex items-center">
+                  <Text className="capitalize">Pulling Enabled</Text>
+                  <Icon
+                    icon={QuestionMarkCircleIcon}
+                    variant="simple"
+                    color="gray"
+                    size="sm"
+                    tooltip={`Whether Keep should try to pull alerts automatically from the provider once in a while`}
+                  />
+                </label>
+              </div>
+              <Button
+                icon={GlobeAltIcon}
+                onClick={callInstallWebhook}
+                variant="secondary"
+                color="orange"
+                className="mt-2.5"
+                disabled={
+                  !installOrUpdateWebhookEnabled || provider.provisioned
+                }
+                tooltip={
+                  !installOrUpdateWebhookEnabled
+                    ? "Fix required webhook scopes and refresh scopes to enable"
+                    : "This uses server saved credentials. If needed, please use the `Update` button first"
+                }
+              >
+                Install/Update Webhook
+              </Button>
+            </>
           )}
           {provider.supports_webhook && (
             <ProviderSemiAutomated
diff --git a/keep-ui/app/providers/providers.tsx b/keep-ui/app/providers/providers.tsx
index 403fde0c8..c3bcc9ae9 100644
--- a/keep-ui/app/providers/providers.tsx
+++ b/keep-ui/app/providers/providers.tsx
@@ -95,6 +95,8 @@ export interface Provider {
   validatedScopes: { [scopeName: string]: boolean | string };
   methods?: ProviderMethod[];
   tags: TProviderLabels[];
+  last_pull_time?: Date;
+  pulling_enabled: boolean;
   alertsDistribution?: AlertDistritbuionData[];
   alertExample?: { [key: string]: string };
   provisioned?: boolean;
@@ -115,4 +117,5 @@ export const defaultProvider: Provider = {
   type: "",
   tags: [],
   validatedScopes: {},
+  pulling_enabled: true,
 };
diff --git a/keep-ui/app/topology/ui/map/service-node.tsx b/keep-ui/app/topology/ui/map/service-node.tsx
index 123303945..a15c7e3c7 100644
--- a/keep-ui/app/topology/ui/map/service-node.tsx
+++ b/keep-ui/app/topology/ui/map/service-node.tsx
@@ -107,7 +107,7 @@ export function ServiceNode({ data, selected }: NodeProps<ServiceNodeType>) {
         onMouseEnter={() => setShowDetails(true)}
         onMouseLeave={() => setShowDetails(false)}
       >
-        <strong className="text-lg">{data.display_name ?? data.service}</strong>
+        <strong className="text-lg">{data.display_name || data.service}</strong>
         {alertCount > 0 && (
           <span
             className={`absolute top-[-20px] right-[-20px] mt-2 mr-2 px-2 py-1 text-white text-xs font-bold rounded-full ${badgeColor} hover:cursor-pointer`}
diff --git a/keep-ui/app/workflows/[workflow_id]/workflow-execution-table.tsx b/keep-ui/app/workflows/[workflow_id]/workflow-execution-table.tsx
index f1b9a225e..752daee02 100644
--- a/keep-ui/app/workflows/[workflow_id]/workflow-execution-table.tsx
+++ b/keep-ui/app/workflows/[workflow_id]/workflow-execution-table.tsx
@@ -14,8 +14,6 @@ import Image from "next/image";
 import {
   CheckCircleIcon,
   EllipsisHorizontalIcon,
-  EyeIcon,
-  WrenchIcon,
   XCircleIcon,
 } from "@heroicons/react/20/solid";
 import TimeAgo, { Formatter, Suffix, Unit } from "react-timeago";
diff --git a/keep/api/models/db/migrations/versions/2024-10-22-10-38_8438f041ee0e.py b/keep/api/models/db/migrations/versions/2024-10-22-10-38_8438f041ee0e.py
new file mode 100644
index 000000000..eae0abfc5
--- /dev/null
+++ b/keep/api/models/db/migrations/versions/2024-10-22-10-38_8438f041ee0e.py
@@ -0,0 +1,32 @@
+"""add pulling_enabled
+
+Revision ID: 8438f041ee0e
+Revises: 83c1020be97d
+Create Date: 2024-10-22 10:38:29.857284
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8438f041ee0e"
+down_revision = "83c1020be97d"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("provider", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("pulling_enabled", sa.Boolean(), nullable=False, default=True)
+        )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("provider", schema=None) as batch_op:
+        batch_op.drop_column("pulling_enabled")
+    # ### end Alembic commands ###
diff --git a/keep/api/models/db/provider.py b/keep/api/models/db/provider.py
index 37d8d05fa..33e200784 100644
--- a/keep/api/models/db/provider.py
+++ b/keep/api/models/db/provider.py
@@ -20,6 +20,7 @@ class Provider(SQLModel, table=True):
         sa_column=Column(JSON)
     )  # scope name is key and value is either True if validated or string with error message, e.g: {"read": True, "write": "error message"}
     consumer: bool = False
+    pulling_enabled: bool = True
     last_pull_time: Optional[datetime]
     provisioned: bool = Field(default=False)
 
diff --git a/keep/api/models/provider.py b/keep/api/models/provider.py
index 93810817c..5c15bfc22 100644
--- a/keep/api/models/provider.py
+++ b/keep/api/models/provider.py
@@ -39,6 +39,7 @@ class Provider(BaseModel):
     methods: list[ProviderMethod] = []
     installed_by: str | None = None
     installation_time: datetime | None = None
+    pulling_enabled: bool = True
     last_pull_time: datetime | None = None
     docs: str | None = None
     tags: list[
diff --git a/keep/api/routes/preset.py b/keep/api/routes/preset.py
index a812d6be1..cfa701681 100644
--- a/keep/api/routes/preset.py
+++ b/keep/api/routes/preset.py
@@ -1,16 +1,18 @@
+import json
 import logging
 import os
 import uuid
 from datetime import datetime
 from typing import Optional
+
 from fastapi import (
     APIRouter,
     BackgroundTasks,
     Depends,
     HTTPException,
+    Query,
     Request,
     Response,
-    Query
 )
 from pydantic import BaseModel
 from sqlmodel import Session, select
@@ -24,7 +26,6 @@
     update_provider_last_pull_time,
 )
 from keep.api.models.alert import AlertDto
-from keep.api.models.time_stamp import TimeStampFilter
 from keep.api.models.db.preset import (
     Preset,
     PresetDto,
@@ -33,6 +34,7 @@
     Tag,
     TagDto,
 )
+from keep.api.models.time_stamp import TimeStampFilter
 from keep.api.tasks.process_event_task import process_event
 from keep.api.tasks.process_topology_task import process_topology
 from keep.contextmanager.contextmanager import ContextManager
@@ -41,7 +43,6 @@
 from keep.providers.base.base_provider import BaseTopologyProvider
 from keep.providers.providers_factory import ProvidersFactory
 from keep.searchengine.searchengine import SearchEngine
-import json
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
@@ -86,6 +87,10 @@ def pull_data_from_providers(
             "trace_id": trace_id,
         }
 
+        if not provider.pulling_enabled:
+            logger.debug("Pulling is disabled for this provider", extra=extra)
+            continue
+
         if provider.last_pull_time is not None:
             now = datetime.now()
             days_passed = (now - provider.last_pull_time).days
@@ -173,9 +178,7 @@ def pull_data_from_providers(
 
 
 # Function to handle the time_stamp query parameter and parse it
-def _get_time_stamp_filter(
-    time_stamp: Optional[str] = Query(None)
-) -> TimeStampFilter:
+def _get_time_stamp_filter(time_stamp: Optional[str] = Query(None)) -> TimeStampFilter:
     if time_stamp:
         try:
             # Parse the JSON string
@@ -186,6 +189,7 @@ def _get_time_stamp_filter(
             raise HTTPException(status_code=400, detail="Invalid time_stamp format")
     return TimeStampFilter()
 
+
 @router.get(
     "",
     description="Get all presets for tenant",
@@ -195,7 +199,7 @@ def get_presets(
         IdentityManagerFactory.get_auth_verifier(["read:preset"])
     ),
     session: Session = Depends(get_session),
-    time_stamp: TimeStampFilter = Depends(_get_time_stamp_filter)
+    time_stamp: TimeStampFilter = Depends(_get_time_stamp_filter),
 ) -> list[PresetDto]:
     tenant_id = authenticated_entity.tenant_id
     logger.info(f"Getting all presets {time_stamp}")
@@ -224,7 +228,9 @@ def get_presets(
     # get the number of alerts + noisy alerts for each preset
     search_engine = SearchEngine(tenant_id=tenant_id)
     # get the preset metatada
-    presets_dto = search_engine.search_preset_alerts(presets=presets_dto, time_stamp=time_stamp)
+    presets_dto = search_engine.search_preset_alerts(
+        presets=presets_dto, time_stamp=time_stamp
+    )
 
     return presets_dto
 
diff --git a/keep/api/routes/providers.py b/keep/api/routes/providers.py
index e30f7c349..e746ce839 100644
--- a/keep/api/routes/providers.py
+++ b/keep/api/routes/providers.py
@@ -490,6 +490,7 @@ async def install_provider(
         provider_id = provider_info.pop("provider_id")
         provider_name = provider_info.pop("provider_name")
         provider_type = provider_info.pop("provider_type", None) or provider_id
+        pulling_enabled = provider_info.pop("pulling_enabled", True)
     except KeyError as e:
         raise HTTPException(
             status_code=400, detail=f"Missing required field: {e.args[0]}"
@@ -507,6 +508,7 @@ async def install_provider(
             provider_name,
             provider_type,
             provider_info,
+            pulling_enabled=pulling_enabled,
         )
         return JSONResponse(status_code=200, content=result)
     except HTTPException as e:
diff --git a/keep/providers/providers_factory.py b/keep/providers/providers_factory.py
index 079467200..fb5d9bdd6 100644
--- a/keep/providers/providers_factory.py
+++ b/keep/providers/providers_factory.py
@@ -411,6 +411,7 @@ def get_installed_providers(
             provider_copy.installation_time = p.installation_time
             provider_copy.last_pull_time = p.last_pull_time
             provider_copy.provisioned = p.provisioned
+            provider_copy.pulling_enabled = p.pulling_enabled
             try:
                 provider_auth = {"name": p.name}
                 if include_details:
diff --git a/keep/providers/providers_service.py b/keep/providers/providers_service.py
index dad1d8d37..453a2ebcf 100644
--- a/keep/providers/providers_service.py
+++ b/keep/providers/providers_service.py
@@ -47,6 +47,7 @@ def install_provider(
         provider_config: Dict[str, Any],
         provisioned: bool = False,
         validate_scopes: bool = True,
+        pulling_enabled: bool = True,
     ) -> Dict[str, Any]:
         provider_unique_id = uuid.uuid4().hex
         logger.info(
@@ -95,6 +96,7 @@ def install_provider(
                 validatedScopes=validated_scopes,
                 consumer=provider.is_consumer,
                 provisioned=provisioned,
+                pulling_enabled=pulling_enabled,
             )
             try:
                 session.add(provider_model)
@@ -148,6 +150,8 @@ def update_provider(
         if provider.provisioned:
             raise HTTPException(403, detail="Cannot update a provisioned provider")
 
+        pulling_enabled = provider_info.pop("pulling_enabled", True)
+
         provider_config = {
             "authentication": provider_info,
             "name": provider.name,
@@ -171,6 +175,7 @@ def update_provider(
 
         provider.installed_by = updated_by
         provider.validatedScopes = validated_scopes
+        provider.pulling_enabled = pulling_enabled
         session.commit()
 
         return {