fixup! feat(jobsdb): introduce new partition_id column

🔒 Scanned for secrets using gitleaks 8.28.0

Author: atzoum

jobsdb/jobsdb.go

@@ -594,6 +594,7 @@ type Handle struct {
 		MaxDSSize                      config.ValueLoader[int]
 		numPartitions                  int // if zero or negative, no partitioning
 		partitionFunction              func(job *JobT) string
+		dbTablesVersion                int // version of the database tables schema (0 means latest)
 
 		migration struct {
 			maxMigrateOnce, maxMigrateDSProbe config.ValueLoader[int]
@@ -826,6 +827,13 @@ func WithNumPartitions(numPartitions int) OptsFunc {
 	}
 }
 
+// withDatabaseTablesVersion sets the database tables version to use (internal use only for verifying database table migrations)
+func withDatabaseTablesVersion(dbVersion int) OptsFunc {
+	return func(jd *Handle) {
+		jd.conf.dbTablesVersion = dbVersion
+	}
+}
+
 func NewForRead(tablePrefix string, opts ...OptsFunc) *Handle {
 	return newOwnerType(Read, tablePrefix, opts...)
 }

jobsdb/jobsdb_tablemigrate_test.go

@@ -0,0 +1,59 @@
+package jobsdb
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+)
+
+func TestJobsDBTableMigrations(t *testing.T) {
+	t.Run("13_partition_id_column", func(t *testing.T) {
+		config.Reset()
+		postgres := startPostgres(t)
+
+		// Start JobsDB with DB version 11 (without partition_id column) to create initial tables
+		jd := NewForWrite("test", WithDBHandle(postgres.DB), withDatabaseTablesVersion(12))
+		require.NoError(t, jd.Start(), "it should be able to start JobsDB")
+		jd.Close()
+
+		// Add some data to ensure migration works fine with existing data
+		_, err := postgres.DB.Exec(`
+			INSERT INTO test_jobs_1 
+				(uuid, workspace_id, user_id, custom_val, parameters, event_payload, event_count) 
+			VALUES 
+				(gen_random_uuid(), 'worskpace-1', 'user-1', 'gw', '{"source_id": "src1"}', '{}', 1)`)
+		require.NoError(t, err, "it should be able to insert initial data")
+
+		// Start JobsDB again to trigger migration to add partition_id column
+		jd = NewForReadWrite("test", WithDBHandle(postgres.DB))
+		require.NoError(t, jd.Start(), "it should be able to start JobsDB with latest table versions")
+		defer jd.Close()
+
+		// Verify that partition_id column exists, but is empty for existing rows
+		unprocessed, err := jd.GetUnprocessed(context.Background(), GetQueryParams{JobsLimit: 1})
+		require.NoError(t, err, "it should be able to get unprocessed jobs")
+		require.Len(t, unprocessed.Jobs, 1)
+		require.Empty(t, unprocessed.Jobs[0].PartitionID)
+
+		// Store a new job and verify partition_id is set
+		require.NoError(t, jd.Store(context.Background(), []*JobT{{
+			UUID:         uuid.New(),
+			WorkspaceId:  "workspace-1",
+			UserID:       "user-1",
+			PartitionID:  "partition-1",
+			CustomVal:    "gw",
+			Parameters:   []byte(`{"source_id": "src2"}`),
+			EventPayload: []byte(`{}`),
+			EventCount:   1,
+		}}))
+
+		unprocessed, err = jd.GetUnprocessed(context.Background(), GetQueryParams{JobsLimit: 2})
+		require.NoError(t, err, "it should be able to get unprocessed jobs")
+		require.Len(t, unprocessed.Jobs, 2)
+		require.Equal(t, "partition-1", unprocessed.Jobs[1].PartitionID)
+	})
+}

jobsdb/jobsdb_test.go

@@ -1774,6 +1774,55 @@ func TestUpdateJobStatus(t *testing.T) {
 	})
 }
 
+func TestPartitionedJobsDB(t *testing.T) {
+	t.Run("without partitioning", func(t *testing.T) {
+		postgres := startPostgres(t)
+		db := NewForReadWrite("partitioning", WithDBHandle(postgres.DB))
+		require.NoError(t, db.Start(), "it should be able to start the jobs db handle")
+		defer db.TearDown()
+		err := db.Store(context.Background(), []*JobT{
+			{
+				UUID:         uuid.New(),
+				WorkspaceId:  "workspace-1",
+				UserID:       "user-1",
+				CustomVal:    "custom-val-1",
+				Parameters:   []byte(`{}`),
+				EventPayload: []byte(`{}`),
+				EventCount:   1,
+			},
+		})
+		require.NoError(t, err, "it should be able to store job")
+		unprocessed, err := db.GetUnprocessed(context.Background(), GetQueryParams{JobsLimit: 1})
+		require.NoError(t, err, "it should be able to get unprocessed jobs")
+		require.Equal(t, 1, len(unprocessed.Jobs), "it should return the stored job")
+		require.Empty(t, unprocessed.Jobs[0].PartitionID, "it should have empty partition ID")
+	})
+
+	t.Run("with partitioning", func(t *testing.T) {
+		postgres := startPostgres(t)
+		db := NewForReadWrite("partitioning", WithNumPartitions(64), WithDBHandle(postgres.DB))
+		require.NoError(t, db.Start(), "it should be able to start the jobs db handle")
+		defer db.TearDown()
+		err := db.Store(context.Background(), []*JobT{
+			{
+				UUID:         uuid.New(),
+				WorkspaceId:  "workspace-1",
+				UserID:       "user-1",
+				CustomVal:    "custom-val-1",
+				Parameters:   []byte(`{}`),
+				EventPayload: []byte(`{}`),
+				EventCount:   1,
+			},
+		})
+		require.NoError(t, err, "it should be able to store job")
+		unprocessed, err := db.GetUnprocessed(context.Background(), GetQueryParams{JobsLimit: 1})
+		require.NoError(t, err, "it should be able to get unprocessed jobs")
+		require.Equal(t, 1, len(unprocessed.Jobs), "it should return the stored job")
+		require.NotEmpty(t, unprocessed.Jobs[0].PartitionID, "it should have non-empty partition ID")
+		require.Equal(t, "workspace-1-62", unprocessed.Jobs[0].PartitionID, "it should have correct partition ID including the workspace ID and partition number")
+	})
+}
+
 type testingT interface {
 	Errorf(format string, args ...interface{})
 	FailNow()

jobsdb/setup.go

@@ -28,6 +28,7 @@ func (jd *Handle) setupDatabaseTables(templateData map[string]interface{}) {
 		Handle:                     jd.dbHandle,
 		MigrationsTable:            jd.SchemaMigrationTable(),
 		ShouldForceSetLowerVersion: jd.config.GetBool("SQLMigrator.forceSetLowerVersion", true),
+		Version:                    uint(jd.conf.dbTablesVersion),
 	}
 	// execute any necessary migrations
 	if err := m.MigrateFromTemplates("jobsdb", templateData); err != nil {

services/sql-migrator/migrator.go

@@ -35,6 +35,9 @@ type Migrator struct {
 
 	// Indicates if all migrations should be run ignoring the current version in MigrationsTable
 	RunAlways bool
+
+	// Migration target version override. 0 means latest.
+	Version uint
 }
 
 var pkgLogger logger.Logger
@@ -161,7 +164,11 @@ func (m *Migrator) MigrateFromTemplates(templatesDir string, context interface{}
 		}
 	}
 
-	err = migration.Up()
+	if m.Version != 0 {
+		err = migration.Migrate(m.Version)
+	} else {
+		err = migration.Up()
+	}
 	if err != nil && err != migrate.ErrNoChange { // migrate library reports that no change was required, using ErrNoChange
 		return fmt.Errorf("run migration from template directory %q, %w", templatesDir, err)
 	}