Add victoriametrics database driver

Author: ljluestc

database/victoria/README.md

@@ -0,0 +1,5 @@
+# VictoriaMetrics
+
+This driver enables golang-migrate to work with [VictoriaMetrics](https://victoriametrics.com/), a high-performance time series database.
+
+## Usage

database/victoria/example/migration_example.down.json

@@ -0,0 +1,2 @@
+-- This is just a placeholder as VictoriaMetrics doesn't have a way to revert imports
+-- Instead, you would typically filter by time range or labels to exclude the data

database/victoria/example/migration_example.up.json

@@ -0,0 +1,3 @@
+{"metric":{"__name__":"up","job":"migrate_test"},"values":[1],"timestamps":[1596698684000]}
+{"metric":{"__name__":"cpu_usage","instance":"server1","job":"migrate_test"},"values":[0.45,0.52,0.48],"timestamps":[1596698684000,1596698694000,1596698704000]}
+{"metric":{"__name__":"memory_usage","instance":"server1","job":"migrate_test"},"values":[0.25,0.28,0.22],"timestamps":[1596698684000,1596698694000,1596698704000]}

database/victoria/export.go

@@ -0,0 +1,59 @@
+package victoria
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+)
+
+// Export reads data from VictoriaMetrics based on label filters and time range
+func (v *Victoria) Export(ctx context.Context, w io.Writer) error {
+	if !v.isOpen {
+		return ErrClosed
+	}
+
+	// Build query parameters
+	query := url.Values{}
+	if v.config.LabelFilter != "" {
+		query.Set("match[]", v.config.LabelFilter)
+	}
+	if v.config.StartTime != "" {
+		query.Set("start", v.config.StartTime)
+	}
+	if v.config.EndTime != "" {
+		query.Set("end", v.config.EndTime)
+	}
+
+	// Create request with context
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, v.exportURL+"?"+query.Encode(), nil)
+	if err != nil {
+		return fmt.Errorf("failed to create export request: %w", err)
+	}
+
+	// Set headers
+	req.Header.Set("Accept", "application/json")
+	req.Header.Set("Accept-Encoding", "gzip")
+
+	// Execute request
+	resp, err := v.client.Do(req)
+	if err != nil {
+		return fmt.Errorf("failed to execute export request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	// Check response status
+	if resp.StatusCode != http.StatusOK {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("export failed with status %d: %s", resp.StatusCode, string(bodyBytes))
+	}
+
+	// Copy response body to writer
+	_, err = io.Copy(w, resp.Body)
+	if err != nil {
+		return fmt.Errorf("failed to read export data: %w", err)
+	}
+
+	return nil
+}

database/victoria/export_test.go

@@ -0,0 +1,82 @@
+package victoria
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestVictoriaExport(t *testing.T) {
+	// Sample export data
+	exportData := `{"metric":{"__name__":"up","job":"test"},"values":[1],"timestamps":[1596698684000]}
+{"metric":{"__name__":"cpu_usage","instance":"server1"},"values":[0.45],"timestamps":[1596698684000]}`
+
+	// Setup test server
+	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/api/v1/export" {
+			// Check query parameters
+			query := r.URL.Query()
+			assert.Equal(t, "{__name__=\"up\"}", query.Get("match[]"))
+			assert.Equal(t, "2020-01-01T00:00:00Z", query.Get("start"))
+			assert.Equal(t, "2020-01-02T00:00:00Z", query.Get("end"))
+
+			// Return export data
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte(exportData))
+		} else {
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer testServer.Close()
+
+	// Parse the URL for our test server
+	serverURL := strings.TrimPrefix(testServer.URL, "http://")
+
+	// Create driver
+	d := &Victoria{}
+	dsn := "victoria://" + serverURL + "?label_filter={__name__=\"up\"}&start=2020-01-01T00:00:00Z&end=2020-01-02T00:00:00Z"
+	// No need to store the returned driver since we're testing the receiver methods directly
+	_, err := d.Open(dsn)
+	assert.NoError(t, err)
+
+	// Test export
+	var buf bytes.Buffer
+	err = d.Export(context.Background(), &buf)
+	assert.NoError(t, err)
+	assert.Equal(t, exportData, buf.String())
+
+	// Test export with closed connection
+	d.Close()
+	err = d.Export(context.Background(), &buf)
+	assert.Error(t, err)
+	assert.Equal(t, ErrClosed, err)
+}
+
+func TestVictoriaExportError(t *testing.T) {
+	// Setup test server that returns an error
+	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("Internal server error"))
+	}))
+	defer testServer.Close()
+
+	// Parse the URL for our test server
+	serverURL := strings.TrimPrefix(testServer.URL, "http://")
+
+	// Create driver
+	d := &Victoria{}
+	dsn := "victoria://" + serverURL
+	_, err := d.Open(dsn)
+	assert.NoError(t, err)
+
+	// Test export
+	var buf bytes.Buffer
+	err = d.Export(context.Background(), &buf)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "status 500")
+}

database/victoria/victoria.go

@@ -3,7 +3,6 @@ package victoria
 import (
 	"bufio"
 	"bytes"
-	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -15,6 +14,12 @@ import (
 	"github.com/golang-migrate/migrate/v4/database"
 )
 
+// Define error constants if they don't exist in the database package
+var (
+	ErrLocked = errors.New("database is locked")
+	ErrClosed = errors.New("database connection is closed")
+)
+
 func init() {
 	database.Register("victoria", &Victoria{})
 }
@@ -110,7 +115,7 @@ func (v *Victoria) Close() error {
 // Lock acquires a database lock (no-op for VictoriaMetrics)
 func (v *Victoria) Lock() error {
 	if !v.isOpen {
-		return database.ErrLocked
+		return ErrLocked
 	}
 	v.isLocked = true
 	return nil
@@ -119,7 +124,7 @@ func (v *Victoria) Lock() error {
 // Unlock releases a database lock (no-op for VictoriaMetrics)
 func (v *Victoria) Unlock() error {
 	if !v.isOpen {
-		return database.ErrLocked
+		return ErrLocked
 	}
 	v.isLocked = false
 	return nil
@@ -128,16 +133,16 @@ func (v *Victoria) Unlock() error {
 // Run executes a migration by importing data into VictoriaMetrics
 func (v *Victoria) Run(migration io.Reader) error {
 	if !v.isOpen {
-		return database.ErrClosed
+		return ErrClosed
 	}
 
 	if !v.isLocked {
-		return database.ErrLocked
+		return ErrLocked
 	}
 
 	// Buffer to collect migration data
 	var migrationBuffer bytes.Buffer
-	
+
 	// Read migration content
 	scanner := bufio.NewScanner(migration)
 	scanner.Buffer(make([]byte, 4*1024*1024), 4*1024*1024) // 4MB buffer
@@ -199,4 +204,3 @@ func (v *Victoria) Drop() error {
 
 // Ensure Victoria implements the database.Driver interface
 var _ database.Driver = (*Victoria)(nil)
-var _ database.Locker = (*Victoria)(nil)