Add OrderByColumn

Author: ilyario

api/v1/dataflow_types.go

@@ -291,6 +291,11 @@ type PostgreSQLSourceSpec struct {
 	// +kubebuilder:default="updated_at"
 	ChangeTrackingColumn string `json:"changeTrackingColumn,omitempty"`
 
+	// OrderByColumn is the secondary sort key for stable pagination (default: id).
+	// Used in ORDER BY together with changeTrackingColumn.
+	// +optional
+	OrderByColumn string `json:"orderByColumn,omitempty"`
+
 	// AutoCreateTable creates the table if it doesn't exist before reading
 	// +optional
 	AutoCreateTable *bool `json:"autoCreateTable,omitempty"`
@@ -337,6 +342,10 @@ type TrinoSourceSpec struct {
 	// TableSecretRef references a Kubernetes secret for table name
 	// +optional
 	TableSecretRef *SecretRef `json:"tableSecretRef,omitempty"`
+
+	// OrderByColumn is the column used for incremental pagination and stable ORDER BY (default: id).
+	// +optional
+	OrderByColumn string `json:"orderByColumn,omitempty"`
 }
 
 // ClickHouseSourceSpec defines ClickHouse source configuration
@@ -362,6 +371,10 @@ type ClickHouseSourceSpec struct {
 	// TableSecretRef references a Kubernetes secret for table name
 	// +optional
 	TableSecretRef *SecretRef `json:"tableSecretRef,omitempty"`
+
+	// OrderByColumn is the column used for incremental pagination and stable ORDER BY (default: id).
+	// +optional
+	OrderByColumn string `json:"orderByColumn,omitempty"`
 }
 
 // NessieSourceSpec defines Nessie (Iceberg REST catalog) source configuration.

api/v1/dataflow_validation.go

@@ -19,6 +19,7 @@ package v1
 import (
 	"encoding/json"
 	"fmt"
+	"regexp"
 	"strconv"
 	"strings"
 
@@ -197,9 +198,25 @@ func validatePostgreSQLSource(p *PostgreSQLSourceSpec, f *field.Path) field.Erro
 	if p.TableSecretRef != nil {
 		all = append(all, validateSecretRef(p.TableSecretRef, f.Child("tableSecretRef"))...)
 	}
+	if err := validateOrderByColumn(p.OrderByColumn, f.Child("orderByColumn")); err != nil {
+		all = append(all, err)
+	}
 	return all
 }
 
+var sqlIdentifierRe = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)
+
+func validateOrderByColumn(col string, f *field.Path) *field.Error {
+	if col == "" {
+		return nil
+	}
+	if !sqlIdentifierRe.MatchString(col) {
+		return field.Invalid(f, col,
+			"must be a valid SQL identifier (letters, digits, underscore; must not start with a digit)")
+	}
+	return nil
+}
+
 func validateTrinoSource(t *TrinoSourceSpec, f *field.Path) field.ErrorList {
 	var all field.ErrorList
 	hasURL := t.ServerURL != "" || t.ServerURLSecretRef != nil
@@ -230,6 +247,9 @@ func validateTrinoSource(t *TrinoSourceSpec, f *field.Path) field.ErrorList {
 	if t.TableSecretRef != nil {
 		all = append(all, validateSecretRef(t.TableSecretRef, f.Child("tableSecretRef"))...)
 	}
+	if err := validateOrderByColumn(t.OrderByColumn, f.Child("orderByColumn")); err != nil {
+		all = append(all, err)
+	}
 	return all
 }
 
@@ -503,6 +523,9 @@ func validateClickHouseSource(c *ClickHouseSourceSpec, f *field.Path) field.Erro
 	if c.TableSecretRef != nil {
 		all = append(all, validateSecretRef(c.TableSecretRef, f.Child("tableSecretRef"))...)
 	}
+	if err := validateOrderByColumn(c.OrderByColumn, f.Child("orderByColumn")); err != nil {
+		all = append(all, err)
+	}
 	return all
 }
 

api/v1/dataflow_validation_orderby_test.go

@@ -0,0 +1,63 @@
+/*
+Copyright 2024.
+
+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 v1
+
+import (
+	"testing"
+
+	"k8s.io/apimachinery/pkg/util/validation/field"
+)
+
+func TestValidateOrderByColumn_sources(t *testing.T) {
+	path := field.NewPath("spec").Child("source").Child("config")
+
+	t.Run("postgresql valid", func(t *testing.T) {
+		spec := &PostgreSQLSourceSpec{
+			ConnectionString: "postgres://x",
+			Table:            "t",
+			OrderByColumn:    "price_id",
+		}
+		if errs := validatePostgreSQLSource(spec, path); len(errs) != 0 {
+			t.Fatalf("expected no errors, got %v", errs)
+		}
+	})
+
+	t.Run("trino invalid", func(t *testing.T) {
+		spec := &TrinoSourceSpec{
+			ServerURL:     "http://trino:8080",
+			Catalog:       "c",
+			Schema:        "s",
+			Table:         "t",
+			OrderByColumn: "bad-col",
+		}
+		errs := validateTrinoSource(spec, path)
+		if len(errs) == 0 {
+			t.Fatal("expected validation error")
+		}
+	})
+
+	t.Run("clickhouse valid", func(t *testing.T) {
+		spec := &ClickHouseSourceSpec{
+			ConnectionString: "clickhouse://x",
+			Table:            "t",
+			OrderByColumn:    "price_id",
+		}
+		if errs := validateClickHouseSource(spec, path); len(errs) != 0 {
+			t.Fatalf("expected no errors, got %v", errs)
+		}
+	})
+}

api/v1/dataflow_validation_postgresql_test.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2024.
+
+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 v1
+
+import (
+	"testing"
+
+	"k8s.io/apimachinery/pkg/util/validation/field"
+)
+
+func TestValidatePostgreSQLSource_orderByColumn(t *testing.T) {
+	path := field.NewPath("spec").Child("source").Child("config")
+
+	t.Run("valid identifier", func(t *testing.T) {
+		spec := &PostgreSQLSourceSpec{
+			ConnectionString: "postgres://x",
+			Table:            "t",
+			OrderByColumn:    "price_id",
+		}
+		if errs := validatePostgreSQLSource(spec, path); len(errs) != 0 {
+			t.Fatalf("expected no errors, got %v", errs)
+		}
+	})
+
+	t.Run("invalid identifier", func(t *testing.T) {
+		spec := &PostgreSQLSourceSpec{
+			ConnectionString: "postgres://x",
+			Table:            "t",
+			OrderByColumn:    "price-id",
+		}
+		errs := validatePostgreSQLSource(spec, path)
+		if len(errs) == 0 {
+			t.Fatal("expected validation error for invalid orderByColumn")
+		}
+	})
+}

internal/connectors/clickhouse.go

@@ -178,16 +178,9 @@ func (c *ClickHouseSourceConnector) readRows(ctx context.Context, msgChan chan *
 
 	var query string
 	if c.config.Query != "" {
-		query = c.config.Query
+		query = c.wrapQueryWithStableOrder(c.config.Query)
 	} else {
-		if lastReadID > 0 {
-			query = fmt.Sprintf("SELECT * FROM %s WHERE id > %d ORDER BY id", c.config.Table, lastReadID)
-		} else if lastReadTime != nil {
-			query = fmt.Sprintf("SELECT * FROM %s WHERE created_at > '%s' ORDER BY created_at",
-				c.config.Table, lastReadTime.Format("2006-01-02 15:04:05"))
-		} else {
-			query = fmt.Sprintf("SELECT * FROM %s", c.config.Table)
-		}
+		query = c.buildTableReadQuery(lastReadID, lastReadTime)
 	}
 
 	c.logger.V(1).Info("Executing ClickHouse query", "query", query, "table", c.config.Table)
@@ -209,16 +202,9 @@ func (c *ClickHouseSourceConnector) readRows(ctx context.Context, msgChan chan *
 		return fmt.Errorf("clickhouse columns: %w", err)
 	}
 
-	var idIndex = -1
-	var createdAtIndex = -1
-	for i, col := range columns {
-		if col == "id" {
-			idIndex = i
-		}
-		if col == "created_at" {
-			createdAtIndex = i
-		}
-	}
+	orderByCol := ResolveOrderByColumn(c.config.OrderByColumn)
+	idIndex := ColumnIndex(columns, orderByCol)
+	createdAtIndex := ColumnIndex(columns, "created_at")
 
 	var maxReadID int64 = lastReadID
 	var maxReadTime *time.Time
@@ -278,7 +264,7 @@ func (c *ClickHouseSourceConnector) readRows(ctx context.Context, msgChan chan *
 		msg := types.NewMessage(jsonData)
 		msg.Metadata["table"] = c.config.Table
 		if idIndex >= 0 && len(values) > idIndex {
-			msg.Metadata["id"] = values[idIndex]
+			SetSourceRowIDMetadata(msg, values[idIndex])
 		}
 		// Ack advances checkpoint only after sink successfully writes; prevents data loss on crash
 		rowID, rowTime := c.extractRowCheckpoint(values, idIndex, createdAtIndex)
@@ -299,6 +285,24 @@ func (c *ClickHouseSourceConnector) readRows(ctx context.Context, msgChan chan *
 	return nil
 }
 
+func (c *ClickHouseSourceConnector) buildTableReadQuery(lastReadID int64, lastReadTime *time.Time) string {
+	table := c.config.Table
+	col := ResolveOrderByColumn(c.config.OrderByColumn)
+	if lastReadID > 0 {
+		return fmt.Sprintf("SELECT * FROM %s WHERE %s > %d ORDER BY %s", table, col, lastReadID, col)
+	}
+	if lastReadTime != nil {
+		return fmt.Sprintf("SELECT * FROM %s WHERE created_at > '%s' ORDER BY created_at, %s",
+			table, lastReadTime.Format("2006-01-02 15:04:05"), col)
+	}
+	return fmt.Sprintf("SELECT * FROM %s ORDER BY created_at, %s", table, col)
+}
+
+func (c *ClickHouseSourceConnector) wrapQueryWithStableOrder(userQuery string) string {
+	col := ResolveOrderByColumn(c.config.OrderByColumn)
+	return WrapQueryStableOrder(userQuery, col)
+}
+
 // extractRowCheckpoint returns (id, created_at) for checkpoint advancement.
 func (c *ClickHouseSourceConnector) extractRowCheckpoint(values []interface{}, idIndex, createdAtIndex int) (int64, *time.Time) {
 	var rowID int64

internal/connectors/clickhouse_test.go

@@ -254,3 +254,50 @@ func TestBuildInsertValues_preservesCreatedAtFromSource(t *testing.T) {
 		})
 	}
 }
+
+func TestClickHouseSourceConnector_buildTableReadQuery(t *testing.T) {
+	spec := &v1.ClickHouseSourceSpec{
+		ConnectionString: "clickhouse://localhost:9000",
+		Table:            "events",
+	}
+	c := NewClickHouseSourceConnector(spec)
+
+	t.Run("first read", func(t *testing.T) {
+		got := c.buildTableReadQuery(0, nil)
+		assert.Contains(t, got, "ORDER BY created_at, id")
+		assert.NotContains(t, got, "WHERE")
+	})
+	t.Run("incremental by id", func(t *testing.T) {
+		got := c.buildTableReadQuery(100, nil)
+		assert.Contains(t, got, "WHERE id > 100")
+		assert.Contains(t, got, "ORDER BY id")
+	})
+	t.Run("incremental by created_at", func(t *testing.T) {
+		ts := time.Date(2024, 6, 1, 12, 0, 0, 0, time.UTC)
+		got := c.buildTableReadQuery(0, &ts)
+		assert.Contains(t, got, "WHERE created_at > '2024-06-01 12:00:00'")
+		assert.Contains(t, got, "ORDER BY created_at, id")
+	})
+	t.Run("custom orderByColumn", func(t *testing.T) {
+		spec := &v1.ClickHouseSourceSpec{
+			ConnectionString: "clickhouse://localhost:9000",
+			Table:            "prices",
+			OrderByColumn:    "price_id",
+		}
+		c := NewClickHouseSourceConnector(spec)
+		got := c.buildTableReadQuery(50, nil)
+		assert.Contains(t, got, "WHERE price_id > 50")
+		assert.Contains(t, got, "ORDER BY price_id")
+	})
+}
+
+func TestClickHouseSourceConnector_wrapQueryWithStableOrder(t *testing.T) {
+	c := NewClickHouseSourceConnector(&v1.ClickHouseSourceSpec{
+		ConnectionString: "clickhouse://localhost:9000",
+		Table:            "t",
+		OrderByColumn:    "price_id",
+	})
+	got := c.wrapQueryWithStableOrder("SELECT * FROM prices")
+	assert.Contains(t, got, "__dataflow_src")
+	assert.Contains(t, got, "ORDER BY price_id")
+}