vstreamclient: framework for robust + simple usage

go/test/endtoend/vstreamclient/main_test.go

@@ -0,0 +1,104 @@
+/*
+Copyright 2025 The Vitess 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 vstreamclient
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"testing"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance   *cluster.LocalProcessCluster
+	vtParams          mysql.ConnParams
+	keyspaceName      = "customer"
+	vtgateGrpcAddress string
+	cell              = "zone1"
+	sqlSchema         = `create table customer(
+			id bigint not null auto_increment,
+			email varchar(128),
+			primary key(id)
+		) ENGINE=InnoDB;`
+
+	vSchema = `{
+						"tables": {
+							"customer": {}
+						}
+					}`
+)
+
+func TestMain(m *testing.M) {
+	flag.Parse()
+
+	exitcode, err := func() (int, error) {
+		clusterInstance = cluster.NewCluster(cell, "localhost")
+		clusterInstance.VtTabletExtraArgs = []string{"--health_check_interval", "1s", "--shutdown_grace_period", "3s"}
+		defer clusterInstance.Teardown()
+
+		// Start topo server
+		err := clusterInstance.StartTopo()
+		if err != nil {
+			return 1, err
+		}
+
+		// Start keyspace
+		customerKeyspace := &cluster.Keyspace{
+			Name:      "customer",
+			SchemaSQL: sqlSchema,
+			VSchema:   vSchema,
+		}
+		err = clusterInstance.StartUnshardedKeyspace(*customerKeyspace, 1, true)
+		if err != nil {
+			return 1, err
+		}
+
+		commerceKeyspace := &cluster.Keyspace{
+			Name:      "commerce",
+			SchemaSQL: "",
+			VSchema:   "",
+		}
+		err = clusterInstance.StartUnshardedKeyspace(*commerceKeyspace, 1, true)
+		if err != nil {
+			return 1, err
+		}
+
+		vtgateInstance := clusterInstance.NewVtgateInstance()
+		// Start vtgate
+		err = vtgateInstance.Setup()
+		if err != nil {
+			return 1, err
+		}
+		// ensure it is torn down during cluster TearDown
+		clusterInstance.VtgateProcess = *vtgateInstance
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+		vtgateGrpcAddress = fmt.Sprintf("%s:%d", clusterInstance.Hostname, clusterInstance.VtgateGrpcPort)
+		return m.Run(), nil
+	}()
+	if err != nil {
+		fmt.Printf("%v\n", err)
+		os.Exit(1)
+	}
+
+	os.Exit(exitcode)
+}

go/test/endtoend/vstreamclient/vstreamclient_test.go

@@ -0,0 +1,203 @@
+/*
+Copyright 2025 The Vitess 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 vstreamclient
+
+import (
+	"context"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/test/endtoend/cluster"
+	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	"vitess.io/vitess/go/vt/vstreamclient"
+)
+
+// Customer is the concrete type that will be built from the stream
+type Customer struct {
+	ID        int64     `vstream:"id"`
+	Email     string    `vstream:"email"`
+	DeletedAt time.Time `vstream:"-"`
+}
+
+func TestVStreamClient(t *testing.T) {
+	flushCount := 0
+	gotCustomers := make([]*Customer, 0)
+	tables := []vstreamclient.TableConfig{{
+		Keyspace:        "customer",
+		Table:           "customer",
+		MaxRowsPerFlush: 7,
+		DataType:        &Customer{},
+		FlushFn: func(ctx context.Context, rows []vstreamclient.Row, meta vstreamclient.FlushMeta) error {
+			flushCount++
+
+			t.Logf("upserting %d customers\n", len(rows))
+			for i, row := range rows {
+				switch {
+				// delete event
+				case row.RowChange.After == nil:
+					customer := row.Data.(*Customer)
+					customer.DeletedAt = time.Now()
+
+					gotCustomers = append(gotCustomers, customer)
+					t.Logf("deleting customer %d: %v\n", i, row)
+
+				// insert event
+				case row.RowChange.Before == nil:
+					gotCustomers = append(gotCustomers, row.Data.(*Customer))
+					t.Logf("inserting customer %d: %v\n", i, row)
+
+				// update event
+				case row.RowChange.Before != nil:
+					gotCustomers = append(gotCustomers, row.Data.(*Customer))
+					t.Logf("updating customer %d: %v\n", i, row)
+				}
+			}
+
+			// a real implementation would do something more meaningful here. For a data warehouse type workload,
+			// it would probably look like streaming rows into the data warehouse, or for more complex versions,
+			// write newline delimited json or a parquet file to object storage, then trigger a load job.
+			return nil
+		},
+	}}
+
+	ctx := context.Background()
+	conn, err := cluster.DialVTGate(ctx, t.Name(), vtgateGrpcAddress, "test_user", "")
+	require.NoError(t, err)
+	defer conn.Close()
+
+	vtgateSession := conn.Session("", nil)
+	qCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	vstreamClient, err := vstreamclient.New(ctx, "bob", conn, tables,
+		vstreamclient.WithMinFlushDuration(500*time.Millisecond),
+		vstreamclient.WithHeartbeatSeconds(1),
+		vstreamclient.WithStateTable("commerce", "vstreams"),
+		vstreamclient.WithEventFunc(func(ctx context.Context, ev *binlogdatapb.VEvent) error {
+			t.Logf("** FIELD EVENT: %v\n", ev)
+			return nil
+		}, binlogdatapb.VEventType_FIELD),
+	)
+	require.NoError(t, err)
+
+	t.Run("inserting rows", func(t *testing.T) {
+		wantCustomers := []*Customer{
+			{ID: 1, Email: "[email protected]"},
+			{ID: 2, Email: "[email protected]"},
+			{ID: 3, Email: "[email protected]"},
+			{ID: 4, Email: "[email protected]"},
+			{ID: 5, Email: "[email protected]"},
+		}
+
+		insertQuery := "insert into customer(id, email) values(:id, :email)"
+		for _, customer := range wantCustomers {
+			bindVariables := map[string]*querypb.BindVariable{
+				"id":    {Type: querypb.Type_UINT64, Value: []byte(strconv.FormatInt(int64(customer.ID), 10))},
+				"email": {Type: querypb.Type_VARCHAR, Value: []byte(customer.Email)},
+			}
+			_, err = vtgateSession.Execute(qCtx, insertQuery, bindVariables, false)
+			require.NoError(t, err)
+		}
+
+		ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		err = vstreamClient.Run(ctx)
+		if err != nil && ctx.Err() == nil {
+			t.Fatalf("failed to run vstreamclient: %v", err)
+		}
+
+		t.Logf("flush count: %d", flushCount)
+
+		assert.Positive(t, flushCount)
+		assert.ElementsMatch(t, gotCustomers, wantCustomers)
+	})
+
+	t.Run("updating rows", func(t *testing.T) {
+		// Clear out gotCustomers
+		gotCustomers = nil
+
+		updateCustomers := []*Customer{
+			{
+				ID:    1,
+				Email: "[email protected]",
+			},
+			{
+				ID:    5,
+				Email: "[email protected]",
+			},
+		}
+		updateQuery := "update customer set email=:email where id=:id"
+		for _, customer := range updateCustomers {
+			bindVariables := map[string]*querypb.BindVariable{
+				"id":    {Type: querypb.Type_UINT64, Value: []byte(strconv.FormatInt(int64(customer.ID), 10))},
+				"email": {Type: querypb.Type_VARCHAR, Value: []byte(customer.Email)},
+			}
+			_, err = vtgateSession.Execute(qCtx, updateQuery, bindVariables, false)
+			require.NoError(t, err)
+		}
+
+		ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		err = vstreamClient.Run(ctx)
+		if err != nil && ctx.Err() == nil {
+			t.Fatalf("failed to run vstreamclient: %v", err)
+		}
+
+		t.Logf("flush count: %d", flushCount)
+
+		assert.ElementsMatch(t, gotCustomers, updateCustomers)
+	})
+
+	t.Run("deleting rows", func(t *testing.T) {
+		// Clear out gotCustomers
+		gotCustomers = nil
+
+		deleteCustomerIDs := []int{1, 5}
+		deleteQuery := "delete from customer where id=:id"
+		for _, id := range deleteCustomerIDs {
+			bindVariables := map[string]*querypb.BindVariable{
+				"id": {Type: querypb.Type_UINT64, Value: []byte(strconv.FormatInt(int64(id), 10))},
+			}
+			_, err = vtgateSession.Execute(qCtx, deleteQuery, bindVariables, false)
+			require.NoError(t, err)
+		}
+
+		ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		err = vstreamClient.Run(ctx)
+		if err != nil && ctx.Err() == nil {
+			t.Fatalf("failed to run vstreamclient: %v", err)
+		}
+
+		t.Logf("flush count: %d", flushCount)
+
+		assert.Len(t, gotCustomers, len(deleteCustomerIDs))
+		// Expect non-zero `DeletedAt` field, as we setting it to
+		// time.Now() in the flushFunc.
+		for _, gotCustomer := range gotCustomers {
+			assert.NotEmpty(t, gotCustomer.DeletedAt)
+		}
+	})
+}

go/vt/vstreamclient/convert.go

@@ -0,0 +1,84 @@
+package vstreamclient
+
+import (
+	"fmt"
+	"reflect"
+	"time"
+
+	"vitess.io/vitess/go/sqltypes"
+	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+)
+
+// VStreamScanner allows for custom scan implementations
+type VStreamScanner interface {
+	VStreamScan(fields []*querypb.Field, row []sqltypes.Value, rowEvent *binlogdatapb.RowEvent, rowChange *binlogdatapb.RowChange) error
+}
+
+// copyRowToStruct builds a customer from a row event
+// TODO: this is very rudimentary mapping that only works for top-level fields
+func copyRowToStruct(shard shardConfig, row []sqltypes.Value, vPtr reflect.Value) error {
+	for fieldName, m := range shard.fieldMap {
+		structField := reflect.Indirect(vPtr).FieldByIndex(m.structIndex)
+
+		switch m.kind {
+		case reflect.Bool:
+			rowVal, err := row[m.rowIndex].ToBool()
+			if err != nil {
+				return fmt.Errorf("error converting row value to bool for field %s: %w", fieldName, err)
+			}
+			structField.SetBool(rowVal)
+
+		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+			rowVal, err := row[m.rowIndex].ToInt64()
+			if err != nil {
+				return fmt.Errorf("error converting row value to int64 for field %s: %w", fieldName, err)
+			}
+			structField.SetInt(rowVal)
+
+		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+			rowVal, err := row[m.rowIndex].ToUint64()
+			if err != nil {
+				return fmt.Errorf("error converting row value to uint64 for field %s: %w", fieldName, err)
+			}
+			structField.SetUint(rowVal)
+
+		case reflect.Float32, reflect.Float64:
+			rowVal, err := row[m.rowIndex].ToFloat64()
+			if err != nil {
+				return fmt.Errorf("error converting row value to float64 for field %s: %w", fieldName, err)
+			}
+			structField.SetFloat(rowVal)
+
+		case reflect.String:
+			rowVal := row[m.rowIndex].ToString()
+			structField.SetString(rowVal)
+
+		case reflect.Struct:
+			switch m.structType.(type) {
+			case time.Time, *time.Time:
+				rowVal, err := row[m.rowIndex].ToTime(time.UTC) // TODO: make timezone configurable
+				if err != nil {
+					return fmt.Errorf("error converting row value to time.Time for field %s: %w", fieldName, err)
+				}
+				structField.Set(reflect.ValueOf(rowVal))
+			}
+
+		case reflect.Pointer,
+			reflect.Slice,
+			reflect.Array,
+			reflect.Invalid,
+			reflect.Uintptr,
+			reflect.Complex64,
+			reflect.Complex128,
+			reflect.Chan,
+			reflect.Func,
+			reflect.Interface,
+			reflect.Map,
+			reflect.UnsafePointer:
+			return fmt.Errorf("vstreamclient: unsupported field type: %s", m.kind.String())
+		}
+	}
+
+	return nil
+}