feat: added dependency config setup (wip)

crawlab-team · Dec 15, 2024 · c5c08df · c5c08df
1 parent 3b54a63
commit c5c08df
Show file tree

Hide file tree

Showing 9 changed files with 295 additions and 88 deletions.
diff --git a/core/grpc/server/dependency_service_server.go b/core/grpc/server/dependency_service_server.go
@@ -145,7 +145,6 @@ func (svr DependencyServiceServer) Sync(_ context.Context, request *grpc.Depende
 }
 
 func (svr DependencyServiceServer) UpdateLogs(stream grpc.DependencyService_UpdateLogsServer) (err error) {
-	var dep *models.Dependency
 	for {
 		// receive message
 		req, err := stream.Recv()
@@ -157,26 +156,19 @@ func (svr DependencyServiceServer) UpdateLogs(stream grpc.DependencyService_Upda
 			return err
 		}
 
-		// if dependency is nil, get dependency
-		if dep == nil {
-			id, err := primitive.ObjectIDFromHex(req.DependencyId)
-			if err != nil {
-				log.Errorf("[DependencyServiceServer] convert dependency id error: %v", err)
-				return err
-			}
-			dep, err = service.NewModelService[models.Dependency]().GetById(id)
-			if err != nil {
-				log.Errorf("[DependencyServiceServer] get dependency error: %v", err)
-				return err
-			}
+		// get id
+		id, err := primitive.ObjectIDFromHex(req.TargetId)
+		if err != nil {
+			log.Errorf("[DependencyServiceServer] convert dependency id error: %v", err)
+			return err
 		}
 
 		// insert dependency logs
 		var depLogs []models.DependencyLog
 		for _, line := range req.Logs {
 			depLog := models.DependencyLog{
-				DependencyId: dep.Id,
-				Content:      line,
+				TargetId: id,
+				Content:  line,
 			}
 			depLogs = append(depLogs, depLog)
 		}
@@ -188,6 +180,57 @@ func (svr DependencyServiceServer) UpdateLogs(stream grpc.DependencyService_Upda
 	}
 }
 
+func (svr DependencyServiceServer) SyncConfigSetup(_ context.Context, request *grpc.DependencyServiceSyncConfigSetupRequest) (response *grpc.Response, err error) {
+	// Get node by node key
+	n, err := service.NewModelService[models.Node]().GetOne(bson.M{"key": request.NodeKey}, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get config
+	cfg, err := service.NewModelService[models.DependencyConfig]().GetOne(bson.M{"key": request.Lang}, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get config setup for the node
+	cs, err := service.NewModelService[models.DependencyConfigSetup]().GetOne(bson.M{
+		"node_id":              n.Id,
+		"dependency_config_id": cfg.Id,
+	}, nil)
+	if err != nil {
+		if !errors.Is(err, mongo.ErrNoDocuments) {
+			log.Errorf("[DependencyService] get dependency config setup from db error: %v", err)
+			return nil, err
+		}
+	}
+
+	if cs == nil {
+		// Create new config setup
+		cs = &models.DependencyConfigSetup{
+			NodeId:             n.Id,
+			DependencyConfigId: cfg.Id,
+			Status:             request.Status,
+			Error:              request.Error,
+		}
+		_, err = service.NewModelService[models.DependencyConfigSetup]().InsertOne(*cs)
+		if err != nil {
+			log.Errorf("[DependencyService] insert dependency config setup error: %v", err)
+			return nil, err
+		}
+	} else {
+		// Update existing config setup
+		cs.Status = request.Status
+		cs.Error = request.Error
+		err = service.NewModelService[models.DependencyConfigSetup]().ReplaceById(cs.Id, *cs)
+		if err != nil {
+			log.Errorf("[DependencyService] update dependency config setup error: %v", err)
+			return nil, err
+		}
+	}
+	return nil, nil
+}
+
 func (svr DependencyServiceServer) GetStream(key string) (stream *grpc.DependencyService_ConnectServer, err error) {
 	svr.mu.Lock()
 	defer svr.mu.Unlock()

diff --git a/core/models/models/dependency.go b/core/models/models/dependency.go
@@ -14,7 +14,6 @@ type Dependency struct {
 	Version               string               `json:"version" bson:"version"`
 	Status                string               `json:"status" bson:"status"`
 	Error                 string               `json:"error,omitempty" bson:"error,omitempty"`
-	Logs                  []string             `json:"logs,omitempty" bson:"logs,omitempty"`
 	NodeIds               []primitive.ObjectID `json:"node_ids,omitempty" bson:"-"`
 	Versions              []string             `json:"versions,omitempty" bson:"-"`
 }
diff --git a/core/models/models/dependency_config.go b/core/models/models/dependency_config.go
@@ -3,10 +3,9 @@ package models
 type DependencyConfig struct {
 	any                         `collection:"dependency_configs"`
 	BaseModel[DependencyConfig] `bson:",inline"`
-	Key                         string   `json:"key" bson:"key"`
-	Name                        string   `json:"name" bson:"name"`
-	Cmd                         string   `json:"cmd" bson:"cmd"`
-	Proxy                       string   `json:"proxy" bson:"proxy"`
-	SetupNodeIds                []string `json:"setup_node_ids" bson:"setup_node_ids"`
-	SetupScriptPath             string   `json:"setup_script_path" bson:"setup_script_path"`
+	Key                         string `json:"key" bson:"key"`
+	Name                        string `json:"name" bson:"name"`
+	ExecCmd                     string `json:"exec_cmd" bson:"exec_cmd"`
+	PkgCmd                      string `json:"pkg_cmd" bson:"pkg_cmd"`
+	Proxy                       string `json:"proxy" bson:"proxy"`
 }
diff --git a/core/models/models/dependency_config_setup.go b/core/models/models/dependency_config_setup.go
@@ -0,0 +1,13 @@
+package models
+
+import "go.mongodb.org/mongo-driver/bson/primitive"
+
+type DependencyConfigSetup struct {
+	any                              `collection:"dependency_config_setups"`
+	BaseModel[DependencyConfigSetup] `bson:",inline"`
+	DependencyConfigId               primitive.ObjectID `json:"dependency_config_id" bson:"dependency_config_id"`
+	NodeId                           primitive.ObjectID `json:"node_id" bson:"node_id"`
+	Version                          string             `json:"version" bson:"version"`
+	Status                           string             `json:"status" bson:"status"`
+	Error                            string             `json:"error,omitempty" bson:"error,omitempty"`
+}
diff --git a/core/models/models/dependency_log.go b/core/models/models/dependency_log.go
@@ -5,6 +5,6 @@ import "go.mongodb.org/mongo-driver/bson/primitive"
 type DependencyLog struct {
 	any                      `collection:"dependency_logs"`
 	BaseModel[DependencyLog] `bson:",inline"`
-	DependencyId             primitive.ObjectID `json:"dependency_id" bson:"dependency_id"`
+	TargetId                 primitive.ObjectID `json:"target_id" bson:"target_id"`
 	Content                  string             `json:"content" bson:"content"`
 }
diff --git a/core/utils/config.go b/core/utils/config.go
@@ -10,24 +10,25 @@ import (
 )
 
 const (
-	DefaultWorkspace           = "crawlab_workspace"
-	DefaultTaskLogPath         = "/var/log/crawlab/tasks"
-	DefaultServerHost          = "0.0.0.0"
-	DefaultServerPort          = 8000
-	DefaultGrpcHost            = "localhost"
-	DefaultGrpcPort            = 9666
-	DefaultGrpcServerHost      = "0.0.0.0"
-	DefaultGrpcServerPort      = 9666
-	DefaultAuthKey             = "Crawlab2024!"
-	DefaultApiEndpoint         = "http://localhost:8000"
-	DefaultApiAllowOrigin      = "*"
-	DefaultApiAllowCredentials = "true"
-	DefaultApiAllowMethods     = "DELETE, POST, OPTIONS, GET, PUT"
-	DefaultApiAllowHeaders     = "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With"
-	DefaultNodeMaxRunners      = 0 // 0 means no limit
-	MetadataConfigDirName      = ".crawlab"
-	MetadataConfigName         = "config.json"
-	PyenvRoot                  = "/root/.pyenv"
+	DefaultWorkspace                 = "crawlab_workspace"
+	DefaultTaskLogPath               = "/var/log/crawlab/tasks"
+	DefaultServerHost                = "0.0.0.0"
+	DefaultServerPort                = 8000
+	DefaultGrpcHost                  = "localhost"
+	DefaultGrpcPort                  = 9666
+	DefaultGrpcServerHost            = "0.0.0.0"
+	DefaultGrpcServerPort            = 9666
+	DefaultAuthKey                   = "Crawlab2024!"
+	DefaultApiEndpoint               = "http://localhost:8000"
+	DefaultApiAllowOrigin            = "*"
+	DefaultApiAllowCredentials       = "true"
+	DefaultApiAllowMethods           = "DELETE, POST, OPTIONS, GET, PUT"
+	DefaultApiAllowHeaders           = "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With"
+	DefaultNodeMaxRunners            = 0 // 0 means no limit
+	DefaultDependencySetupScriptRoot = "/app/install"
+	MetadataConfigDirName            = ".crawlab"
+	MetadataConfigName               = "config.json"
+	PyenvRoot                        = "/root/.pyenv"
 )
 
 func IsDev() bool {
@@ -221,3 +222,10 @@ func GetMetadataConfigPath() string {
 
 	return filepath.Join(homeDirPath, MetadataConfigDirName, MetadataConfigName)
 }
+
+func GetDependencySetupScriptRoot() string {
+	if res := viper.GetString("dependency.setupScriptRoot"); res != "" {
+		return res
+	}
+	return DefaultDependencySetupScriptRoot
+}