Refactor checkpoint acknowledgment handling across connectors to use AssignCompositeSourceAck function for improved consistency and clarity. Added unit tests for order-by-only checkpoint acknowledgment in ClickHouse, PostgreSQL, and Trino connectors.

Author: ilyario

internal/connectors/clickhouse.go

@@ -200,12 +200,7 @@ func (c *ClickHouseSourceConnector) readRows(ctx context.Context, msgChan chan *
 			if idIndex >= 0 && len(values) > idIndex {
 				SetSourceRowIDMetadata(msg, values[idIndex])
 			}
-			if changeTime != nil {
-				ct := *changeTime
-				msg.Ack = c.cp.MakeAck(&ct, orderByVal, true)
-			} else if orderByVal != nil {
-				msg.Ack = c.cp.MakeAck(nil, orderByVal, false)
-			}
+			AssignCompositeSourceAck(msg, &c.cp, changeTime, orderByVal)
 
 			select {
 			case msgChan <- msg:

internal/connectors/clickhouse_test.go

@@ -24,6 +24,7 @@ import (
 
 	v1 "github.com/dataflow-operator/dataflow/api/v1"
 	"github.com/dataflow-operator/dataflow/internal/checkpoint"
+	"github.com/dataflow-operator/dataflow/internal/types"
 	"github.com/go-logr/logr"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -285,6 +286,25 @@ func TestClickHouseSourceConnector_buildReadQuery(t *testing.T) {
 	})
 }
 
+func TestClickHouseSourceConnector_orderByOnlyCheckpointAck(t *testing.T) {
+	c := NewClickHouseSourceConnector(&v1.ClickHouseSourceSpec{
+		ConnectionString:     "clickhouse://localhost:9000",
+		Table:                "mv_one_p_prices_migration",
+		ChangeTrackingColumn: "material_id",
+		OrderByColumn:        "material_id",
+	})
+	msg := types.NewMessage([]byte(`{"material_id":200019}`))
+	AssignCompositeSourceAck(msg, &c.cp, nil, int64(200019))
+	require.NotNil(t, msg.Ack)
+	msg.Ack()
+	snap := c.cp.Snapshot()
+	require.Equal(t, int64(200019), snap.OrderByValue)
+	assert.Nil(t, snap.ChangeTime)
+
+	got := c.buildReadQuery()
+	assert.Contains(t, got, "WHERE `material_id` > 200019")
+}
+
 func TestClickHouseSourceConnector_applyInitialCheckpoint_legacy(t *testing.T) {
 	opts := &SourceConnectorOptions{
 		InitialCheckpoint: []byte(`{"lastReadID":100,"lastReadTime":"2024-06-01 12:00:00"}`),

internal/connectors/composite_source.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/dataflow-operator/dataflow/internal/checkpoint"
+	"github.com/dataflow-operator/dataflow/internal/types"
 )
 
 // SQLDialect formats SQL literals and identifiers for a database engine.
@@ -292,6 +293,20 @@ func (h *CompositeCheckpointHolder) Advance(next checkpoint.Composite, requireTi
 	}
 }
 
+// AssignCompositeSourceAck sets msg.Ack so the composite checkpoint advances after a successful sink write.
+// When changeTime is nil but orderByVal is set (e.g. non-timestamp changeTrackingColumn), orderBy-only ack is used.
+func AssignCompositeSourceAck(msg *types.Message, cp *CompositeCheckpointHolder, changeTime *time.Time, orderByVal interface{}) {
+	if msg == nil || cp == nil {
+		return
+	}
+	if changeTime != nil {
+		ct := *changeTime
+		msg.Ack = cp.MakeAck(&ct, orderByVal, true)
+	} else if orderByVal != nil {
+		msg.Ack = cp.MakeAck(nil, orderByVal, false)
+	}
+}
+
 // MakeAck returns a callback that advances the checkpoint after sink write.
 func (h *CompositeCheckpointHolder) MakeAck(changeTime *time.Time, orderBy interface{}, requireTime bool) func() {
 	return func() {

internal/connectors/composite_source_test.go

@@ -26,6 +26,7 @@ import (
 	"k8s.io/client-go/kubernetes/fake"
 
 	"github.com/dataflow-operator/dataflow/internal/checkpoint"
+	"github.com/dataflow-operator/dataflow/internal/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -374,6 +375,17 @@ func TestCompositeCheckpointHolder_AdvanceAndSyncFlush(t *testing.T) {
 	assert.Contains(t, string(loaded), "5042")
 }
 
+func TestAssignCompositeSourceAck_orderByOnly(t *testing.T) {
+	h := &CompositeCheckpointHolder{}
+	msg := types.NewMessage([]byte(`{}`))
+	AssignCompositeSourceAck(msg, h, nil, int64(200019))
+	require.NotNil(t, msg.Ack)
+	msg.Ack()
+	snap := h.Snapshot()
+	require.Equal(t, int64(200019), snap.OrderByValue)
+	assert.Nil(t, snap.ChangeTime)
+}
+
 func TestCompositeCheckpointHolder_orderOnlyLegacy(t *testing.T) {
 	h := &CompositeCheckpointHolder{}
 	h.ApplyInitial([]byte(`{"lastReadOrderByValue":50}`))

internal/connectors/postgresql.go

@@ -255,12 +255,7 @@ func (p *PostgreSQLSourceConnector) readRows(ctx context.Context, msgChan chan *
 				SetSourceRowIDMetadata(msg, values[idIndex])
 			}
 			msg.Metadata["operation"] = operation
-			if changeTime != nil {
-				ct := *changeTime
-				msg.Ack = p.cp.MakeAck(&ct, orderByVal, true)
-			} else if orderByVal != nil {
-				msg.Ack = p.cp.MakeAck(nil, orderByVal, false)
-			}
+			AssignCompositeSourceAck(msg, &p.cp, changeTime, orderByVal)
 
 			select {
 			case msgChan <- msg:

internal/connectors/postgresql_test.go

@@ -176,8 +176,7 @@ func TestPostgreSQLSourceConnector_orderByOnlyCheckpointAck(t *testing.T) {
 		OrderByColumn:        "material_id",
 	})
 	msg := types.NewMessage([]byte(`{"material_id":200019}`))
-	orderByVal := int64(200019)
-	msg.Ack = p.cp.MakeAck(nil, orderByVal, false)
+	AssignCompositeSourceAck(msg, &p.cp, nil, int64(200019))
 	require.NotNil(t, msg.Ack)
 	msg.Ack()
 	snap := p.cp.Snapshot()

internal/connectors/trino.go

@@ -165,12 +165,7 @@ func (t *TrinoSourceConnector) readRows(ctx context.Context, msgChan chan *types
 			msg.Metadata["table"] = t.config.Table
 			SetSourceRowIDMetadata(msg, orderByVal)
 
-			if changeTime != nil {
-				ct := *changeTime
-				msg.Ack = t.cp.MakeAck(&ct, orderByVal, true)
-			} else if orderByVal != nil {
-				msg.Ack = t.cp.MakeAck(nil, orderByVal, false)
-			}
+			AssignCompositeSourceAck(msg, &t.cp, changeTime, orderByVal)
 
 			select {
 			case msgChan <- msg:

internal/connectors/trino_test.go

@@ -611,6 +611,26 @@ func TestTrinoSourceConnector_buildReadQuery(t *testing.T) {
 	})
 }
 
+func TestTrinoSourceConnector_orderByOnlyCheckpointAck(t *testing.T) {
+	c := NewTrinoSourceConnector(&v1.TrinoSourceSpec{
+		Catalog:              "cat",
+		Schema:               "sch",
+		Table:                "mv_one_p_prices_migration",
+		ChangeTrackingColumn: "material_id",
+		OrderByColumn:        "material_id",
+	})
+	msg := types.NewMessage([]byte(`{"material_id":200019}`))
+	AssignCompositeSourceAck(msg, &c.cp, nil, int64(200019))
+	require.NotNil(t, msg.Ack)
+	msg.Ack()
+	snap := c.cp.Snapshot()
+	require.Equal(t, int64(200019), snap.OrderByValue)
+	assert.Nil(t, snap.ChangeTime)
+
+	got := c.buildReadQuery()
+	assert.Contains(t, got, "WHERE material_id > 200019")
+}
+
 func TestTrinoSourceConnector_wrapQueryWithStableOrder(t *testing.T) {
 	c := NewTrinoSourceConnector(&v1.TrinoSourceSpec{
 		Catalog:       "c",