fix: missing data source issue

core/controllers/data_source_v2.go

@@ -0,0 +1,120 @@
+package controllers
+
+import (
+	"github.com/crawlab-team/crawlab/core/ds"
+	"github.com/crawlab-team/crawlab/core/errors"
+	"github.com/crawlab-team/crawlab/core/models/models"
+	"github.com/crawlab-team/crawlab/core/models/service"
+	"github.com/gin-gonic/gin"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+)
+
+func PostDataSource(c *gin.Context) {
+	// data source
+	var payload struct {
+		Name        string            `json:"name"`
+		Type        string            `json:"type"`
+		Description string            `json:"description"`
+		Host        string            `json:"host"`
+		Port        string            `json:"port"`
+		Url         string            `json:"url"`
+		Hosts       []string          `json:"hosts"`
+		Database    string            `json:"database"`
+		Username    string            `json:"username"`
+		Password    string            `json:"-,omitempty"`
+		ConnectType string            `json:"connect_type"`
+		Status      string            `json:"status"`
+		Error       string            `json:"error"`
+		Extra       map[string]string `json:"extra,omitempty"`
+	}
+	if err := c.ShouldBindJSON(&payload); err != nil {
+		HandleErrorBadRequest(c, err)
+		return
+	}
+
+	u := GetUserFromContextV2(c)
+
+	// add data source to db
+	dataSource := models.DataSourceV2{
+		Name:        payload.Name,
+		Type:        payload.Type,
+		Description: payload.Description,
+		Host:        payload.Host,
+		Port:        payload.Port,
+		Url:         payload.Url,
+		Hosts:       payload.Hosts,
+		Database:    payload.Database,
+		Username:    payload.Username,
+		Password:    payload.Password,
+		ConnectType: payload.ConnectType,
+		Status:      payload.Status,
+		Error:       payload.Error,
+		Extra:       payload.Extra,
+	}
+	dataSource.SetCreated(u.Id)
+	dataSource.SetUpdated(u.Id)
+	id, err := service.NewModelServiceV2[models.DataSourceV2]().InsertOne(dataSource)
+	if err != nil {
+		HandleErrorInternalServerError(c, err)
+		return
+	}
+	dataSource.Id = id
+
+	// check data source status
+	go func() {
+		_ = ds.GetDataSourceServiceV2().CheckStatus(id)
+	}()
+
+	HandleSuccessWithData(c, dataSource)
+}
+
+func PutDataSourceById(c *gin.Context) {
+	id, err := primitive.ObjectIDFromHex(c.Param("id"))
+	if err != nil {
+		HandleErrorInternalServerError(c, err)
+		return
+	}
+
+	// data source
+	var dataSource models.DataSourceV2
+	if err := c.ShouldBindJSON(&dataSource); err != nil {
+		HandleErrorBadRequest(c, err)
+		return
+	}
+
+	err = service.NewModelServiceV2[models.DataSourceV2]().ReplaceById(id, dataSource)
+	if err != nil {
+		HandleErrorInternalServerError(c, err)
+		return
+	}
+
+	// check data source status
+	go func() {
+		_ = ds.GetDataSourceServiceV2().CheckStatus(id)
+	}()
+}
+
+func PostDataSourceChangePassword(c *gin.Context) {
+	id, err := primitive.ObjectIDFromHex(c.Param("id"))
+	if err != nil {
+		HandleErrorBadRequest(c, err)
+		return
+	}
+	var payload struct {
+		Password string `json:"password"`
+	}
+	if err := c.ShouldBindJSON(&payload); err != nil {
+		HandleErrorBadRequest(c, err)
+		return
+	}
+	if payload.Password == "" {
+		HandleErrorBadRequest(c, errors.ErrorDataSourceMissingRequiredFields)
+		return
+	}
+	u := GetUserFromContextV2(c)
+	if err := ds.GetDataSourceServiceV2().ChangePassword(id, payload.Password, u.Id); err != nil {
+		HandleErrorInternalServerError(c, err)
+		return
+	}
+	HandleSuccess(c)
+}

core/controllers/router_v2.go

@@ -56,10 +56,26 @@ func InitRoutes(app *gin.Engine) (err error) {
 	groups := NewRouterGroups(app)
 
 	RegisterController(groups.AuthGroup, "/data/collections", NewControllerV2[models.DataCollectionV2]())
-	RegisterController(groups.AuthGroup, "/data-sources", NewControllerV2[models.DataSourceV2]())
+	RegisterController(groups.AuthGroup, "/data-sources", NewControllerV2[models.DataSourceV2]([]Action{
+		{
+			Method:      http.MethodPost,
+			Path:        "",
+			HandlerFunc: PostDataSource,
+		},
+		{
+			Method:      http.MethodPut,
+			Path:        "/:id",
+			HandlerFunc: PutDataSourceById,
+		},
+		{
+			Method:      http.MethodPost,
+			Path:        "/:id/change-password",
+			HandlerFunc: PostDataSourceChangePassword,
+		},
+	}...))
 	RegisterController(groups.AuthGroup, "/environments", NewControllerV2[models.EnvironmentV2]())
 	RegisterController(groups.AuthGroup, "/nodes", NewControllerV2[models.NodeV2]())
-	RegisterController(groups.AuthGroup, "/notifications/settings", NewControllerV2[models.SettingV2]())
+	RegisterController(groups.AuthGroup, "/notifications/settings", NewControllerV2[models.NotificationSettingV2]())
 	RegisterController(groups.AuthGroup, "/permissions", NewControllerV2[models.PermissionV2]())
 	RegisterController(groups.AuthGroup, "/projects", NewControllerV2[models.ProjectV2]([]Action{
 		{

core/ds/service_v2.go

@@ -0,0 +1,271 @@
+package ds
+
+import (
+	"github.com/apex/log"
+	"github.com/crawlab-team/crawlab/core/constants"
+	constants2 "github.com/crawlab-team/crawlab/core/constants"
+	"github.com/crawlab-team/crawlab/core/models/models"
+	"github.com/crawlab-team/crawlab/core/models/service"
+	"github.com/crawlab-team/crawlab/core/result"
+	"github.com/crawlab-team/crawlab/core/utils"
+	utils2 "github.com/crawlab-team/crawlab/core/utils"
+	"github.com/crawlab-team/crawlab/trace"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"sync"
+	"time"
+)
+
+type ServiceV2 struct {
+	// internals
+	timeout         time.Duration
+	monitorInterval time.Duration
+	stopped         bool
+}
+
+func (svc *ServiceV2) Init() {
+	// result service registry
+	reg := result.GetResultServiceRegistry()
+
+	// register result services
+	reg.Register(constants.DataSourceTypeMongo, NewDataSourceMongoService)
+	reg.Register(constants.DataSourceTypeMysql, NewDataSourceMysqlService)
+	reg.Register(constants.DataSourceTypePostgresql, NewDataSourcePostgresqlService)
+	reg.Register(constants.DataSourceTypeMssql, NewDataSourceMssqlService)
+	reg.Register(constants.DataSourceTypeSqlite, NewDataSourceSqliteService)
+	reg.Register(constants.DataSourceTypeCockroachdb, NewDataSourceCockroachdbService)
+	reg.Register(constants.DataSourceTypeElasticSearch, NewDataSourceElasticsearchService)
+	reg.Register(constants.DataSourceTypeKafka, NewDataSourceKafkaService)
+}
+
+func (svc *ServiceV2) Start() {
+	// start monitoring
+	go svc.Monitor()
+}
+
+func (svc *ServiceV2) Wait() {
+	utils.DefaultWait()
+}
+
+func (svc *ServiceV2) Stop() {
+	svc.stopped = true
+}
+
+func (svc *ServiceV2) ChangePassword(id primitive.ObjectID, password string, by primitive.ObjectID) (err error) {
+	dataSource, err := service.NewModelServiceV2[models.DataSourceV2]().GetById(id)
+	if err != nil {
+		return err
+	}
+	dataSource.Password, err = utils.EncryptAES(password)
+	if err != nil {
+		return err
+	}
+	dataSource.SetUpdated(by)
+	err = service.NewModelServiceV2[models.DataSourceV2]().ReplaceById(id, *dataSource)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (svc *ServiceV2) Monitor() {
+	for {
+		// return if stopped
+		if svc.stopped {
+			return
+		}
+
+		// monitor
+		if err := svc.monitor(); err != nil {
+			trace.PrintError(err)
+		}
+
+		// wait
+		time.Sleep(svc.monitorInterval)
+	}
+}
+
+func (svc *ServiceV2) CheckStatus(id primitive.ObjectID) (err error) {
+	ds, err := service.NewModelServiceV2[models.DataSourceV2]().GetById(id)
+	if err != nil {
+		return err
+	}
+	return svc.checkStatus(ds, true)
+}
+
+func (svc *ServiceV2) SetTimeout(duration time.Duration) {
+	svc.timeout = duration
+}
+
+func (svc *ServiceV2) SetMonitorInterval(duration time.Duration) {
+	svc.monitorInterval = duration
+}
+
+func (svc *ServiceV2) monitor() (err error) {
+	// start
+	tic := time.Now()
+	log.Debugf("[DataSourceService] start monitoring")
+
+	// data source list
+	dataSources, err := service.NewModelServiceV2[models.DataSourceV2]().GetMany(nil, nil)
+	if err != nil {
+		return err
+	}
+
+	// waiting group
+	wg := sync.WaitGroup{}
+	wg.Add(len(dataSources))
+
+	// iterate data source list
+	for _, ds := range dataSources {
+		// async operation
+		go func(ds *models.DataSourceV2) {
+			// check status and save
+			_ = svc.checkStatus(ds, true)
+
+			// release
+			wg.Done()
+		}(&ds)
+	}
+
+	// wait
+	wg.Wait()
+
+	// finish
+	toc := time.Now()
+	log.Debugf("[DataSourceService] finished monitoring. elapsed: %d ms", (toc.Sub(tic)).Milliseconds())
+
+	return nil
+}
+
+func (svc *ServiceV2) checkStatus(ds *models.DataSourceV2, save bool) (err error) {
+	// check status
+	if err := svc._checkStatus(ds); err != nil {
+		ds.Status = constants2.DataSourceStatusOffline
+		ds.Error = err.Error()
+	} else {
+		ds.Status = constants2.DataSourceStatusOnline
+		ds.Error = ""
+	}
+
+	// save
+	if save {
+		return svc._save(ds)
+	}
+
+	return nil
+}
+
+func (svc *ServiceV2) _save(ds *models.DataSourceV2) (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)
+}
+
+func (svc *ServiceV2) _checkStatus(ds *models.DataSourceV2) (err error) {
+	switch ds.Type {
+	case constants.DataSourceTypeMongo:
+		_, err := utils2.GetMongoClientWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+	case constants.DataSourceTypeMysql:
+		s, err := utils2.GetMysqlSessionWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			err := s.Close()
+			if err != nil {
+				return err
+			}
+		}
+	case constants.DataSourceTypePostgresql:
+		s, err := utils2.GetPostgresqlSessionWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			err := s.Close()
+			if err != nil {
+				return err
+			}
+		}
+	case constants.DataSourceTypeMssql:
+		s, err := utils2.GetMssqlSessionWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			err := s.Close()
+			if err != nil {
+				return err
+			}
+		}
+	case constants.DataSourceTypeSqlite:
+		s, err := utils2.GetSqliteSessionWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			err := s.Close()
+			if err != nil {
+				return err
+			}
+		}
+	case constants.DataSourceTypeCockroachdb:
+		s, err := utils2.GetCockroachdbSessionWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+		if s != nil {
+			err := s.Close()
+			if err != nil {
+				return err
+			}
+		}
+	case constants.DataSourceTypeElasticSearch:
+		_, err := utils2.GetElasticsearchClientWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+	case constants.DataSourceTypeKafka:
+		c, err := utils2.GetKafkaConnectionWithTimeoutV2(ds, svc.timeout)
+		if err != nil {
+			return err
+		}
+		if c != nil {
+			err := c.Close()
+			if err != nil {
+				return err
+			}
+		}
+	default:
+		log.Warnf("[DataSourceService] invalid data source type: %s", ds.Type)
+	}
+	return nil
+}
+
+func NewDataSourceServiceV2() *ServiceV2 {
+	// service
+	svc := &ServiceV2{
+		monitorInterval: 15 * time.Second,
+		timeout:         10 * time.Second,
+	}
+
+	// initialize
+	svc.Init()
+
+	// start
+	svc.Start()
+
+	return svc
+}
+
+var _dsSvcV2 *ServiceV2
+
+func GetDataSourceServiceV2() *ServiceV2 {
+	if _dsSvcV2 != nil {
+		return _dsSvcV2
+	}
+	_dsSvcV2 = NewDataSourceServiceV2()
+	return _dsSvcV2
+}

core/grpc/server/model_base_service_v2_server.go

@@ -26,6 +26,7 @@ var (
 		*new(models.EnvironmentV2),
 		*new(models.GitV2),
 		*new(models.NodeV2),
+		*new(models.NotificationSettingV2),
 		*new(models.PermissionV2),
 		*new(models.ProjectV2),
 		*new(models.RolePermissionV2),

core/models/models/data_source_v2.go

@@ -12,7 +12,7 @@ type DataSourceV2 struct {
 	Hosts                   []string          `json:"hosts" bson:"hosts"`
 	Database                string            `json:"database" bson:"database"`
 	Username                string            `json:"username" bson:"username"`
-	Password                string            `json:"password,omitempty" bson:"-"`
+	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"`

core/models/models/notification_setting_v2.go

@@ -3,6 +3,8 @@ package models
 import "go.mongodb.org/mongo-driver/bson/primitive"
 
 type NotificationSettingV2 struct {
+	any                                `collection:"notification_settings"`
+	BaseModelV2[NotificationSettingV2] `bson:",inline"`
 	Id                                 primitive.ObjectID        `json:"_id" bson:"_id"`
 	Type                               string                    `json:"type" bson:"type"`
 	Name                               string                    `json:"name" bson:"name"`

core/models/models/spider_v2.go

@@ -5,7 +5,7 @@ import (
 )
 
 type SpiderV2 struct {
-	any                   `collection:"spiders"` // spider id
+	any                   `collection:"spiders"`
 	BaseModelV2[SpiderV2] `bson:",inline"`
 	Name                  string               `json:"name" bson:"name"`                     // spider name
 	Type                  string               `json:"type" bson:"type"`                     // spider type

core/node/service/master_service_v2.go

@@ -36,9 +36,9 @@ type MasterServiceV2 struct {
 	schedulerSvc    *scheduler.ServiceV2
 	handlerSvc      *handler.ServiceV2
 	scheduleSvc     *schedule.ServiceV2
-	notificationSvc *notification.Service
+	notificationSvc *notification.ServiceV2
 	spiderAdminSvc  *admin.ServiceV2
-	systemSvc       *system.Service
+	systemSvc       *system.ServiceV2
 
 	// settings
 	cfgPath         string
@@ -368,7 +368,7 @@ func NewMasterServiceV2() (res interfaces.NodeMasterService, err error) {
 	}
 
 	// notification service
-	svc.notificationSvc = notification.GetService()
+	svc.notificationSvc = notification.GetServiceV2()
 
 	// spider admin service
 	svc.spiderAdminSvc, err = admin.GetSpiderAdminServiceV2()
@@ -377,7 +377,7 @@ func NewMasterServiceV2() (res interfaces.NodeMasterService, err error) {
 	}
 
 	// system service
-	svc.systemSvc = system.GetService()
+	svc.systemSvc = system.GetServiceV2()
 
 	// init
 	if err := svc.Init(); err != nil {

core/notification/service_v2.go

@@ -41,7 +41,6 @@ func (svc *ServiceV2) initData() (err error) {
 	// data to initialize
 	settings := []models.NotificationSettingV2{
 		{
-			Id:          primitive.NewObjectID(),
 			Type:        TypeMail,
 			Enabled:     true,
 			Name:        "任务通知（邮件）",
@@ -77,7 +76,6 @@ func (svc *ServiceV2) initData() (err error) {
 			},
 		},
 		{
-			Id:          primitive.NewObjectID(),
 			Type:        TypeMail,
 			Enabled:     true,
 			Name:        "Task Change (Mail)",
@@ -113,7 +111,6 @@ Please find the task data as below.
 			},
 		},
 		{
-			Id:          primitive.NewObjectID(),
 			Type:        TypeMobile,
 			Enabled:     true,
 			Name:        "任务通知（移动端）",
@@ -142,7 +139,6 @@ Please find the task data as below.
 			Mobile: models.NotificationSettingMobile{},
 		},
 		{
-			Id:          primitive.NewObjectID(),
 			Type:        TypeMobile,
 			Enabled:     true,
 			Name:        "Task Change (Mobile)",

core/system/service_v2.go

@@ -0,0 +1,67 @@
+package system
+
+import (
+	"github.com/crawlab-team/crawlab/core/models/models"
+	"github.com/crawlab-team/crawlab/core/models/service"
+	"go.mongodb.org/mongo-driver/bson"
+)
+
+type ServiceV2 struct {
+}
+
+func (svc *ServiceV2) Init() (err error) {
+	// initialize data
+	if err := svc.initData(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (svc *ServiceV2) initData() (err error) {
+	total, err := service.NewModelServiceV2[models.SettingV2]().Count(bson.M{
+		"key": "site_title",
+	})
+	if err != nil {
+		return err
+	}
+	if total > 0 {
+		return nil
+	}
+
+	// data to initialize
+	settings := []models.SettingV2{
+		{
+			Key: "site_title",
+			Value: bson.M{
+				"customize_site_title": false,
+				"site_title":           "",
+			},
+		},
+	}
+	_, err = service.NewModelServiceV2[models.SettingV2]().InsertMany(settings)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func NewServiceV2() *ServiceV2 {
+	// service
+	svc := &ServiceV2{}
+
+	if err := svc.Init(); err != nil {
+		panic(err)
+	}
+
+	return svc
+}
+
+var _serviceV2 *ServiceV2
+
+func GetServiceV2() *ServiceV2 {
+	if _serviceV2 == nil {
+		_serviceV2 = NewServiceV2()
+	}
+	return _serviceV2
+}

core/utils/cockroachdb.go

@@ -58,3 +58,48 @@ func getCockroachdbSession(ctx context.Context, ds *models.DataSource) (s db.Ses
 
 	return s, err
 }
+
+func GetCockroachdbSessionWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (s db.Session, err error) {
+	// normalize settings
+	host := ds.Host
+	port := ds.Port
+	if ds.Host == "" {
+		host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		port = constants.DefaultCockroachdbPort
+	}
+
+	// connect settings
+	settings := mssql.ConnectionURL{
+		User:     ds.Username,
+		Password: ds.Password,
+		Database: ds.Database,
+		Host:     fmt.Sprintf("%s:%s", host, port),
+		Options:  nil,
+	}
+
+	// session
+	done := make(chan struct{})
+	go func() {
+		s, err = mssql.Open(settings)
+		close(done)
+	}()
+
+	// wait for done
+	select {
+	case <-ctx.Done():
+		if ctx.Err() != nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
+	return s, err
+}

core/utils/es.go

@@ -50,34 +50,75 @@ func getElasticsearchClient(ctx context.Context, ds *models.DataSource) (c *elas
 		Addresses: addresses,
 		Username:  ds.Username,
 		Password:  ds.Password,
-		//CloudID:                 "",
-		//APIKey:                  "",
-		//ServiceToken:            "",
-		//CertificateFingerprint:  "",
-		//Header:                  nil,
-		//CACert:                  nil,
-		//RetryOnStatus:           nil,
-		//DisableRetry:            false,
-		//EnableRetryOnTimeout:    false,
-		//MaxRetries:              0,
-		//CompressRequestBody:     false,
-		//DiscoverNodesOnStart:    false,
-		//DiscoverNodesInterval:   0,
-		//EnableMetrics:           false,
-		//EnableDebugLogger:       false,
-		//EnableCompatibilityMode: false,
-		//DisableMetaHeader:       false,
-		//UseResponseCheckOnly:    false,
 		RetryBackoff: func(i int) time.Duration {
 			if i == 1 {
 				rb.Reset()
 			}
 			return rb.NextBackOff()
 		},
-		//Transport:               nil,
-		//Logger:                  nil,
-		//Selector:                nil,
-		//ConnectionPoolFunc:      nil,
+	}
+
+	// es client
+	done := make(chan struct{})
+	go func() {
+		c, err = elasticsearch.NewClient(cfg)
+		if err != nil {
+			return
+		}
+		var res *esapi.Response
+		res, err = c.Info()
+		fmt.Println(res)
+		close(done)
+	}()
+
+	// wait for done
+	select {
+	case <-ctx.Done():
+		if ctx.Err() != nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
+	return c, err
+}
+
+func GetElasticsearchClientWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (c *elasticsearch.Client, err error) {
+	// normalize settings
+	host := ds.Host
+	port := ds.Port
+	if ds.Host == "" {
+		host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		port = constants.DefaultElasticsearchPort
+	}
+
+	// es hosts
+	addresses := []string{
+		fmt.Sprintf("http://%s:%s", host, port),
+	}
+
+	// retry backoff
+	rb := backoff.NewExponentialBackOff()
+
+	// es client options
+	cfg := elasticsearch.Config{
+		Addresses: addresses,
+		Username:  ds.Username,
+		Password:  ds.Password,
+		RetryBackoff: func(i int) time.Duration {
+			if i == 1 {
+				rb.Reset()
+			}
+			return rb.NextBackOff()
+		},
 	}
 
 	// es client

core/utils/kafka.go

@@ -39,3 +39,30 @@ func getKafkaConnection(ctx context.Context, ds *models.DataSource) (c *kafka.Co
 	// kafka connection
 	return kafka.DialLeader(ctx, network, address, topic, partition)
 }
+
+func GetKafkaConnectionWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (c *kafka.Conn, err error) {
+	// normalize settings
+	host := ds.Host
+	port := ds.Port
+	if ds.Host == "" {
+		host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		port = constants.DefaultKafkaPort
+	}
+
+	// kafka connection address
+	network := "tcp"
+	address := fmt.Sprintf("%s:%s", host, port)
+	topic := ds.Database
+	partition := 0 // TODO: parameterize
+
+	// kafka connection
+	return kafka.DialLeader(ctx, network, address, topic, partition)
+}

core/utils/mongo.go

@@ -54,6 +54,12 @@ func GetMongoClientWithTimeout(ds *models.DataSource, timeout time.Duration) (c
 	return getMongoClient(ctx, ds)
 }
 
+func GetMongoClientWithTimeoutV2(ds *models.DataSourceV2, timeout time.Duration) (c *mongo2.Client, err error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	return getMongoClientV2(ctx, ds)
+}
+
 func getMongoClient(ctx context.Context, ds *models.DataSource) (c *mongo2.Client, err error) {
 	// normalize settings
 	if ds.Host == "" {
@@ -92,3 +98,42 @@ func getMongoClient(ctx context.Context, ds *models.DataSource) (c *mongo2.Clien
 	// client
 	return mongo.GetMongoClient(opts...)
 }
+
+func getMongoClientV2(ctx context.Context, ds *models.DataSourceV2) (c *mongo2.Client, err error) {
+	// normalize settings
+	if ds.Host == "" {
+		ds.Host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		ds.Port = constants.DefaultMongoPort
+	}
+
+	// options
+	var opts []mongo.ClientOption
+	opts = append(opts, mongo.WithContext(ctx))
+	opts = append(opts, mongo.WithUri(ds.Url))
+	opts = append(opts, mongo.WithHost(ds.Host))
+	opts = append(opts, mongo.WithPort(ds.Port))
+	opts = append(opts, mongo.WithDb(ds.Database))
+	opts = append(opts, mongo.WithUsername(ds.Username))
+	opts = append(opts, mongo.WithPassword(ds.Password))
+	opts = append(opts, mongo.WithHosts(ds.Hosts))
+
+	// extra
+	if ds.Extra != nil {
+		// auth source
+		authSource, ok := ds.Extra["auth_source"]
+		if ok {
+			opts = append(opts, mongo.WithAuthSource(authSource))
+		}
+
+		// auth mechanism
+		authMechanism, ok := ds.Extra["auth_mechanism"]
+		if ok {
+			opts = append(opts, mongo.WithAuthMechanism(authMechanism))
+		}
+	}
+
+	// client
+	return mongo.GetMongoClient(opts...)
+}

core/utils/mssql.go

@@ -58,3 +58,48 @@ func getMssqlSession(ctx context.Context, ds *models.DataSource) (s db.Session,
 
 	return s, err
 }
+
+func GetMssqlSessionWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (s db.Session, err error) {
+	// normalize settings
+	host := ds.Host
+	port := ds.Port
+	if ds.Host == "" {
+		host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		port = constants.DefaultMssqlPort
+	}
+
+	// connect settings
+	settings := mssql.ConnectionURL{
+		User:     ds.Username,
+		Password: ds.Password,
+		Database: ds.Database,
+		Host:     fmt.Sprintf("%s:%s", host, port),
+		Options:  nil,
+	}
+
+	// session
+	done := make(chan struct{})
+	go func() {
+		s, err = mssql.Open(settings)
+		close(done)
+	}()
+
+	// wait for done
+	select {
+	case <-ctx.Done():
+		if ctx.Err() != nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
+	return s, err
+}

core/utils/mysql.go

@@ -58,3 +58,48 @@ func getMysqlSession(ctx context.Context, ds *models.DataSource) (s db.Session,
 
 	return s, err
 }
+
+func GetMysqlSessionWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (s db.Session, err error) {
+	// normalize settings
+	host := ds.Host
+	port := ds.Port
+	if ds.Host == "" {
+		host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		port = constants.DefaultMysqlPort
+	}
+
+	// connect settings
+	settings := mysql.ConnectionURL{
+		User:     ds.Username,
+		Password: ds.Password,
+		Database: ds.Database,
+		Host:     fmt.Sprintf("%s:%s", host, port),
+		Options:  nil,
+	}
+
+	// session
+	done := make(chan struct{})
+	go func() {
+		s, err = mysql.Open(settings)
+		close(done)
+	}()
+
+	// wait for done
+	select {
+	case <-ctx.Done():
+		if ctx.Err() != nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
+	return s, err
+}

core/utils/postgresql.go

@@ -58,3 +58,48 @@ func getPostgresqlSession(ctx context.Context, ds *models.DataSource) (s db.Sess
 
 	return s, err
 }
+
+func GetPostgresqlSessionWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (s db.Session, err error) {
+	// normalize settings
+	host := ds.Host
+	port := ds.Port
+	if ds.Host == "" {
+		host = constants.DefaultHost
+	}
+	if ds.Port == "" {
+		port = constants.DefaultPostgresqlPort
+	}
+
+	// connect settings
+	settings := postgresql.ConnectionURL{
+		User:     ds.Username,
+		Password: ds.Password,
+		Database: ds.Database,
+		Host:     fmt.Sprintf("%s:%s", host, port),
+		Options:  nil,
+	}
+
+	// session
+	done := make(chan struct{})
+	go func() {
+		s, err = postgresql.Open(settings)
+		close(done)
+	}()
+
+	// wait for done
+	select {
+	case <-ctx.Done():
+		if ctx.Err() != nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
+	return s, err
+}

core/utils/sqlite.go

@@ -43,3 +43,35 @@ func getSqliteSession(ctx context.Context, ds *models.DataSource) (s db.Session,
 
 	return s, err
 }
+
+func GetSqliteSessionWithTimeoutV2(ds *models.DataSourceV2, 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 *models.DataSourceV2) (s db.Session, err error) {
+	// connect settings
+	settings := sqlite.ConnectionURL{
+		Database: ds.Database,
+		Options:  nil,
+	}
+
+	// session
+	done := make(chan struct{})
+	go func() {
+		s, err = sqlite.Open(settings)
+		close(done)
+	}()
+
+	// wait for done
+	select {
+	case <-ctx.Done():
+		if ctx.Err() != nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
+	return s, err
+}