fix: add load and remove demo data

Author: Arta Asadi

services/integration/api/integration-types/api.go

@@ -3,6 +3,8 @@ package integration_types
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/google/uuid"
+	"github.com/jackc/pgtype"
 	"github.com/opengovern/og-util/pkg/config"
 	"github.com/opengovern/og-util/pkg/httpclient"
 	"github.com/opengovern/og-util/pkg/opengovernance-es-sdk"
@@ -85,6 +87,7 @@ func (a *API) Register(e *echo.Group) {
 	plugin.POST("/:id/healthcheck", httpserver.AuthorizeHandler(a.HealthCheck, api.ViewerRole))
 	plugin.GET("/tables", httpserver.AuthorizeHandler(a.GetPluginsTables, api.ViewerRole))
 	plugin.PUT("/:id/demo/load", httpserver.AuthorizeHandler(a.LoadPluginDemoData, api.EditorRole))
+	plugin.PUT("/:id/demo/remove", httpserver.AuthorizeHandler(a.RemovePluginDemoData, api.EditorRole))
 }
 
 // List godoc
@@ -1343,6 +1346,10 @@ func (a *API) LoadPluginDemoData(c echo.Context) error {
 		return echo.NewHTTPError(http.StatusNotFound, "plugin not found")
 	}
 
+	if plugin.DemoDataLoaded {
+		return echo.NewHTTPError(http.StatusConflict, "plugin demo data already loaded")
+	}
+
 	if plugin.DemoDataURL == "" {
 		return echo.NewHTTPError(http.StatusNotFound, "plugin does not contain demo data")
 	}
@@ -1354,7 +1361,111 @@ func (a *API) LoadPluginDemoData(c echo.Context) error {
 		ElasticsearchPass: a.elasticConfig.Password,
 		ElasticsearchAddr: a.elasticConfig.Address,
 	}
-	err = demo_import.LoadDemoData(demoImportConfig, a.logger)
+	integrations, err := demo_import.LoadDemoData(demoImportConfig, a.logger)
+	if err != nil {
+		a.logger.Error("failed to load demo data", zap.Error(err))
+		return echo.NewHTTPError(http.StatusInternalServerError, "failed to load demo data")
+	}
+
+	dummyCredentialID := uuid.New()
+	dummyCredential := models2.Credential{
+		ID:              dummyCredentialID,
+		IntegrationType: plugin.IntegrationType,
+		CredentialType:  "",
+		Secret:          "",
+		Metadata: func() pgtype.JSONB {
+			var jsonb pgtype.JSONB
+			if err := jsonb.Set([]byte("{}")); err != nil {
+				a.logger.Error("failed to convert WidgetProps to JSONB", zap.Error(err))
+			}
+			return jsonb
+		}(),
+		MaskedSecret: func() pgtype.JSONB {
+			var jsonb pgtype.JSONB
+			if err := jsonb.Set([]byte("{}")); err != nil {
+				a.logger.Error("failed to convert WidgetProps to JSONB", zap.Error(err))
+			}
+			return jsonb
+		}(),
+		Description: "dummy credential for demo integrations",
+	}
+
+	err = a.database.CreateCredential(&dummyCredential)
+	if err != nil {
+		a.logger.Error("failed to create credential", zap.Error(err))
+		return echo.NewHTTPError(http.StatusInternalServerError, "failed to create credential")
+	}
+
+	for _, i := range integrations {
+		integrationId, err := uuid.Parse(i.IntegrationID)
+		if err != nil {
+			a.logger.Error("failed to parse integration id", zap.Error(err))
+			return echo.NewHTTPError(http.StatusInternalServerError, "failed to parse integration id")
+		}
+		dbIntegration := models2.Integration{
+			Integration: integration.Integration{
+				IntegrationID:   integrationId,
+				ProviderID:      i.ProviderID,
+				Name:            i.Name,
+				IntegrationType: plugin.IntegrationType,
+				Annotations: func() pgtype.JSONB {
+					var jsonb pgtype.JSONB
+					if err := jsonb.Set(i.Annotations); err != nil {
+						a.logger.Error("failed to convert WidgetProps to JSONB", zap.Error(err))
+					}
+					return jsonb
+				}(),
+				Labels: func() pgtype.JSONB {
+					var jsonb pgtype.JSONB
+					if err := jsonb.Set(i.Labels); err != nil {
+						a.logger.Error("failed to convert WidgetProps to JSONB", zap.Error(err))
+					}
+					return jsonb
+				}(),
+				CredentialID: dummyCredentialID,
+				State:        integration.IntegrationStateSample,
+			},
+		}
+		err = a.database.CreateIntegration(&dbIntegration)
+		if err != nil {
+			a.logger.Error("failed to create integration", zap.Error(err))
+			return echo.NewHTTPError(http.StatusInternalServerError, "failed to create integration")
+		}
+	}
+
+	return c.NoContent(http.StatusOK)
+}
+
+// RemovePluginDemoData godoc
+//
+// @Summary			Remove demo data for plugin
+// @Description		Remove demo data for plugin by the given url
+// @Security		BearerToken
+// @Tags			integration_types
+// @Produce			json
+// @Param			id	path	string	true	"plugin id"
+// @Success			200
+// @Router			/integration/api/v1/plugin/{id}/demo/remove [put]
+func (a *API) RemovePluginDemoData(c echo.Context) error {
+	id := c.Param("id")
+
+	plugin, err := a.database.GetPluginByID(id)
+	if err != nil {
+		a.logger.Error("failed to get plugin", zap.Error(err))
+		return echo.NewHTTPError(http.StatusInternalServerError, "failed to get plugin")
+	}
+	if plugin == nil {
+		return echo.NewHTTPError(http.StatusNotFound, "plugin not found")
+	}
+
+	if !plugin.DemoDataLoaded {
+		return echo.NewHTTPError(http.StatusConflict, "plugin demo data not loaded")
+	}
+
+	if err = a.database.DeletePluginSampleIntegrations(plugin.IntegrationType); err != nil {
+		a.logger.Error("failed to delete plugin sample integration", zap.Error(err))
+		return echo.NewHTTPError(http.StatusInternalServerError, "failed to delete plugin sample integration")
+	}
 
 	return c.NoContent(http.StatusOK)
 }

services/integration/db/integration.go

@@ -39,6 +39,20 @@ func (db Database) DeleteIntegration(IntegrationID uuid.UUID) error {
 	return nil
 }
 
+// DeletePluginSampleIntegrations deletes sample integrations for a plugin
+func (db Database) DeletePluginSampleIntegrations(integrationType integration.Type) error {
+	tx := db.Orm.
+		Where("state = ?", integration.IntegrationStateSample).
+		Where("integration_type = ?", integrationType).
+		Unscoped().
+		Delete(&models.Integration{})
+	if tx.Error != nil {
+		return tx.Error
+	}
+
+	return nil
+}
+
 // DeleteSampleIntegrations deletes sample integrations
 func (db Database) DeleteSampleIntegrations() error {
 	tx := db.Orm.

services/integration/db/integration_plugin.go

@@ -32,6 +32,30 @@ func (db Database) GetPluginByID(pluginID string) (*models.IntegrationPlugin, er
 	return &plugin, nil
 }
 
+func (db Database) PluginDemoDataLoad(pluginID string) error {
+	err := db.IntegrationTypeOrm.Model(models.IntegrationPlugin{}).Where("plugin_id = ?", pluginID).
+		Update("demo_data_loaded = ?", true).Error
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil
+		}
+		return err
+	}
+	return nil
+}
+
+func (db Database) PluginDemoDataUnLoad(pluginID string) error {
+	err := db.IntegrationTypeOrm.Model(models.IntegrationPlugin{}).Where("plugin_id = ?", pluginID).
+		Update("demo_data_loaded = ?", false).Error
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil
+		}
+		return err
+	}
+	return nil
+}
+
 func (db Database) UpdatePluginInstallTimedOut(pluginInstallTime int64) error {
 	tx := db.IntegrationTypeOrm.
 		Model(&models.IntegrationPlugin{}).

services/integration/demo-import/import.go

@@ -1,7 +1,9 @@
 package demo_import
 
 import (
+	"encoding/json"
 	"fmt"
+	"github.com/opengovern/og-util/pkg/integration"
 	"go.uber.org/zap"
 	"io"
 	"net/http"
@@ -25,14 +27,23 @@ type Config struct {
 	ElasticsearchAddr string
 }
 
-func LoadDemoData(cfg Config, logger *zap.Logger) error {
+type Integration struct {
+	IntegrationID   string
+	ProviderID      string
+	Name            string
+	IntegrationType integration.Type
+	Annotations     map[string]string
+	Labels          map[string]string
+}
+
+func LoadDemoData(cfg Config, logger *zap.Logger) ([]Integration, error) {
 	logger.Info("Starting data loading process with provided config...")
 
 	wd, err := os.Getwd()
 	if err != nil {
 		// Log error before returning
 		logger.Error("Failed to get current working directory", zap.Error(err))
-		return fmt.Errorf("failed to get current working directory: %w", err)
+		return nil, fmt.Errorf("failed to get current working directory: %w", err)
 	}
 	workDir := wd
 
@@ -42,7 +53,13 @@ func LoadDemoData(cfg Config, logger *zap.Logger) error {
 	decryptedFilePath := filepath.Join(workDir, "demo_data.tar.gz")
 	// Assuming tar extracts relative paths inside the archive into the current dir (workDir).
 	inputPathForDump := filepath.Join(workDir, "/demo-data/es-demo/")
-	//:= filepath.Join(workDir, "/demo-data/integrations.json")
+	integrationsJsonFilePath := filepath.Join(workDir, "/demo-data/integrations.json")
+
+	integrations, err := loadIntegrationsFromJSON(integrationsJsonFilePath)
+	if err != nil {
+		logger.Error("Failed to load integrations.", zap.Error(err))
+		return nil, fmt.Errorf("failed to load integrations.: %w", err)
+	}
 
 	// --- Ensure cleanup of intermediate files ---
 	defer func() {
@@ -67,7 +84,7 @@ func LoadDemoData(cfg Config, logger *zap.Logger) error {
 		zap.String("destination", encryptedFilePath))
 	err = downloadFile(encryptedFilePath, cfg.DemoDataURL, logger)
 	if err != nil {
-		return fmt.Errorf("failed to download file: %w", err)
+		return nil, fmt.Errorf("failed to download file: %w", err)
 	}
 	logger.Info("Successfully downloaded file", zap.String("path", encryptedFilePath))
 
@@ -84,7 +101,7 @@ func LoadDemoData(cfg Config, logger *zap.Logger) error {
 	}
 	err = runCommand(workDir, "openssl", logger, opensslArgs...) // Pass logger
 	if err != nil {
-		return fmt.Errorf("failed to run openssl command: %w", err)
+		return nil, fmt.Errorf("failed to run openssl command: %w", err)
 	}
 	logger.Info("Successfully decrypted file", zap.String("path", decryptedFilePath))
 
@@ -95,7 +112,7 @@ func LoadDemoData(cfg Config, logger *zap.Logger) error {
 	tarArgs := []string{"-xvf", decryptedFilePath}
 	err = runCommand(workDir, "tar", logger, tarArgs...) // Pass logger
 	if err != nil {
-		return fmt.Errorf("failed to run tar command: %w", err)
+		return nil, fmt.Errorf("failed to run tar command: %w", err)
 	}
 	logger.Info("Successfully extracted tarball", zap.String("path", decryptedFilePath))
 
@@ -140,12 +157,12 @@ func LoadDemoData(cfg Config, logger *zap.Logger) error {
 	err = cmd.Run()
 	if err != nil {
 		logger.Error("multielasticdump command failed", zap.Error(err))
-		return fmt.Errorf("failed to run multielasticdump command: %w", err)
+		return nil, fmt.Errorf("failed to run multielasticdump command: %w", err)
 	}
 	logger.Info("Successfully ran multielasticdump.")
 
 	logger.Info("Data loading process completed successfully.")
-	return nil // Success
+	return integrations, nil
 }
 
 // downloadFile downloads a file from a URL and saves it locally, using zap for logging.
@@ -219,3 +236,23 @@ func runCommand(workDir string, name string, logger *zap.Logger, args ...string)
 	logger.Debug("Command executed successfully", zap.String("command", name))
 	return nil
 }
+
+func loadIntegrationsFromJSON(filePath string) ([]Integration, error) {
+	jsonData, err := os.ReadFile(filePath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read file %s: %w", filePath, err)
+	}
+
+	if len(jsonData) == 0 {
+		return []Integration{}, nil
+	}
+
+	var integrations []Integration
+
+	err = json.Unmarshal(jsonData, &integrations)
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal JSON from %s: %w", filePath, err)
+	}
+
+	return integrations, nil
+}

services/integration/models/integration_plugin.go

@@ -57,6 +57,7 @@ type IntegrationPlugin struct {
 	URL                      string
 	DescriberURL             string
 	DemoDataURL              string
+	DemoDataLoaded           bool
 	DescriberTag             string
 	OperationalStatusUpdates pgtype.JSONB `gorm:"default:'[]'"`
 	Tags                     pgtype.JSONB

services/tasks/http.go

@@ -364,7 +364,7 @@ func (r *httpRoutes) ListTaskRunResults(ctx echo.Context) error {
 		var result map[string]interface{}
 		err = json.Unmarshal(task.Result.Bytes, &result)
 		if err != nil {
-			r.logger.Error("failed to unmarshal result", zap.Error(err))
+			r.logger.Error("failed to unmarshal result json", zap.Error(err))
 			return ctx.JSON(http.StatusInternalServerError, "failed to unmarshal result")
 		}
 		taskRunResponses = append(taskRunResponses, api.TaskRun{