MiFaZhan
206 lines
7.1 KiB
Go
206 lines
7.1 KiB
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
|
|
}
|