add tdengine support (#25)

* add tdengine support

---------

Co-authored-by: Rick <[email protected]>

Author: LinuxSuRen

README.md

@@ -3,12 +3,12 @@
 # atest-ext-store-orm
 ORM database Store Extension for API Testing
 
-This project provides an ORM-based database store extension for API testing, simplifying data storage and retrieval operations. It supports various databases including SQLite, MySQL, and others, making it versatile for different testing environments.
+This project provides an ORM-based database store extension for API testing, simplifying data storage and retrieval operations. It supports various databases including SQLite, MySQL, PostgreSQL, TDengine, and others, making it versatile for different testing environments.
 
 ## Features
 - Simplified database operations using ORM.
 - Integration with API testing frameworks.
-- Support for multiple databases (SQLite, MySQL, etc.).
+- Support for multiple databases (SQLite, MySQL, PostgreSQL, TDengine, etc.).
 
 ## Usage
 To use this extension in your API testing project, follow these steps:

cmd/root.go

@@ -38,6 +38,12 @@ func NewRootCommand() (c *cobra.Command) {
 }
 
 func (o *option) runE(c *cobra.Command, args []string) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			c.Println(r)
+		}
+	}()
+
 	if o.version {
 		c.Println(version.GetVersion())
 		c.Println(version.GetDate())

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/linuxsuren/api-testing v0.0.19-0.20250220092001-ed58adc30f20
 	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
+	github.com/taosdata/driver-go/v3 v3.6.0
 	google.golang.org/protobuf v1.33.0
 	gorm.io/driver/mysql v1.5.2
 	gorm.io/driver/postgres v1.5.4
@@ -29,9 +30,6 @@ require (
 	github.com/blang/semver/v4 v4.0.0 // indirect
 	github.com/bufbuild/protocompile v0.6.0 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
-	github.com/cucumber/gherkin-go/v19 v19.0.3 // indirect
-	github.com/cucumber/godog v0.12.6 // indirect
-	github.com/cucumber/messages-go/v16 v16.0.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/expr-lang/expr v1.15.6 // indirect
 	github.com/flopp/go-findfont v0.1.0 // indirect
@@ -43,33 +41,30 @@ require (
 	github.com/go-openapi/spec v0.21.0 // indirect
 	github.com/go-openapi/swag v0.23.0 // indirect
 	github.com/go-sql-driver/mysql v1.7.0 // indirect
-	github.com/gofrs/uuid v4.2.0+incompatible // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/gorilla/mux v1.8.1 // indirect
+	github.com/gorilla/websocket v1.5.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0 // indirect
-	github.com/hashicorp/go-immutable-radix v1.3.1 // indirect
-	github.com/hashicorp/go-memdb v1.3.2 // indirect
-	github.com/hashicorp/golang-lru v0.5.4 // indirect
 	github.com/huandu/xstrings v1.4.0 // indirect
-	github.com/iancoleman/orderedmap v0.3.0 // indirect
 	github.com/imdario/mergo v0.3.16 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
-	github.com/invopop/jsonschema v0.7.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
 	github.com/jackc/pgx/v5 v5.4.3 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/linuxsuren/go-fake-runtime v0.0.4 // indirect
-	github.com/linuxsuren/go-service v0.0.0-20231225060426-efabcd3a5161 // indirect
 	github.com/linuxsuren/oauth-hub v0.0.0-20240809060240-e78c21b5d8d4 // indirect
 	github.com/linuxsuren/unstructured v0.0.1 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-sqlite3 v1.14.22 // indirect
 	github.com/mitchellh/copystructure v1.2.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.2 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/phpdave11/gofpdi v1.0.14-0.20211212211723-1f10f9844311 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect

go.sum

@@ -17,9 +17,11 @@ package pkg
 
 import (
 	"context"
-
+	"fmt"
 	"github.com/linuxsuren/api-testing/pkg/server"
 	"gorm.io/gorm"
+	"reflect"
+	"time"
 )
 
 func (s *dbserver) Query(ctx context.Context, query *server.DataQuery) (result *server.DataQueryResult, err error) {
@@ -32,18 +34,31 @@ func (s *dbserver) Query(ctx context.Context, query *server.DataQuery) (result *
 	if err != nil {
 		return
 	}
-	defer rows.Close()
-
-	columns, err := rows.Columns()
-	if err != nil {
-		return
-	}
+	defer func() {
+		if rows != nil {
+			rows.Close()
+		}
+	}()
 
 	result = &server.DataQueryResult{
 		Data:  []*server.Pair{},
 		Items: make([]*server.Pairs, 0),
 	}
 
+	if rows == nil {
+		if rows, err = db.ConnPool.QueryContext(ctx, query.Sql); err != nil {
+			return
+		} else if rows == nil {
+			fmt.Println("no rows found")
+			return
+		}
+	}
+
+	columns, err := rows.Columns()
+	if err != nil {
+		return
+	}
+
 	for rows.Next() {
 		// Create a slice of interface{}'s to represent each column,
 		// and a second slice to contain pointers to each item in the columns slice.
@@ -63,10 +78,21 @@ func (s *dbserver) Query(ctx context.Context, query *server.DataQuery) (result *
 		for i, colName := range columns {
 			rowData := &server.Pair{}
 			val := columnsData[i]
-			b, ok := val.([]byte)
-			if ok {
-				rowData.Key = colName
-				rowData.Value = string(b)
+
+			rowData.Key = colName
+			switch v := val.(type) {
+			case []byte:
+				rowData.Value = string(v)
+			case string:
+				rowData.Value = v
+			case int, uint64, uint32, int32, int64:
+				rowData.Value = fmt.Sprintf("%d", v)
+			case float32, float64:
+				rowData.Value = fmt.Sprintf("%f", v)
+			case time.Time:
+				rowData.Value = v.String()
+			default:
+				fmt.Println("column", colName, "type", reflect.TypeOf(v))
 			}
 
 			// Append the map to our slice of maps.

pkg/data_query.go

@@ -68,6 +68,9 @@ func createDB(user, password, address, database, driver string) (db *gorm.DB, er
 		}
 		dsn = fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s sslmode=disable TimeZone=Asia/Shanghai", host, user, password, database, port)
 		dialector = postgres.Open(dsn)
+	case "tdengine":
+		dsn = fmt.Sprintf("%s:%s@ws(%s)/", user, password, address)
+		dialector = NewTDengineDialector(dsn)
 	default:
 		err = fmt.Errorf("invalid database driver %q", driver)
 		return
@@ -82,9 +85,11 @@ func createDB(user, password, address, database, driver string) (db *gorm.DB, er
 		return
 	}
 
-	err = errors.Join(err, db.AutoMigrate(&TestCase{}))
-	err = errors.Join(err, db.AutoMigrate(&TestSuite{}))
-	err = errors.Join(err, db.AutoMigrate(&HistoryTestResult{}))
+	if driver != "tdengine" {
+		err = errors.Join(err, db.AutoMigrate(&TestCase{}))
+		err = errors.Join(err, db.AutoMigrate(&TestSuite{}))
+		err = errors.Join(err, db.AutoMigrate(&HistoryTestResult{}))
+	}
 	return
 }
 

pkg/server.go

@@ -0,0 +1,119 @@
+/*
+Copyright 2025 API Testing Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package pkg
+
+import (
+	"database/sql"
+	"fmt"
+
+	_ "github.com/taosdata/driver-go/v3/taosWS"
+	"gorm.io/driver/mysql"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+	"gorm.io/gorm/logger"
+	"gorm.io/gorm/schema"
+)
+
+var _ gorm.Dialector = &tdengineDialector{}
+
+type tdengineDialector struct {
+	DSN string
+}
+
+func (d tdengineDialector) Name() string {
+	return "tdengine"
+}
+
+func (d tdengineDialector) Initialize(db *gorm.DB) (err error) {
+	// Initialize the TDengine connection here
+	if db.ConnPool == nil {
+		db.ConnPool, err = sql.Open("taosWS", d.DSN)
+	}
+	return
+}
+
+func (d tdengineDialector) Migrator(db *gorm.DB) gorm.Migrator {
+	// Return the TDengine migrator here
+	return &mysql.Migrator{}
+}
+
+func (d tdengineDialector) DataTypeOf(field *schema.Field) string {
+	// Return the TDengine data type for the given field
+	switch field.DataType {
+	case schema.Bool:
+		return "bool"
+	case schema.Int, schema.Uint:
+		sqlType := "bigint"
+		switch {
+		case field.Size <= 8:
+			sqlType = "tinyint"
+		case field.Size <= 16:
+			sqlType = "smallint"
+		case field.Size <= 32:
+			sqlType = "int"
+		}
+		return sqlType
+	case schema.Float:
+		if field.Size <= 32 {
+			return "float"
+		}
+		return "double"
+	case schema.String:
+		size := field.Size
+		if size == 0 {
+			size = 64
+		}
+		return fmt.Sprintf("NCHAR(%d)", size)
+	case schema.Time:
+		return "TIMESTAMP"
+	case schema.Bytes:
+		size := field.Size
+		if size == 0 {
+			size = 64
+		}
+		return fmt.Sprintf("BINARY(%d)", size)
+	}
+
+	return string(field.DataType)
+}
+
+func (d tdengineDialector) DefaultValueOf(field *schema.Field) clause.Expression {
+	return clause.Expr{SQL: "NULL"}
+}
+
+func (d tdengineDialector) BindVarTo(writer clause.Writer, stmt *gorm.Statement, v interface{}) {
+	// Bind variables for TDengine
+	switch v.(type) {
+	case string:
+		writer.WriteString("'?'")
+	default:
+		writer.WriteByte('?')
+	}
+}
+
+func (d tdengineDialector) QuoteTo(writer clause.Writer, str string) {
+	// Quote identifiers for TDengine
+	writer.WriteString(str)
+}
+
+func (d tdengineDialector) Explain(sql string, vars ...interface{}) string {
+	// Explain the SQL query for TDengine
+	return logger.ExplainSQL(sql, nil, "'", vars...)
+}
+
+func NewTDengineDialector(dsn string) gorm.Dialector {
+	return tdengineDialector{DSN: dsn}
+}

pkg/tdengine_dialector.go

@@ -1,3 +1,18 @@
+/*
+Copyright 2025 API Testing Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
 package pkg
 
 type TestCase struct {