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 }