vtctl: Add interface for atomic VSchema operations

go/vt/topo/vschema.go

@@ -327,3 +327,43 @@ func (ts *Server) SaveMirrorRules(ctx context.Context, mirrorRules *vschemapb.Mi
 	_, err = ts.globalCell.Update(ctx, MirrorRulesFile, data, nil)
 	return err
 }
+
+type VSchemaInfo struct {
+	keyspace string
+	version  Version
+	*vschemapb.Keyspace
+}
+
+func (ts *Server) GetVSchema(ctx context.Context, keyspace string) (*VSchemaInfo, error) {
+	ksvs, err := ts.GetVSchema(ctx, keyspace)
+	if err != nil {
+		return nil, err
+	}
+	return &VSchemaInfo{
+		keyspace: ksvs.Name,
+		version:  ksvs.version,
+		Keyspace: ksvs.Keyspace,
+	}, nil
+}
+
+func (ts *Server) SaveVSchema(ctx context.Context, keyspace string, vschema *VSchemaInfo) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+
+	nodePath := path.Join(KeyspacesPath, keyspace, VSchemaFile)
+	data, err := vschema.MarshalVT()
+	if err != nil {
+		return err
+	}
+
+	version, err := ts.globalCell.Update(ctx, nodePath, data, vschema.version)
+	if err != nil {
+		log.Errorf("failed to update vschema for keyspace %s: %v", keyspace, err)
+		return err
+	}
+	vschema.version = version
+	log.Infof("successfully updated vschema for keyspace %s: %+v", keyspace, vschema.Keyspace)
+
+	return nil
+}

go/vt/vtctl/vtctlclient/client_test.go

@@ -0,0 +1,14 @@
+package vtctlclient
+
+import (
+	"testing"
+)
+
+func TestVSchemaAtomicOperations(t *testing.T) {
+	// Test cases for:
+	// - Concurrent modifications
+	// - Version conflicts
+	// - Successful atomic updates
+	// - Invalid inputs
+	// - Error conditions
+} 

go/vt/vtctl/vtctlclient/grpc_client.go

@@ -0,0 +1,34 @@
+package vtctlclient
+
+import (
+	"context"
+
+	"github.com/your-project/vschemapb"
+)
+
+type gRPCVSchemaOperations struct {
+	client *gRPCClient
+	keyspace string
+}
+
+func (ops *gRPCVSchemaOperations) AddTable(ctx context.Context, keyspace, tableName string, tableSpec interface{}) error {
+	// 1. Get current VSchema with version
+	vschema, version, err := ops.client.GetVSchema(ctx, keyspace)
+	if err != nil {
+		return err
+	}
+
+	// 2. Validate table spec
+	spec, err := validateTableSpec(tableSpec)
+	if err != nil {
+		return err
+	}
+
+	// 3. Modify VSchema
+	// Add table to vschema.Tables map
+
+	// 4. Save with version check
+	return ops.client.SaveVSchema(ctx, keyspace, vschema, version)
+}
+
+// Implement other VSchemaOperations methods similarly 

go/vt/vtctl/vtctlclient/interface.go

@@ -55,6 +55,35 @@ type VtctlClient interface {
 	Close()
 }
 
+
+// VSchemaOperations defines the interface for atomic VSchema modifications
+type VSchemaOperations interface {
+	// AddTable atomically adds a table to the VSchema
+	AddTable(ctx context.Context, keyspace, tableName string, tableSpec interface{}) error
+
+	// RemoveTable atomically removes a table from the VSchema
+	RemoveTable(ctx context.Context, keyspace, tableName string) error
+
+	// AddVindex atomically adds a vindex to the VSchema
+	AddVindex(ctx context.Context, keyspace, vindexName string, vindexSpec interface{}) error
+
+	// RemoveVindex atomically removes a vindex from the VSchema
+	RemoveVindex(ctx context.Context, keyspace, vindexName string) error
+}
+
+// VtctlClient defines the interface used to send remote vtctl commands
+type VtctlClient interface {
+	// ExecuteVtctlCommand will execute the command remotely
+	ExecuteVtctlCommand(ctx context.Context, args []string, actionTimeout time.Duration) (logutil.EventStream, error)
+
+	// VSchema returns the VSchema operations interface
+	VSchema() VSchemaOperations
+
+	// Close will terminate the connection. This object won't be
+	// used after this.
+	Close()
+}
+
 // Factory functions are registered by client implementations
 type Factory func(ctx context.Context, addr string) (VtctlClient, error)
 

go/vt/vtctl/vtctlclient/vschema_test.go

@@ -0,0 +1,205 @@
+package vtctlclient
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/topo/memorytopo"
+	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
+)
+
+func TestVSchemaAtomicOperations(t *testing.T) {
+	ctx := context.Background()
+	ts := memorytopo.NewServer("")
+	client := newTestClient(ts)
+
+	tests := []struct {
+		name        string
+		setup       func(t *testing.T)
+		operation   func(t *testing.T) error
+		concurrent  func(t *testing.T)
+		wantErr     bool
+		errContains string
+	}{
+		{
+			name: "successful_add_table",
+			setup: func(t *testing.T) {
+				// Initialize empty vschema
+				initEmptyVSchema(t, ts, "test_keyspace")
+			},
+			operation: func(t *testing.T) error {
+				return client.VSchema().AddTable(ctx, "test_keyspace", "test_table", &vschemapb.Table{})
+			},
+			wantErr: false,
+		},
+		{
+			name: "concurrent_modification",
+			setup: func(t *testing.T) {
+				initEmptyVSchema(t, ts, "test_keyspace")
+			},
+			operation: func(t *testing.T) error {
+				return client.VSchema().AddTable(ctx, "test_keyspace", "test_table", &vschemapb.Table{})
+			},
+			concurrent: func(t *testing.T) {
+				// Simulate concurrent modification
+				client2 := newTestClient(ts)
+				err := client2.VSchema().AddTable(ctx, "test_keyspace", "other_table", &vschemapb.Table{})
+				require.NoError(t, err)
+			},
+			wantErr:     true,
+			errContains: "version mismatch",
+		},
+		{
+			name: "retry_success",
+			setup: func(t *testing.T) {
+				initEmptyVSchema(t, ts, "test_keyspace")
+			},
+			operation: func(t *testing.T) error {
+				return client.VSchema().AddTable(ctx, "test_keyspace", "test_table", &vschemapb.Table{})
+			},
+			concurrent: func(t *testing.T) {
+				// Simulate temporary interference
+				time.Sleep(10 * time.Millisecond)
+				client2 := newTestClient(ts)
+				err := client2.VSchema().AddTable(ctx, "test_keyspace", "other_table", &vschemapb.Table{})
+				require.NoError(t, err)
+			},
+			wantErr: false, // Should succeed after retry
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.setup(t)
+
+			if tt.concurrent != nil {
+				go tt.concurrent(t)
+			}
+
+			err := tt.operation(t)
+
+			if tt.wantErr {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.errContains)
+			} else {
+				require.NoError(t, err)
+				// Verify final state
+				vschema, err := ts.GetVSchema(ctx, "test_keyspace")
+				require.NoError(t, err)
+				assert.NotNil(t, vschema)
+			}
+		})
+	}
+}
+
+func TestVSchemaVersionConflicts(t *testing.T) {
+	ctx := context.Background()
+	ts := memorytopo.NewServer("")
+	client := newTestClient(ts)
+
+	// Test multiple concurrent modifications
+	t.Run("multiple_concurrent_modifications", func(t *testing.T) {
+		initEmptyVSchema(t, ts, "test_keyspace")
+
+		// Launch multiple goroutines attempting modifications
+		const numGoroutines = 5
+		errChan := make(chan error, numGoroutines)
+
+		for i := 0; i < numGoroutines; i++ {
+			go func(idx int) {
+				err := client.VSchema().AddTable(ctx, "test_keyspace", 
+					fmt.Sprintf("table_%d", idx), &vschemapb.Table{})
+				errChan <- err
+			}(i)
+		}
+
+		// Collect results
+		successCount := 0
+		for i := 0; i < numGoroutines; i++ {
+			err := <-errChan
+			if err == nil {
+				successCount++
+			}
+		}
+
+		// At least one should succeed, others should fail with version conflict
+		assert.True(t, successCount >= 1)
+		assert.True(t, successCount < numGoroutines)
+	})
+}
+
+func TestVSchemaErrorConditions(t *testing.T) {
+	ctx := context.Background()
+	ts := memorytopo.NewServer("")
+	client := newTestClient(ts)
+
+	tests := []struct {
+		name        string
+		operation   func() error
+		errContains string
+	}{
+		{
+			name: "invalid_keyspace",
+			operation: func() error {
+				return client.VSchema().AddTable(ctx, "", "table", &vschemapb.Table{})
+			},
+			errContains: "invalid keyspace",
+		},
+		{
+			name: "invalid_table_name",
+			operation: func() error {
+				return client.VSchema().AddTable(ctx, "keyspace", "", &vschemapb.Table{})
+			},
+			errContains: "invalid table name",
+		},
+		{
+			name: "nil_table_spec",
+			operation: func() error {
+				return client.VSchema().AddTable(ctx, "keyspace", "table", nil)
+			},
+			errContains: "nil table spec",
+		},
+		{
+			name: "keyspace_not_found",
+			operation: func() error {
+				return client.VSchema().AddTable(ctx, "nonexistent", "table", &vschemapb.Table{})
+			},
+			errContains: "keyspace not found",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := tt.operation()
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), tt.errContains)
+		})
+	}
+}
+
+// Helper functions
+
+func initEmptyVSchema(t *testing.T, ts *memorytopo.Server, keyspace string) {
+	ctx := context.Background()
+	err := ts.CreateKeyspace(ctx, keyspace, &vschemapb.Keyspace{})
+	require.NoError(t, err)
+}
+
+func newTestClient(ts *memorytopo.Server) VtctlClient {
+	// Create a test client implementation
+	return &testVtctlClient{
+		ts: ts,
+	}
+}
+
+type testVtctlClient struct {
+	ts *memorytopo.Server
+}
+
+// Implement VtctlClient interface methods... 

go/vt/vtctl/vtctlserver/server.go

@@ -0,0 +1,16 @@
+package vtctlserver
+
+import (
+	"context"
+
+	"github.com/your-project/vtctlservicepb"
+)
+
+type VtctlServer struct {
+	// Add any necessary fields here
+}
+
+func (s *VtctlServer) AddVSchemaTable(ctx context.Context, req *vtctlservicepb.AddVSchemaTableRequest) (*vtctlservicepb.AddVSchemaTableResponse, error) {
+	// Implement server-side handling of atomic VSchema modifications
+	return nil, nil // Placeholder return, actual implementation needed
+}