mirror of
https://github.com/crawlab-team/crawlab.git
synced 2026-01-21 17:21:09 +01:00
e2cb99e56a
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.
107 lines
2.3 KiB
Go
107 lines
2.3 KiB
Go
package utils
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"github.com/crawlab-team/crawlab/core/constants"
|
|
"github.com/crawlab-team/crawlab/core/models/models"
|
|
models2 "github.com/crawlab-team/crawlab/core/models/models/v2"
|
|
"github.com/upper/db/v4"
|
|
"github.com/upper/db/v4/adapter/mssql"
|
|
"time"
|
|
)
|
|
|
|
func GetCockroachdbSession(ds *models.DataSource) (s db.Session, err error) {
|
|
return getCockroachdbSession(context.Background(), ds)
|
|
}
|
|
|
|
func GetCockroachdbSessionWithTimeout(ds *models.DataSource, timeout time.Duration) (s db.Session, err error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
return getCockroachdbSession(ctx, ds)
|
|
}
|
|
|
|
func getCockroachdbSession(ctx context.Context, ds *models.DataSource) (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
|
|
}
|
|
|
|
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.DatabaseV2) (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
|
|
}
|