fix: add progress to task

Author: Arta Asadi

discovery/task/run-task.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/opengovern/og-describer-kubernetes/discovery/envs"
 	"github.com/opengovern/og-describer-kubernetes/discovery/pkg/orchestrator"
 	authApi "github.com/opengovern/og-util/pkg/api"
 	"github.com/opengovern/og-util/pkg/describe"
@@ -54,6 +55,24 @@ func NewTaskRunner(ctx context.Context, jq *jq.JobQueue, coreServiceEndpoint str
 }
 
 type TaskResult struct {
+	AllIntegrations             []string                      `json:"all_integrations"`
+	AllIntegrationsCount        int                           `json:"all_integrations_count"`
+	ProgressedIntegrations      map[string]*IntegrationResult `json:"proposed_integrations"`
+	ProgressedIntegrationsCount int                           `json:"proposed_integrations_count"`
+}
+
+type IntegrationResult struct {
+	IntegrationID              string               `json:"integration_id"`
+	AllResourceTypes           []string             `json:"all_resource_types"`
+	AllResourceTypesCount      int                  `json:"all_resource_types_count"`
+	ResourceTypeResults        []ResourceTypeResult `json:"resource_type_results"`
+	FinishedResourceTypesCount int                  `json:"finished_resource_types_count"`
+}
+
+type ResourceTypeResult struct {
+	ResourceType  string `json:"resource_type"`
+	Error         string `json:"error"`
+	ResourceCount int    `json:"resource_count"`
 }
 
 type ResourceType struct {
@@ -86,8 +105,15 @@ func (tr *TaskRunner) RunTask(ctx context.Context) error {
 		}
 	}
 
+	taskResult.AllIntegrations = make([]string, len(integrations))
+	for _, i := range integrations {
+		taskResult.AllIntegrations = append(taskResult.AllIntegrations, i.IntegrationID)
+	}
+	taskResult.AllIntegrationsCount = len(integrations)
+	taskResult.ProgressedIntegrations = make(map[string]*IntegrationResult)
+
 	for _, i := range integrations {
-		err = tr.describeIntegrationResourceTypes(ctx, i)
+		err = tr.describeIntegrationResourceTypes(ctx, i, taskResult)
 		if err != nil {
 			tr.logger.Error("Error describing integrations", zap.Error(err))
 			return err
@@ -104,7 +130,12 @@ func (tr *TaskRunner) RunTask(ctx context.Context) error {
 	return nil
 }
 
-func (tr *TaskRunner) describeIntegrationResourceTypes(ctx context.Context, i Integration) error {
+func (tr *TaskRunner) describeIntegrationResourceTypes(ctx context.Context, i Integration, taskResult *TaskResult) error {
+	taskResult.ProgressedIntegrations[i.IntegrationID] = &IntegrationResult{
+		IntegrationID: i.IntegrationID,
+	}
+	taskResult.ProgressedIntegrationsCount++
+
 	var resourceTypes []ResourceType
 	var err error
 
@@ -124,6 +155,13 @@ func (tr *TaskRunner) describeIntegrationResourceTypes(ctx context.Context, i In
 		}
 	}
 
+	taskResult.ProgressedIntegrations[i.IntegrationID].AllResourceTypes = make([]string, len(resourceTypes))
+	for _, rt := range resourceTypes {
+		taskResult.ProgressedIntegrations[i.IntegrationID].AllResourceTypes = append(taskResult.ProgressedIntegrations[i.IntegrationID].AllResourceTypes, rt.Name)
+	}
+	taskResult.ProgressedIntegrations[i.IntegrationID].AllResourceTypesCount = len(resourceTypes)
+	taskResult.ProgressedIntegrations[i.IntegrationID].ResourceTypeResults = make([]ResourceTypeResult, len(resourceTypes))
+
 	for _, rt := range resourceTypes {
 		params := make(map[string]string)
 		for key, value := range tr.request.TaskDefinition.Params {
@@ -144,10 +182,36 @@ func (tr *TaskRunner) describeIntegrationResourceTypes(ctx context.Context, i In
 			IntegrationLabels:      i.Labels,
 			IntegrationAnnotations: i.Annotations,
 		}
-		_, err = orchestrator.Describe(ctx, tr.logger, job, params, config, tr.request.EsDeliverEndpoint,
+		resources, err := orchestrator.Describe(ctx, tr.logger, job, params, config, tr.request.EsDeliverEndpoint,
 			tr.request.IngestionPipelineEndpoint, tr.describeToken, tr.request.UseOpenSearch)
+		errMsg := ""
 		if err != nil {
 			tr.logger.Error("Error describing job", zap.Error(err))
+			errMsg = err.Error()
+		}
+		taskResult.ProgressedIntegrations[i.IntegrationID].ResourceTypeResults = append(
+			taskResult.ProgressedIntegrations[i.IntegrationID].ResourceTypeResults,
+			ResourceTypeResult{
+				ResourceType:  rt.Name,
+				Error:         errMsg,
+				ResourceCount: len(resources),
+			})
+
+		taskResult.ProgressedIntegrations[i.IntegrationID].FinishedResourceTypesCount = len(taskResult.ProgressedIntegrations[i.IntegrationID].ResourceTypeResults)
+		jsonBytes, err := json.Marshal(taskResult)
+		if err != nil {
+			err = fmt.Errorf("failed Marshaling task result: %s", err.Error())
+			return err
+		}
+		tr.response.Result = jsonBytes
+		responseJson, marshalErr := json.Marshal(tr.response)
+		if marshalErr != nil {
+			tr.logger.Error("failed to create final job result json", zap.Error(marshalErr))
+			return marshalErr
+		}
+		msgId := fmt.Sprintf("task-run-result-%d", tr.request.TaskDefinition.RunID)
+		if _, err = tr.jq.Produce(ctx, envs.ResultTopicName, responseJson, msgId); err != nil { // Use original ctx
+			tr.logger.Error("failed to publish initial InProgress job status", zap.String("response", string(responseJson)), zap.Error(err))
 			return err
 		}
 	}

discovery/worker/worker.go

@@ -2,11 +2,9 @@ package worker
 
 import (
 	"context"
-	"encoding/base64"
 	"encoding/json"
 	"errors" // Added import
 	"fmt"
-	"github.com/golang-jwt/jwt/v5"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
 	"github.com/opengovern/og-describer-kubernetes/discovery/envs"
@@ -17,7 +15,6 @@ import (
 	"github.com/opengovern/opensecurity/services/tasks/db/models"
 	"github.com/opengovern/opensecurity/services/tasks/scheduler"
 	"go.uber.org/zap"
-	"os"
 	"strconv"
 	"time"
 )
@@ -231,11 +228,6 @@ func (w *Worker) ProcessMessage(ctx context.Context, msg jetstream.Msg) (err err
 		}
 	}()
 
-	//token, err := getJWTAuthToken()
-	//if err != nil {
-	//	return fmt.Errorf("failed to get JWT token: %w", err)
-	//}
-
 	msgLogger.Info("Starting task execution")
 	taskRunner, err := task.NewTaskRunner(ctxWithCancel, w.jq, envs.InventoryServiceEndpoint, "", w.esClient, msgLogger, request, response)
 	if err != nil {
@@ -251,28 +243,3 @@ func (w *Worker) ProcessMessage(ctx context.Context, msg jetstream.Msg) (err err
 
 	return err
 }
-
-func getJWTAuthToken() (string, error) {
-	privateKey, ok := os.LookupEnv("JWT_PRIVATE_KEY")
-	if !ok {
-		return "", fmt.Errorf("JWT_PRIVATE_KEY not set")
-	}
-
-	privateKeyBytes, err := base64.StdEncoding.DecodeString(privateKey)
-	if err != nil {
-		return "", fmt.Errorf("JWT_PRIVATE_KEY not base64 encoded")
-	}
-
-	pk, err := jwt.ParseRSAPrivateKeyFromPEM(privateKeyBytes)
-	if err != nil {
-		return "", fmt.Errorf("JWT_PRIVATE_KEY not valid")
-	}
-
-	token, err := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
-		"email": "[email protected]",
-	}).SignedString(pk)
-	if err != nil {
-		return "", fmt.Errorf("JWT token generation failed %v", err)
-	}
-	return token, nil
-}