feat: Update models to use DatabaseV2 instead of DataSourceV2

The code changes update the models and related functions to use the new DatabaseV2 struct instead of the deprecated DataSourceV2 struct. This change ensures consistency and clarity in the codebase.
crawlab-team · Aug 4, 2024 · e2cb99e · e2cb99e
1 parent a7da3a0
commit e2cb99e
Show file tree

Hide file tree

Showing 18 changed files with 62 additions and 67 deletions.
diff --git a/core/controllers/data_source_v2.go → core/controllers/database_v2.go b/core/controllers/data_source_v2.go → core/controllers/database_v2.go
@@ -9,7 +9,7 @@ import (
 	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
-func PostDataSource(c *gin.Context) {
+func PostDatabase(c *gin.Context) {
 	// data source
 	var payload struct {
 		Name        string            `json:"name"`
@@ -35,7 +35,7 @@ func PostDataSource(c *gin.Context) {
 	u := GetUserFromContextV2(c)
 
 	// add data source to db
-	dataSource := models.DataSourceV2{
+	dataSource := models.DatabaseV2{
 		Name:        payload.Name,
 		Type:        payload.Type,
 		Description: payload.Description,
@@ -53,7 +53,7 @@ func PostDataSource(c *gin.Context) {
 	}
 	dataSource.SetCreated(u.Id)
 	dataSource.SetUpdated(u.Id)
-	id, err := service.NewModelServiceV2[models.DataSourceV2]().InsertOne(dataSource)
+	id, err := service.NewModelServiceV2[models.DatabaseV2]().InsertOne(dataSource)
 	if err != nil {
 		HandleErrorInternalServerError(c, err)
 		return
@@ -68,21 +68,21 @@ func PostDataSource(c *gin.Context) {
 	HandleSuccessWithData(c, dataSource)
 }
 
-func PutDataSourceById(c *gin.Context) {
+func PutDatabaseById(c *gin.Context) {
 	id, err := primitive.ObjectIDFromHex(c.Param("id"))
 	if err != nil {
 		HandleErrorInternalServerError(c, err)
 		return
 	}
 
 	// data source
-	var dataSource models.DataSourceV2
+	var dataSource models.DatabaseV2
 	if err := c.ShouldBindJSON(&dataSource); err != nil {
 		HandleErrorBadRequest(c, err)
 		return
 	}
 
-	err = service.NewModelServiceV2[models.DataSourceV2]().ReplaceById(id, dataSource)
+	err = service.NewModelServiceV2[models.DatabaseV2]().ReplaceById(id, dataSource)
 	if err != nil {
 		HandleErrorInternalServerError(c, err)
 		return
@@ -94,7 +94,7 @@ func PutDataSourceById(c *gin.Context) {
 	}()
 }
 
-func PostDataSourceChangePassword(c *gin.Context) {
+func PostDatabaseChangePassword(c *gin.Context) {
 	id, err := primitive.ObjectIDFromHex(c.Param("id"))
 	if err != nil {
 		HandleErrorBadRequest(c, err)

diff --git a/core/controllers/result_v2.go b/core/controllers/result_v2.go
@@ -39,10 +39,10 @@ func GetResultList(c *gin.Context) {
 	}
 
 	// data source
-	ds, err := service.NewModelServiceV2[models2.DataSourceV2]().GetById(dsId)
+	ds, err := service.NewModelServiceV2[models2.DatabaseV2]().GetById(dsId)
 	if err != nil {
 		if err.Error() == mongo2.ErrNoDocuments.Error() {
-			ds = &models2.DataSourceV2{}
+			ds = &models2.DatabaseV2{}
 		} else {
 			HandleErrorInternalServerError(c, err)
 			return

diff --git a/core/controllers/router_v2.go b/core/controllers/router_v2.go
@@ -56,21 +56,16 @@ func InitRoutes(app *gin.Engine) (err error) {
 	groups := NewRouterGroups(app)
 
 	RegisterController(groups.AuthGroup, "/data/collections", NewControllerV2[models2.DataCollectionV2]())
-	RegisterController(groups.AuthGroup, "/data-sources", NewControllerV2[models2.DataSourceV2]([]Action{
+	RegisterController(groups.AuthGroup, "/databases", NewControllerV2[models2.DatabaseV2]([]Action{
 		{
 			Method:      http.MethodPost,
 			Path:        "",
-			HandlerFunc: PostDataSource,
+			HandlerFunc: PostDatabase,
 		},
 		{
 			Method:      http.MethodPut,
 			Path:        "/:id",
-			HandlerFunc: PutDataSourceById,
-		},
-		{
-			Method:      http.MethodPost,
-			Path:        "/:id/change-password",
-			HandlerFunc: PostDataSourceChangePassword,
+			HandlerFunc: PutDatabaseById,
 		},
 	}...))
 	RegisterController(groups.AuthGroup, "/environments", NewControllerV2[models2.EnvironmentV2]())

diff --git a/core/controllers/spider_v2.go b/core/controllers/spider_v2.go
@@ -699,7 +699,7 @@ func GetSpiderDataSource(c *gin.Context) {
 	}
 
 	// data source
-	ds, err := service.NewModelServiceV2[models2.DataSourceV2]().GetById(s.DataSourceId)
+	ds, err := service.NewModelServiceV2[models2.DatabaseV2]().GetById(s.DataSourceId)
 	if err != nil {
 		if err.Error() == mongo2.ErrNoDocuments.Error() {
 			HandleSuccess(c)
@@ -736,7 +736,7 @@ func PostSpiderDataSource(c *gin.Context) {
 
 	// data source
 	if !dsId.IsZero() {
-		_, err = service.NewModelServiceV2[models2.DataSourceV2]().GetById(dsId)
+		_, err = service.NewModelServiceV2[models2.DatabaseV2]().GetById(dsId)
 		if err != nil {
 			HandleErrorInternalServerError(c, err)
 			return

diff --git a/core/ds/service_v2.go b/core/ds/service_v2.go
@@ -51,7 +51,7 @@ func (svc *ServiceV2) Stop() {
 }
 
 func (svc *ServiceV2) ChangePassword(id primitive.ObjectID, password string, by primitive.ObjectID) (err error) {
-	dataSource, err := service.NewModelServiceV2[models.DataSourceV2]().GetById(id)
+	dataSource, err := service.NewModelServiceV2[models.DatabaseV2]().GetById(id)
 	if err != nil {
 		return err
 	}
@@ -60,7 +60,7 @@ func (svc *ServiceV2) ChangePassword(id primitive.ObjectID, password string, by
 		return err
 	}
 	dataSource.SetUpdated(by)
-	err = service.NewModelServiceV2[models.DataSourceV2]().ReplaceById(id, *dataSource)
+	err = service.NewModelServiceV2[models.DatabaseV2]().ReplaceById(id, *dataSource)
 	if err != nil {
 		return err
 	}
@@ -85,7 +85,7 @@ func (svc *ServiceV2) Monitor() {
 }
 
 func (svc *ServiceV2) CheckStatus(id primitive.ObjectID) (err error) {
-	ds, err := service.NewModelServiceV2[models.DataSourceV2]().GetById(id)
+	ds, err := service.NewModelServiceV2[models.DatabaseV2]().GetById(id)
 	if err != nil {
 		return err
 	}
@@ -106,7 +106,7 @@ func (svc *ServiceV2) monitor() (err error) {
 	log.Debugf("[DataSourceService] start monitoring")
 
 	// data source list
-	dataSources, err := service.NewModelServiceV2[models.DataSourceV2]().GetMany(nil, nil)
+	dataSources, err := service.NewModelServiceV2[models.DatabaseV2]().GetMany(nil, nil)
 	if err != nil {
 		return err
 	}
@@ -118,7 +118,7 @@ func (svc *ServiceV2) monitor() (err error) {
 	// iterate data source list
 	for _, ds := range dataSources {
 		// async operation
-		go func(ds *models.DataSourceV2) {
+		go func(ds *models.DatabaseV2) {
 			// check status and save
 			_ = svc.checkStatus(ds, true)
 
@@ -137,7 +137,7 @@ func (svc *ServiceV2) monitor() (err error) {
 	return nil
 }
 
-func (svc *ServiceV2) checkStatus(ds *models.DataSourceV2, save bool) (err error) {
+func (svc *ServiceV2) checkStatus(ds *models.DatabaseV2, save bool) (err error) {
 	// check status
 	if err := svc._checkStatus(ds); err != nil {
 		ds.Status = constants2.DataSourceStatusOffline
@@ -155,12 +155,12 @@ func (svc *ServiceV2) checkStatus(ds *models.DataSourceV2, save bool) (err error
 	return nil
 }
 
-func (svc *ServiceV2) _save(ds *models.DataSourceV2) (err error) {
+func (svc *ServiceV2) _save(ds *models.DatabaseV2) (err error) {
 	log.Debugf("[DataSourceService] saving data source: name=%s, type=%s, status=%s, error=%s", ds.Name, ds.Type, ds.Status, ds.Error)
-	return service.NewModelServiceV2[models.DataSourceV2]().ReplaceById(ds.Id, *ds)
+	return service.NewModelServiceV2[models.DatabaseV2]().ReplaceById(ds.Id, *ds)
 }
 
-func (svc *ServiceV2) _checkStatus(ds *models.DataSourceV2) (err error) {
+func (svc *ServiceV2) _checkStatus(ds *models.DatabaseV2) (err error) {
 	switch ds.Type {
 	case constants.DataSourceTypeMongo:
 		_, err := utils2.GetMongoClientWithTimeoutV2(ds, svc.timeout)

diff --git a/core/grpc/server/model_base_service_v2_server.go b/core/grpc/server/model_base_service_v2_server.go
@@ -18,7 +18,7 @@ var (
 	typeOneInstances    = []any{
 		*new(models2.TestModelV2),
 		*new(models2.DataCollectionV2),
-		*new(models2.DataSourceV2),
+		*new(models2.DatabaseV2),
 		*new(models2.DependencyV2),
 		*new(models2.DependencyLogV2),
 		*new(models2.DependencySettingV2),

diff --git a/core/models/common/index_service_v2.go b/core/models/common/index_service_v2.go
@@ -82,7 +82,7 @@ func CreateIndexesV2() {
 	})
 
 	// data sources
-	mongo.GetMongoCol(service.GetCollectionNameByInstance(models2.DataSourceV2{})).MustCreateIndexes([]mongo2.IndexModel{
+	mongo.GetMongoCol(service.GetCollectionNameByInstance(models2.DatabaseV2{})).MustCreateIndexes([]mongo2.IndexModel{
 		{Keys: bson.M{"name": 1}},
 	})
 

diff --git a/core/models/models/v2/data_source_v2.go b/core/models/models/v2/data_source_v2.go
diff --git a/core/models/models/v2/database_v2.go b/core/models/models/v2/database_v2.go
@@ -0,0 +1,20 @@
+package models
+
+type DatabaseV2 struct {
+	any                     `collection:"databases"`
+	BaseModelV2[DatabaseV2] `bson:",inline"`
+	Name                    string            `json:"name" bson:"name"`
+	Type                    string            `json:"type" bson:"type"`
+	Description             string            `json:"description" bson:"description"`
+	Host                    string            `json:"host" bson:"host"`
+	Port                    string            `json:"port" bson:"port"`
+	Url                     string            `json:"url" bson:"url"`
+	Hosts                   []string          `json:"hosts" bson:"hosts"`
+	Database                string            `json:"database" bson:"database"`
+	Username                string            `json:"username" bson:"username"`
+	Password                string            `json:"-,omitempty" bson:"password"`
+	ConnectType             string            `json:"connect_type" bson:"connect_type"`
+	Status                  string            `json:"status" bson:"status"`
+	Error                   string            `json:"error" bson:"error"`
+	Extra                   map[string]string `json:"extra,omitempty" bson:"extra,omitempty"`
+}
diff --git a/core/models/models/v2/spider_v2.go b/core/models/models/v2/spider_v2.go
@@ -11,7 +11,7 @@ type SpiderV2 struct {
 	ColId                 primitive.ObjectID   `json:"col_id" bson:"col_id"`                 // data collection id
 	ColName               string               `json:"col_name,omitempty" bson:"-"`          // data collection name
 	DataSourceId          primitive.ObjectID   `json:"data_source_id" bson:"data_source_id"` // data source id
-	DataSource            *DataSourceV2        `json:"data_source,omitempty" bson:"-"`       // data source
+	DataSource            *DatabaseV2          `json:"data_source,omitempty" bson:"-"`       // data source
 	Description           string               `json:"description" bson:"description"`       // description
 	ProjectId             primitive.ObjectID   `json:"project_id" bson:"project_id"`         // Project.Id
 	Mode                  string               `json:"mode" bson:"mode"`                     // default Task.Mode

diff --git a/core/utils/cockroachdb.go b/core/utils/cockroachdb.go
@@ -60,13 +60,13 @@ func getCockroachdbSession(ctx context.Context, ds *models.DataSource) (s db.Ses
 	return s, err
 }
 
-func GetCockroachdbSessionWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (s db.Session, err error) {
+func GetCockroachdbSessionWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (s db.Session, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getCockroachdbSessionV2(ctx, ds)
 }
 
-func getCockroachdbSessionV2(ctx context.Context, ds *models2.DataSourceV2) (s db.Session, err error) {
+func getCockroachdbSessionV2(ctx context.Context, ds *models2.DatabaseV2) (s db.Session, err error) {
 	// normalize settings
 	host := ds.Host
 	port := ds.Port

diff --git a/core/utils/es.go b/core/utils/es.go
@@ -84,13 +84,13 @@ func getElasticsearchClient(ctx context.Context, ds *models.DataSource) (c *elas
 	return c, err
 }
 
-func GetElasticsearchClientWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (c *elasticsearch.Client, err error) {
+func GetElasticsearchClientWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (c *elasticsearch.Client, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getElasticsearchClientV2(ctx, ds)
 }
 
-func getElasticsearchClientV2(ctx context.Context, ds *models2.DataSourceV2) (c *elasticsearch.Client, err error) {
+func getElasticsearchClientV2(ctx context.Context, ds *models2.DatabaseV2) (c *elasticsearch.Client, err error) {
 	// normalize settings
 	host := ds.Host
 	port := ds.Port

diff --git a/core/utils/kafka.go b/core/utils/kafka.go
@@ -41,13 +41,13 @@ func getKafkaConnection(ctx context.Context, ds *models.DataSource) (c *kafka.Co
 	return kafka.DialLeader(ctx, network, address, topic, partition)
 }
 
-func GetKafkaConnectionWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (c *kafka.Conn, err error) {
+func GetKafkaConnectionWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (c *kafka.Conn, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getKafkaConnectionV2(ctx, ds)
 }
 
-func getKafkaConnectionV2(ctx context.Context, ds *models2.DataSourceV2) (c *kafka.Conn, err error) {
+func getKafkaConnectionV2(ctx context.Context, ds *models2.DatabaseV2) (c *kafka.Conn, err error) {
 	// normalize settings
 	host := ds.Host
 	port := ds.Port

diff --git a/core/utils/mongo.go b/core/utils/mongo.go
@@ -55,7 +55,7 @@ func GetMongoClientWithTimeout(ds *models.DataSource, timeout time.Duration) (c
 	return getMongoClient(ctx, ds)
 }
 
-func GetMongoClientWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (c *mongo2.Client, err error) {
+func GetMongoClientWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (c *mongo2.Client, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getMongoClientV2(ctx, ds)
@@ -100,7 +100,7 @@ func getMongoClient(ctx context.Context, ds *models.DataSource) (c *mongo2.Clien
 	return mongo.GetMongoClient(opts...)
 }
 
-func getMongoClientV2(ctx context.Context, ds *models2.DataSourceV2) (c *mongo2.Client, err error) {
+func getMongoClientV2(ctx context.Context, ds *models2.DatabaseV2) (c *mongo2.Client, err error) {
 	// normalize settings
 	if ds.Host == "" {
 		ds.Host = constants.DefaultHost

diff --git a/core/utils/mssql.go b/core/utils/mssql.go
@@ -60,13 +60,13 @@ func getMssqlSession(ctx context.Context, ds *models.DataSource) (s db.Session,
 	return s, err
 }
 
-func GetMssqlSessionWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (s db.Session, err error) {
+func GetMssqlSessionWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (s db.Session, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getMssqlSessionV2(ctx, ds)
 }
 
-func getMssqlSessionV2(ctx context.Context, ds *models2.DataSourceV2) (s db.Session, err error) {
+func getMssqlSessionV2(ctx context.Context, ds *models2.DatabaseV2) (s db.Session, err error) {
 	// normalize settings
 	host := ds.Host
 	port := ds.Port

diff --git a/core/utils/mysql.go b/core/utils/mysql.go
@@ -60,13 +60,13 @@ func getMysqlSession(ctx context.Context, ds *models.DataSource) (s db.Session,
 	return s, err
 }
 
-func GetMysqlSessionWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (s db.Session, err error) {
+func GetMysqlSessionWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (s db.Session, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getMysqlSessionV2(ctx, ds)
 }
 
-func getMysqlSessionV2(ctx context.Context, ds *models2.DataSourceV2) (s db.Session, err error) {
+func getMysqlSessionV2(ctx context.Context, ds *models2.DatabaseV2) (s db.Session, err error) {
 	// normalize settings
 	host := ds.Host
 	port := ds.Port

diff --git a/core/utils/postgresql.go b/core/utils/postgresql.go
@@ -60,13 +60,13 @@ func getPostgresqlSession(ctx context.Context, ds *models.DataSource) (s db.Sess
 	return s, err
 }
 
-func GetPostgresqlSessionWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (s db.Session, err error) {
+func GetPostgresqlSessionWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (s db.Session, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getPostgresqlSessionV2(ctx, ds)
 }
 
-func getPostgresqlSessionV2(ctx context.Context, ds *models2.DataSourceV2) (s db.Session, err error) {
+func getPostgresqlSessionV2(ctx context.Context, ds *models2.DatabaseV2) (s db.Session, err error) {
 	// normalize settings
 	host := ds.Host
 	port := ds.Port

diff --git a/core/utils/sqlite.go b/core/utils/sqlite.go
@@ -45,13 +45,13 @@ func getSqliteSession(ctx context.Context, ds *models.DataSource) (s db.Session,
 	return s, err
 }
 
-func GetSqliteSessionWithTimeoutV2(ds *models2.DataSourceV2, timeout time.Duration) (s db.Session, err error) {
+func GetSqliteSessionWithTimeoutV2(ds *models2.DatabaseV2, timeout time.Duration) (s db.Session, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 	return getSqliteSessionV2(ctx, ds)
 }
 
-func getSqliteSessionV2(ctx context.Context, ds *models2.DataSourceV2) (s db.Session, err error) {
+func getSqliteSessionV2(ctx context.Context, ds *models2.DatabaseV2) (s db.Session, err error) {
 	// connect settings
 	settings := sqlite.ConnectionURL{
 		Database: ds.Database,