bring back in changes

Author: jakeheft

bttest/cmv.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"os"
 	"regexp"
-	"sort"
 	"strconv"
 	"strings"
 )
@@ -23,7 +22,7 @@ type CMVConfig struct {
 	// KeyMapping defines how source key components map to CMV key components.
 	// Each entry is the 0-based index into the SPLIT result of the source key.
 	// The CMV row key is built by joining the mapped components with KeySeparator.
-	// Example: [3,4,1,2,0,5] means CMV key = source[3]#source[4]#source[1]#...
+	// Example: [3,4,1,2,0] means CMV key = source[3]#source[4]#source[1]#source[2]#source[0]
 	KeyMapping []int `json:"key_mapping"`
 	// IncludeFamilies lists the column families to carry from source to CMV.
 	// An empty list means all families are included.
@@ -60,13 +59,13 @@ func LoadCMVConfigs(path string) ([]CMVConfig, error) {
 }
 
 // ParseCMVConfigFromSQL attempts to extract a CMVConfig from a CMV SQL query.
-// This handles the common pattern used by Bigtable CMVs:
+// It parses the standard Bigtable CMV SQL pattern:
 //
-//	SELECT CAST(SPLIT(_key, '#')[SAFE_OFFSET(n)] AS STRING) AS col, ...
+//	SELECT SPLIT(_key, '#')[SAFE_OFFSET(n)] AS col, ..., _key AS source_key, family AS family
 //	FROM `table_name`
-//	ORDER BY col1, col2, ...
+//	ORDER BY col1, col2, ..., source_key
 //
-// Returns a CMVConfig with the key mapping derived from the ORDER BY + SELECT.
+// _key (optionally aliased) in ORDER BY sets AppendSourceKey = true.
 func ParseCMVConfigFromSQL(viewID, query string) (*CMVConfig, error) {
 	cfg := &CMVConfig{
 		ViewID:       viewID,
@@ -81,9 +80,6 @@ func ParseCMVConfigFromSQL(viewID, query string) (*CMVConfig, error) {
 	cfg.SourceTable = fromMatch[1]
 
 	// Extract column aliases and their SAFE_OFFSET indices from SELECT.
-	// Handles patterns like:
-	//   CAST(SPLIT(_key, '#')[SAFE_OFFSET(n)] AS STRING) AS alias
-	//   SPLIT(_key, '#')[SAFE_OFFSET(n)] AS alias
 	offsetRe := regexp.MustCompile(`SPLIT\(_key,\s*'([^']+)'\)\[SAFE_OFFSET\((\d+)\)\].*?\bAS\s+\w+\)\s+AS\s+(\w+)|SPLIT\(_key,\s*'([^']+)'\)\[SAFE_OFFSET\((\d+)\)\]\s+AS\s+(\w+)`)
 	matches := offsetRe.FindAllStringSubmatch(query, -1)
 	if len(matches) == 0 {
@@ -92,20 +88,27 @@ func ParseCMVConfigFromSQL(viewID, query string) (*CMVConfig, error) {
 
 	colMap := make(map[string]int) // alias → SAFE_OFFSET index
 	for _, m := range matches {
-		// First alternative matched: CAST(...) AS alias
 		if m[1] != "" {
+			// CAST(...) AS alias form
 			cfg.KeySeparator = m[1]
 			idx, _ := strconv.Atoi(m[2])
 			colMap[m[3]] = idx
 		} else {
-			// Second alternative matched: direct SPLIT AS alias
+			// Direct SPLIT AS alias form
 			cfg.KeySeparator = m[4]
 			idx, _ := strconv.Atoi(m[5])
 			colMap[m[6]] = idx
 		}
 	}
 
-	// Extract ORDER BY columns (excluding _key if present).
+	// Detect "_key AS <alias>" in SELECT so we can recognise the alias in ORDER BY.
+	var sourceKeyAlias string
+	keyAliasRe := regexp.MustCompile(`(?i)\b_key\s+AS\s+(\w+)`)
+	if m := keyAliasRe.FindStringSubmatch(query); m != nil {
+		sourceKeyAlias = m[1]
+	}
+
+	// Extract ORDER BY columns; _key or its alias sets AppendSourceKey.
 	orderRe := regexp.MustCompile(`(?i)ORDER\s+BY\s+(.+)$`)
 	orderMatch := orderRe.FindStringSubmatch(strings.TrimSpace(query))
 	if orderMatch == nil {
@@ -114,7 +117,7 @@ func ParseCMVConfigFromSQL(viewID, query string) (*CMVConfig, error) {
 	orderCols := strings.Split(orderMatch[1], ",")
 	for _, col := range orderCols {
 		col = strings.TrimSpace(col)
-		if col == "_key" {
+		if col == "_key" || (sourceKeyAlias != "" && col == sourceKeyAlias) {
 			cfg.AppendSourceKey = true
 			continue
 		}
@@ -129,7 +132,7 @@ func ParseCMVConfigFromSQL(viewID, query string) (*CMVConfig, error) {
 		return nil, fmt.Errorf("no key mapping columns found in ORDER BY")
 	}
 
-	// Extract included column families from SELECT (pattern: family_name AS family_name).
+	// Extract included column families from SELECT (pattern: family AS family).
 	famRe := regexp.MustCompile(`(?:,\s*)(\w+)\s+AS\s+(\w+)`)
 	famMatches := famRe.FindAllStringSubmatch(query, -1)
 	for _, m := range famMatches {
@@ -144,12 +147,9 @@ func ParseCMVConfigFromSQL(viewID, query string) (*CMVConfig, error) {
 	return cfg, nil
 }
 
-// cmvRegistry maps source table IDs to their CMV definitions.
-// The emulator uses fully-qualified table names like
-// "projects/{project}/instances/{instance}/tables/{table}", but the CMV config
-// only specifies plain table IDs. The registry matches by suffix.
+// cmvRegistry maps plain source table IDs to CMV definitions.
+// Lookups match by table ID suffix against fully-qualified table names.
 type cmvRegistry struct {
-	// configs stores raw CMV configs indexed by source table ID (plain name).
 	configs map[string][]CMVConfig
 }
 
@@ -159,13 +159,10 @@ func newCMVRegistry() *cmvRegistry {
 	}
 }
 
-// register adds a CMV definition keyed by its source table ID.
 func (r *cmvRegistry) register(cfg CMVConfig) {
 	r.configs[cfg.SourceTable] = append(r.configs[cfg.SourceTable], cfg)
 }
 
-// cmvsForTable returns all CMV instances for a fully-qualified table name.
-// Extracts the parent prefix and plain table ID from the FQ name.
 func (r *cmvRegistry) cmvsForTable(fqTable string) []*cmvInstance {
 	parent, tableID := splitFQTable(fqTable)
 	cfgs, ok := r.configs[tableID]
@@ -193,7 +190,6 @@ type cmvInstance struct {
 	parent string // e.g., projects/p/instances/i
 }
 
-// transformKey applies the CMV key mapping to produce a new row key.
 func (c *cmvInstance) transformKey(sourceKey string) string {
 	parts := strings.Split(sourceKey, c.config.KeySeparator)
 	var newParts []string
@@ -210,8 +206,7 @@ func (c *cmvInstance) transformKey(sourceKey string) string {
 	return strings.Join(newParts, c.config.KeySeparator)
 }
 
-// shouldIncludeFamily returns true if the given family name should be carried
-// to the CMV. Returns true for all families if IncludeFamilies is empty.
+// shouldIncludeFamily returns true for all families when IncludeFamilies is empty.
 func (c *cmvInstance) shouldIncludeFamily(famName string) bool {
 	if len(c.config.IncludeFamilies) == 0 {
 		return true
@@ -224,8 +219,7 @@ func (c *cmvInstance) shouldIncludeFamily(famName string) bool {
 	return false
 }
 
-// buildCMVRow creates a new row for the CMV table from a source row.
-// It copies all included families and their cells.
+// buildCMVRow builds a re-keyed CMV row by copying all included families from the source row.
 func (c *cmvInstance) buildCMVRow(sourceRow *row) *row {
 	newKey := c.transformKey(sourceRow.key)
 	cmvRow := newRow(newKey)
@@ -250,18 +244,7 @@ func (c *cmvInstance) buildCMVRow(sourceRow *row) *row {
 	return cmvRow
 }
 
-// deriveCMVKey computes what the CMV row key would be for a given source key.
-// Used for delete propagation — we need to find the CMV row to delete.
+// deriveCMVKey returns the CMV row key for a given source key.
 func (c *cmvInstance) deriveCMVKey(sourceKey string) string {
 	return c.transformKey(sourceKey)
 }
-
-// sortedCMVConfigs returns configs sorted by ViewID for deterministic behavior.
-func sortedCMVConfigs(configs []CMVConfig) []CMVConfig {
-	sorted := make([]CMVConfig, len(configs))
-	copy(sorted, configs)
-	sort.Slice(sorted, func(i, j int) bool {
-		return sorted[i].ViewID < sorted[j].ViewID
-	})
-	return sorted
-}

bttest/cmv_test.go

@@ -32,19 +32,25 @@ func newTestServerWithCMV(t *testing.T, configs []CMVConfig) *server {
 }
 
 func TestCMVTransformKey(t *testing.T) {
+	// Mirrors the production CMV config: key_mapping [3,4,1,2,0] + append_source_key.
+	//   SELECT SPLIT(_key,'#')[SAFE_OFFSET(3)] AS provider_id, ..., _key AS source_key
+	//   ORDER BY provider_id, client_id, rev_ts, event_category, bitlink_id, source_key
+	// CMV key = provider_id#client_id#rev_ts#event_category#bitlink_id#<full_source_key>
 	inst := &cmvInstance{
 		config: CMVConfig{
 			SourceTable:     "attributions_conversion_events",
 			ViewID:          "attributions_conversion_events_by_client",
 			KeySeparator:    "#",
-			KeyMapping:      []int{3, 4, 1, 2, 0, 5},
+			KeyMapping:      []int{3, 4, 1, 2, 0},
 			AppendSourceKey: true,
 		},
 	}
 
 	sourceKey := "bitlink123#9999999999#order#shopify#myshop.myshopify.com#01HWXYZ"
 	got := inst.transformKey(sourceKey)
-	want := "shopify#myshop.myshopify.com#9999999999#order#bitlink123#01HWXYZ#bitlink123#9999999999#order#shopify#myshop.myshopify.com#01HWXYZ"
+	// parts: [0]bitlink123 [1]9999999999 [2]order [3]shopify [4]myshop.myshopify.com [5]01HWXYZ
+	// mapped: parts[3]#parts[4]#parts[1]#parts[2]#parts[0] + full source key
+	want := "shopify#myshop.myshopify.com#9999999999#order#bitlink123#bitlink123#9999999999#order#shopify#myshop.myshopify.com#01HWXYZ"
 	assert.Equal(t, want, got)
 }
 
@@ -261,7 +267,7 @@ func TestLoadCMVConfigs(t *testing.T) {
 			"source_table": "attributions_conversion_events",
 			"view_id": "attributions_conversion_events_by_client",
 			"key_separator": "#",
-			"key_mapping": [3, 4, 1, 2, 0, 5],
+			"key_mapping": [3, 4, 1, 2, 0],
 			"include_families": ["cr", "d", "m"],
 			"append_source_key": true
 		}
@@ -276,34 +282,100 @@ func TestLoadCMVConfigs(t *testing.T) {
 	assert.Equal(t, "attributions_conversion_events", configs[0].SourceTable)
 	assert.Equal(t, "attributions_conversion_events_by_client", configs[0].ViewID)
 	assert.Equal(t, "#", configs[0].KeySeparator)
-	assert.Equal(t, []int{3, 4, 1, 2, 0, 5}, configs[0].KeyMapping)
+	assert.Equal(t, []int{3, 4, 1, 2, 0}, configs[0].KeyMapping)
 	assert.Equal(t, []string{"cr", "d", "m"}, configs[0].IncludeFamilies)
 	assert.True(t, configs[0].AppendSourceKey)
 }
 
 func TestParseCMVConfigFromSQL(t *testing.T) {
+	// Mirrors the production Terraform CMV SQL exactly: no CAST, _key aliased as
+	// source_key, attribution_source_id is embedded in source_key not a separate component.
 	query := `SELECT
-  CAST(SPLIT(_key, '#')[SAFE_OFFSET(3)] AS STRING) AS provider_id,
-  CAST(SPLIT(_key, '#')[SAFE_OFFSET(4)] AS STRING) AS client_id,
-  CAST(SPLIT(_key, '#')[SAFE_OFFSET(1)] AS STRING) AS rev_ts,
-  CAST(SPLIT(_key, '#')[SAFE_OFFSET(2)] AS STRING) AS conversion_type,
-  CAST(SPLIT(_key, '#')[SAFE_OFFSET(0)] AS STRING) AS bitlink_id,
-  CAST(SPLIT(_key, '#')[SAFE_OFFSET(5)] AS STRING) AS attribution_source_id,
-  _key,
+  SPLIT(_key, '#')[SAFE_OFFSET(3)] AS provider_id,
+  SPLIT(_key, '#')[SAFE_OFFSET(4)] AS client_id,
+  SPLIT(_key, '#')[SAFE_OFFSET(1)] AS rev_ts,
+  SPLIT(_key, '#')[SAFE_OFFSET(2)] AS event_category,
+  SPLIT(_key, '#')[SAFE_OFFSET(0)] AS bitlink_id,
+  _key AS source_key,
   cr AS cr,
   d AS d,
   m AS m
 FROM ` + "`attributions_conversion_events`" + `
-ORDER BY provider_id, client_id, rev_ts, conversion_type, bitlink_id, attribution_source_id, _key`
+ORDER BY provider_id, client_id, rev_ts, event_category, bitlink_id, source_key`
 
 	cfg, err := ParseCMVConfigFromSQL("attributions_conversion_events_by_client", query)
 	require.NoError(t, err)
 	assert.Equal(t, "attributions_conversion_events", cfg.SourceTable)
 	assert.Equal(t, "attributions_conversion_events_by_client", cfg.ViewID)
 	assert.Equal(t, "#", cfg.KeySeparator)
-	assert.Equal(t, []int{3, 4, 1, 2, 0, 5}, cfg.KeyMapping)
+	assert.Equal(t, []int{3, 4, 1, 2, 0}, cfg.KeyMapping)
 	assert.True(t, cfg.AppendSourceKey)
 	assert.Contains(t, cfg.IncludeFamilies, "cr")
 	assert.Contains(t, cfg.IncludeFamilies, "d")
 	assert.Contains(t, cfg.IncludeFamilies, "m")
 }
+
+func TestCMVDropRowRangePrefix(t *testing.T) {
+	ctx := context.Background()
+	// Key mapping [1,0]: CMV key = source[1]#source[0] (swap two components).
+	configs := []CMVConfig{{
+		SourceTable:  "src_table",
+		ViewID:       "src_cmv",
+		KeySeparator: "#",
+		KeyMapping:   []int{1, 0},
+	}}
+	s := newTestServerWithCMV(t, configs)
+
+	parent := "projects/test/instances/test"
+	fqSrc := parent + "/tables/src_table"
+	fqCMV := parent + "/tables/src_cmv"
+
+	_, err := s.CreateTable(ctx, &btapb.CreateTableRequest{
+		Parent:  parent,
+		TableId: "src_table",
+		Table: &btapb.Table{
+			ColumnFamilies: map[string]*btapb.ColumnFamily{
+				"cf1": {},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	// Write three rows: two share the prefix "alpha#" and one does not.
+	for _, key := range []string{"alpha#one", "alpha#two", "beta#three"} {
+		_, err = s.MutateRow(ctx, &btpb.MutateRowRequest{
+			TableName: fqSrc,
+			RowKey:    []byte(key),
+			Mutations: []*btpb.Mutation{{
+				Mutation: &btpb.Mutation_SetCell_{
+					SetCell: &btpb.Mutation_SetCell{
+						FamilyName:      "cf1",
+						ColumnQualifier: []byte("c"),
+						TimestampMicros: 1000,
+						Value:           []byte("v"),
+					},
+				},
+			}},
+		})
+		require.NoError(t, err)
+	}
+
+	s.mu.Lock()
+	cmvTbl := s.tables[fqCMV]
+	s.mu.Unlock()
+	require.NotNil(t, cmvTbl)
+	assert.Equal(t, 3, cmvTbl.rows.Len())
+
+	// Drop source rows with prefix "alpha#".
+	_, err = s.DropRowRange(ctx, &btapb.DropRowRangeRequest{
+		Name:   fqSrc,
+		Target: &btapb.DropRowRangeRequest_RowKeyPrefix{RowKeyPrefix: []byte("alpha#")},
+	})
+	require.NoError(t, err)
+
+	// CMV should now have only 1 row: "three#beta" (from "beta#three").
+	assert.Equal(t, 1, cmvTbl.rows.Len())
+	// The remaining CMV row should be the transformed "beta#three" → "three#beta".
+	cmvRow := cmvTbl.rows.Get(btreeKey("three#beta"))
+	assert.NotNil(t, cmvRow, "CMV row for non-deleted source should still exist")
+}