feat: add support for BearerToken and TOTP authentication

.github/workflows/enterprise.yml

@@ -0,0 +1,89 @@
+name: enterprise.yml
+on:
+  pull_request:
+    branches:
+      - '3.0'
+      - 'main'
+  push:
+    branches:
+      - '3.0'
+      - 'main'
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ 'ubuntu-latest' ]
+        go: [ 'stable' ]
+    name: enterprise-${{ matrix.os }}-${{ matrix.go }}
+    steps:
+      - name: get branch push
+        id: get-branch
+        run: |
+          if [[ "${{ github.event_name }}" == "push" ]]; then
+            echo "branch=${{ github.ref_name }}" >> $GITHUB_OUTPUT
+          elif [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "branch=${{ github.base_ref }}" >> $GITHUB_OUTPUT
+          else
+            echo "branch=main" >> $GITHUB_OUTPUT
+          fi
+
+      - name: get TDengine
+        run: |
+          wget "${{ secrets.NIGHTLY_TDENGINE_ENTERPRISE_BASE_URL }}/tsdb-nightly-${{ steps.get-branch.outputs.branch }}.tar.gz?v=$(date +%s)" -O tsdb-nightly-${{ steps.get-branch.outputs.branch }}.tar.gz
+
+      - name: install
+        run: |
+          tar -zxf tsdb-nightly-${{ steps.get-branch.outputs.branch }}.tar.gz
+          cd tsdb-nightly-${{ steps.get-branch.outputs.branch }}
+          ls -al
+          sudo ./install.sh
+
+      - name: checkout
+        uses: actions/checkout@v4
+
+      - name: copy taos cfg
+        run: |
+          sudo mkdir -p /etc/taos
+          sudo cp ./.github/workflows/taos.cfg /etc/taos/taos.cfg
+
+      - name: shell
+        run: |
+          cat >start.sh<<EOF 
+          ulimit -n 65535 && taosd
+          EOF
+
+      - name: taosd
+        run: nohup sudo sh ./start.sh &
+
+      - name: taosadapter
+        run: nohup sudo taosadapter &
+
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ matrix.go }}
+          cache-dependency-path: go.sum
+
+      - name: Test
+        id: test
+        env:
+          GO_TEST_DISABLE_DISPLAY_LOG: "true"
+          GO_TEST_ENTERPRISE: "true"
+        run: |
+          go mod download
+          go test -coverpkg=./... -v --count=1 -coverprofile=coverage.txt -covermode=atomic ./...
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v5
+        with:
+          files: ./coverage.txt
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_ORG_TOKEN }}
+          OS: linux
+
+      - uses: actions/upload-artifact@v4
+        if: always() && (steps.test.outcome == 'failure' || steps.test.outcome == 'cancelled')
+        with:
+          name: enterprise-${{ matrix.os }}-${{ matrix.go }}-log
+          path: /var/log/taos/

af/conn_test.go

@@ -1053,6 +1053,10 @@ func TestSelectDB(t *testing.T) {
 
 func TestGetTableVGroupID(t *testing.T) {
 	db := testDatabase(t)
+	defer func() {
+		err := db.Close()
+		assert.NoError(t, err)
+	}()
 	_, err := exec(db, "create table test_vg (ts timestamp,v int)")
 	assert.NoError(t, err)
 	vgID, err := db.GetTableVGroupID("test_af", "test_vg")
@@ -1062,6 +1066,10 @@ func TestGetTableVGroupID(t *testing.T) {
 
 func TestDecimal(t *testing.T) {
 	db := testDatabase(t)
+	defer func() {
+		err := db.Close()
+		assert.NoError(t, err)
+	}()
 	_, err := exec(db, "create table test_decimal (ts timestamp, v1 decimal(10,2), v2 decimal (20,4))")
 	assert.NoError(t, err)
 	_, err = exec(db, "insert into test_decimal values (now, 10.2, 20.4)")
@@ -1095,7 +1103,15 @@ func TestDecimal(t *testing.T) {
 
 func TestTimezone(t *testing.T) {
 	db := testDatabase(t)
+	defer func() {
+		err := db.Close()
+		assert.NoError(t, err)
+	}()
 	dbParis := testDatabase(t)
+	defer func() {
+		err := dbParis.Close()
+		assert.NoError(t, err)
+	}()
 	tz := "Europe/Paris"
 	timezone, err := time.LoadLocation(tz)
 	require.NoError(t, err)

af/stmt_test.go

@@ -16,6 +16,10 @@ import (
 
 func TestNewStmt(t *testing.T) {
 	db := testDatabase(t)
+	defer func() {
+		err := db.Close()
+		assert.NoError(t, err)
+	}()
 	_, err := exec(db, "create table test_stmt (ts timestamp,v int)")
 	assert.NoError(t, err)
 	stmt := db.Stmt()
@@ -53,8 +57,6 @@ func TestNewStmt(t *testing.T) {
 	assert.NoError(t, err)
 	err = stmt.Close()
 	assert.NoError(t, err)
-	err = db.Close()
-	assert.NoError(t, err)
 }
 
 func TestStmtQueryResultWithDecimal(t *testing.T) {
@@ -122,6 +124,10 @@ func TestStmtQueryResultWithDecimal(t *testing.T) {
 func TestStmtTimezone(t *testing.T) {
 	_, is3360 := os.LookupEnv("TD_3360_TEST")
 	db := testDatabase(t)
+	defer func() {
+		err := db.Close()
+		assert.NoError(t, err)
+	}()
 	tz := "Europe/Paris"
 	timezone, err := time.LoadLocation(tz)
 	require.NoError(t, err)
@@ -171,6 +177,4 @@ func TestStmtTimezone(t *testing.T) {
 	assert.NoError(t, err)
 	err = stmt.Close()
 	assert.NoError(t, err)
-	err = db.Close()
-	assert.NoError(t, err)
 }

common/process.go

@@ -1,9 +1,12 @@
 package common
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"sync"
+
+	"github.com/taosdata/driver-go/v3/version"
 )
 
 var processName string
@@ -19,3 +22,8 @@ func GetProcessName() string {
 	})
 	return processName
 }
+
+// GetConnectorInfo returns connector information string
+func GetConnectorInfo(connectType string) string {
+	return fmt.Sprintf("go-%s-%s-%s", connectType, version.Tag, version.Commit)
+}

common/testenv/env.go

@@ -0,0 +1,12 @@
+package testenv
+
+import "os"
+
+const EnterpriseTestEnvVar = "GO_TEST_ENTERPRISE"
+
+func IsEnterpriseTest() bool {
+	if _, ok := os.LookupEnv(EnterpriseTestEnvVar); ok {
+		return true
+	}
+	return false
+}

common/testenv/env_test.go

@@ -0,0 +1,25 @@
+package testenv
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsEnterpriseTest(t *testing.T) {
+	old, exists := os.LookupEnv(EnterpriseTestEnvVar)
+	if exists {
+		defer func() {
+			_ = os.Setenv(EnterpriseTestEnvVar, old)
+		}()
+	} else {
+		defer func() {
+			_ = os.Unsetenv(EnterpriseTestEnvVar)
+		}()
+	}
+	_ = os.Setenv(EnterpriseTestEnvVar, "true")
+	assert.True(t, IsEnterpriseTest())
+	_ = os.Unsetenv(EnterpriseTestEnvVar)
+	assert.False(t, IsEnterpriseTest())
+}

common/testtool/sql.go

@@ -0,0 +1,26 @@
+package testtool
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/taosdata/driver-go/v3/common"
+)
+
+func HTTPQuery(payload string) (*common.TDEngineRestfulResp, error) {
+	body := strings.NewReader(payload)
+	req, _ := http.NewRequest(http.MethodPost, "http://127.0.0.1:6041/rest/sql", body)
+	req.Header.Set("Authorization", "Taosd /KfeAzX/f9na8qdtNZmtONryp201ma04bEl8LcvLUd7a8qdtNZmtONryp201ma04")
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		_ = resp.Body.Close()
+	}()
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("http code: %d", resp.StatusCode)
+	}
+	return common.UnmarshalRestfulBody(resp.Body, 512)
+}

common/testtool/sql_test.go

@@ -0,0 +1,13 @@
+package testtool
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHTTPQuery(t *testing.T) {
+	resp, err := HTTPQuery("select 1")
+	assert.NoError(t, err)
+	assert.Equal(t, int64(1), resp.Data[0][0])
+}

common/totp.go

@@ -0,0 +1,39 @@
+package common
+
+import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"crypto/sha256"
+	"encoding/base32"
+	"encoding/binary"
+)
+
+// GenerateTOTPCode generates a TOTP code based on the provided key, counter, and number of digits.
+func GenerateTOTPCode(key []byte, counter uint64, digits int) int {
+	h := hmac.New(sha1.New, key)
+	counterBytes := make([]byte, 8)
+	binary.BigEndian.PutUint64(counterBytes, counter)
+	h.Write(counterBytes)
+	sum := h.Sum(nil)
+	offset := sum[len(sum)-1] & 0x0F
+	v := binary.BigEndian.Uint32(sum[offset:]) & 0x7FFFFFFF
+	d := uint32(1)
+	for i := 0; i < digits && i < 8; i++ {
+		d *= 10
+	}
+	return int(v % d)
+}
+
+// GenerateTOTPSecret generates a TOTP secret using HMAC-SHA256 with the provided seed.
+func GenerateTOTPSecret(seed []byte) []byte {
+	h := hmac.New(sha256.New, nil)
+	h.Write(seed)
+	hmacResult := h.Sum(nil)
+	return hmacResult
+}
+
+// TOTPSecretStr converts the given TOTP secret bytes into a base32-encoded
+// string without padding and returns the resulting string representation.
+func TOTPSecretStr(secret []byte) string {
+	return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(secret)
+}