iotlcm/server/internal/library/tdengine/tdengine.go

package tdengine

import (
	"context"
	"database/sql"
	"sync"

	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
	_ "github.com/taosdata/driver-go/v3/taosRestful"
	_ "github.com/taosdata/driver-go/v3/taosWS"
)

const (
	// dailyStable 和 monthlyStable 是 TDengine 中已经提前创建好的超级表名称。
	// 业务启动时只连接和读写，不负责自动建库建表。
	dailyStable                = "t_sim_card_usage_daily"
	monthlyStable              = "t_sim_card_usage_monthly"
	poolDailyStable            = "t_flow_pool_usage_daily"
	poolMonthlyStable          = "t_flow_pool_usage_monthly"
	poolCardCountDailyStable   = "t_flow_pool_card_count_daily"
	poolCardCountMonthlyStable = "t_flow_pool_card_count_monthly"
	defaultBatchSize           = 200
	// defaultDailyDbName / defaultMonthlyDbName 是配置缺失时的兑底库名。
	// 日流量与月流量分库（duration 不同，而 duration 只能建库时确定）。
	defaultDailyDbName   = "iotlcm_daily"
	defaultMonthlyDbName = "iotlcm_monthly"
)

// Config 是项目内 TDengine 连接和流量超级表配置。
//
// 日流量与月流量分库存储（iotlcm_daily / iotlcm_monthly）：两库的 duration 需要
// 不同——日流量按天分片、月流量按月分片——而 TDengine 的 duration 只能建库时确定、
// 事后无法 alter。DbName 保留为兼容位，指向不支持分库的流量池等表的落库位置。
type Config struct {
	Enabled                    bool
	DriverName                 string
	Dsn                        string
	DbName                     string
	DailyDbName                string
	MonthlyDbName              string
	DailyStable                string
	MonthlyStable              string
	PoolDailyStable            string
	PoolMonthlyStable          string
	PoolCardCountDailyStable   string
	PoolCardCountMonthlyStable string
	MaxOpenConns               int
	MaxIdleConns               int
	BatchSize                  int
}

// pools 按库名缓存连接池。DSN 不携带库名（形如 ws(host:6041)/），
// 因此同一个 DSN 可以服务多个 database，日/月分库无需两套连接串。
var (
	poolsMu sync.Mutex
	pools   = map[string]*sql.DB{}
)

// GetConfig 从 GoFrame 配置中读取 TDengine 配置。
// 表名保持在代码中固定，避免业务配置文件里暴露过多 TDengine 内部结构。
func GetConfig(ctx context.Context) Config {
	database := g.Cfg().MustGet(ctx, "tsdb.database", "TdEngine").String()
	// dbName 只是连接池的键名与兑底值，不决定实际访问哪个库：
	// DSN 末尾留空、库名由 SQL 里的 库名.表名 指定。
	// 默认值取日库，避免兑底到一个不存在的库名。
	dbName := g.Cfg().MustGet(ctx, "tsdb.tdengine.dbName", defaultDailyDbName).String()
	// dailyDbName / monthlyDbName 未配置时回退到 dbName。
	dailyDbName := g.Cfg().MustGet(ctx, "tsdb.tdengine.dailyDbName", dbName).String()
	monthlyDbName := g.Cfg().MustGet(ctx, "tsdb.tdengine.monthlyDbName", dbName).String()
	return Config{
		Enabled:                    database == "TdEngine",
		DriverName:                 g.Cfg().MustGet(ctx, "tsdb.tdengine.type", "taosWS").String(),
		Dsn:                        g.Cfg().MustGet(ctx, "tsdb.tdengine.dsn", "").String(),
		DbName:                     dbName,
		DailyDbName:                dailyDbName,
		MonthlyDbName:              monthlyDbName,
		DailyStable:                dailyStable,
		MonthlyStable:              monthlyStable,
		PoolDailyStable:            poolDailyStable,
		PoolMonthlyStable:          poolMonthlyStable,
		PoolCardCountDailyStable:   poolCardCountDailyStable,
		PoolCardCountMonthlyStable: poolCardCountMonthlyStable,
		MaxOpenConns:               g.Cfg().MustGet(ctx, "tsdb.tdengine.maxOpenConns", 100).Int(),
		MaxIdleConns:               g.Cfg().MustGet(ctx, "tsdb.tdengine.maxIdleConns", 10).Int(),
		BatchSize:                  g.Cfg().MustGet(ctx, "tsdb.tdengine.batchSize", defaultBatchSize).Int(),
	}
}

// DbNameFor 返回指定超级表应该写入的 database 名。
//
// 日流量与月流量分库：两库的 TDengine duration 不同（日库按天分片、月库按月分片），
// 而 duration 只能建库时确定。未单独配置 dailyDbName/monthlyDbName 时回退到 DbName。
func (c Config) DbNameFor(stable string) string {
	switch stable {
	case c.DailyStable, c.PoolDailyStable, c.PoolCardCountDailyStable:
		if c.DailyDbName != "" {
			return c.DailyDbName
		}
	case c.MonthlyStable, c.PoolMonthlyStable, c.PoolCardCountMonthlyStable:
		if c.MonthlyDbName != "" {
			return c.MonthlyDbName
		}
	}
	return c.DbName
}

// Enabled 只判断当前项目是否启用 TDengine，并且 DSN 是否已经配置。
func Enabled(ctx context.Context) bool {
	cfg := GetConfig(ctx)
	return cfg.Enabled && cfg.Dsn != ""
}

// DB 懒加载并复用默认库（DbName）的连接池。
// 注意：这里不会创建 database/stable，库表需要先通过 SQL 脚本准备好。
func DB(ctx context.Context) (*sql.DB, error) {
	return DBForDb(ctx, GetConfig(ctx).DbName)
}

// DBForDb 返回指定 database 的连接池；同一库名复用同一个池。
func DBForDb(ctx context.Context, dbName string) (*sql.DB, error) {
	cfg := GetConfig(ctx)
	if !cfg.Enabled {
		return nil, gerror.New("tdengine is disabled")
	}
	if cfg.Dsn == "" {
		return nil, gerror.New("tdengine dsn is empty")
	}
	if dbName == "" {
		dbName = cfg.DbName
	}

	poolsMu.Lock()
	defer poolsMu.Unlock()
	if pool, ok := pools[dbName]; ok {
		return pool, nil
	}

	driverName := cfg.DriverName
	if driverName == "" {
		driverName = "taosWS"
	}
	pool, err := sql.Open(driverName, cfg.Dsn)
	if err != nil {
		return nil, err
	}
	pool.SetMaxIdleConns(cfg.MaxIdleConns)
	pool.SetMaxOpenConns(cfg.MaxOpenConns)
	if err := pool.PingContext(ctx); err != nil {
		_ = pool.Close()
		return nil, err
	}
	pools[dbName] = pool
	return pool, nil
}

// Exec 执行 TDengine 默认库的写入、更新、删除类 SQL。
func Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
	conn, err := DB(ctx)
	if err != nil {
		return nil, err
	}
	return conn.ExecContext(ctx, query, args...)
}

// ExecForDb 在指定 database 的连接池上执行写入类 SQL。
func ExecForDb(ctx context.Context, dbName, query string, args ...any) (sql.Result, error) {
	conn, err := DBForDb(ctx, dbName)
	if err != nil {
		return nil, err
	}
	return conn.ExecContext(ctx, query, args...)
}

// Query 执行 TDengine 查询 SQL，调用方负责关闭返回的 rows。
func Query(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
	conn, err := DB(ctx)
	if err != nil {
		return nil, err
	}
	return conn.QueryContext(ctx, query, args...)
}

// QueryRow 执行 TDengine 单行查询 SQL。
func QueryRow(ctx context.Context, query string, args ...any) (*sql.Row, error) {
	conn, err := DB(ctx)
	if err != nil {
		return nil, err
	}
	return conn.QueryRowContext(ctx, query, args...), nil
}

// Close 在应用退出时关闭所有 TDengine 连接池。
func Close() error {
	poolsMu.Lock()
	defer poolsMu.Unlock()
	var firstErr error
	for name, pool := range pools {
		if err := pool.Close(); err != nil && firstErr == nil {
			firstErr = err
		}
		delete(pools, name)
	}
	return firstErr
}