crawlab/core/ds/service.go

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/container"
	"github.com/crawlab-team/crawlab/core/interfaces"
	"github.com/crawlab-team/crawlab/core/models/delegate"
	"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"
	"go.mongodb.org/mongo-driver/mongo"
	"sync"
	"time"
)

type Service struct {
	// dependencies
	modelSvc service.ModelService

	// internals
	timeout         time.Duration
	monitorInterval time.Duration
	stopped         bool
}

func (svc *Service) 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 *Service) Start() {
	// start monitoring
	go svc.Monitor()
}

func (svc *Service) Wait() {
	utils.DefaultWait()
}

func (svc *Service) Stop() {
	svc.stopped = true
}

func (svc *Service) ChangePassword(id primitive.ObjectID, password string) (err error) {
	p, err := svc.modelSvc.GetPasswordById(id)
	if err == nil {
		// exists, save
		encryptedPassword, err := utils.EncryptAES(password)
		if err != nil {
			return err
		}
		p.Password = encryptedPassword
		if err := delegate.NewModelDelegate(p).Save(); err != nil {
			return err
		}
		return nil
	} else if err.Error() == mongo.ErrNoDocuments.Error() {
		// not exists, add
		encryptedPassword, err := utils.EncryptAES(password)
		if err != nil {
			return err
		}
		p = &models.Password{
			Id:       id,
			Password: encryptedPassword,
		}
		if err := delegate.NewModelDelegate(p).Add(); err != nil {
			return err
		}
		return nil
	} else {
		// error
		return err
	}
}

func (svc *Service) 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 *Service) CheckStatus(id primitive.ObjectID) (err error) {
	ds, err := svc.modelSvc.GetDataSourceById(id)
	if err != nil {
		return err
	}
	return svc.checkStatus(ds, true)
}

func (svc *Service) SetTimeout(duration time.Duration) {
	svc.timeout = duration
}

func (svc *Service) SetMonitorInterval(duration time.Duration) {
	svc.monitorInterval = duration
}

func (svc *Service) monitor() (err error) {
	// start
	tic := time.Now()
	log.Debugf("[DataSourceService] start monitoring")

	// data source list
	dsList, err := svc.modelSvc.GetDataSourceList(nil, nil)
	if err != nil {
		return err
	}

	// waiting group
	wg := sync.WaitGroup{}
	wg.Add(len(dsList))

	// iterate data source list
	for _, ds := range dsList {
		// async operation
		go func(ds models.DataSource) {
			// 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 *Service) checkStatus(ds *models.DataSource, save bool) (err error) {
	// password
	if ds.Password == "" {
		pwd, err := svc.modelSvc.GetPasswordById(ds.Id)
		if err == nil {
			ds.Password, err = utils.DecryptAES(pwd.Password)
			if err != nil {
				return err
			}
		} else if err.Error() != mongo.ErrNoDocuments.Error() {
			return trace.TraceError(err)
		}
	}

	// 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 *Service) _save(ds *models.DataSource) (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 delegate.NewModelDelegate(ds).Save()
}

func (svc *Service) _checkStatus(ds *models.DataSource) (err error) {
	switch ds.Type {
	case constants.DataSourceTypeMongo:
		_, err := utils2.GetMongoClientWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
	case constants.DataSourceTypeMysql:
		s, err := utils2.GetMysqlSessionWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
		if s != nil {
			s.Close()
		}
	case constants.DataSourceTypePostgresql:
		s, err := utils2.GetPostgresqlSessionWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
		if s != nil {
			s.Close()
		}
	case constants.DataSourceTypeMssql:
		s, err := utils2.GetMssqlSessionWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
		if s != nil {
			s.Close()
		}
	case constants.DataSourceTypeSqlite:
		s, err := utils2.GetSqliteSessionWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
		if s != nil {
			s.Close()
		}
	case constants.DataSourceTypeCockroachdb:
		s, err := utils2.GetCockroachdbSessionWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
		if s != nil {
			s.Close()
		}
	case constants.DataSourceTypeElasticSearch:
		_, err := utils2.GetElasticsearchClientWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
	case constants.DataSourceTypeKafka:
		c, err := utils2.GetKafkaConnectionWithTimeout(ds, svc.timeout)
		if err != nil {
			return err
		}
		if c != nil {
			c.Close()
		}
	default:
		log.Warnf("[DataSourceService] invalid data source type: %s", ds.Type)
	}
	return nil
}

func NewDataSourceService(opts ...DataSourceServiceOption) (svc2 interfaces.DataSourceService, err error) {
	// service
	svc := &Service{
		monitorInterval: 15 * time.Second,
		timeout:         10 * time.Second,
	}

	// apply options
	for _, opt := range opts {
		opt(svc)
	}

	// dependency injection
	if err := container.GetContainer().Invoke(func(modelSvc service.ModelService) {
		svc.modelSvc = modelSvc
	}); err != nil {
		return nil, trace.TraceError(err)
	}

	// initialize
	svc.Init()

	// start
	svc.Start()

	return svc, nil
}

var _dsSvc interfaces.DataSourceService

func GetDataSourceService() (svc interfaces.DataSourceService, err error) {
	if _dsSvc != nil {
		return _dsSvc, nil
	}
	svc, err = NewDataSourceService()
	if err != nil {
		return nil, err
	}
	_dsSvc = svc
	return svc, nil
}