Merge #143456 #143489

143456: crosscluster/producer: allow table level auth for REPLICATIONSOURCE r=jeffswenson a=msbutler

Epic: none

Release note (ops change): previously, the user provided in the source URI in
the LDR stream required the REPLICATIONSOURCE priv at the system level. With
this change, the user only needs this priv on the source tables (i.e. a table
level priv).

143489: cspann: add ChildKey de-duplicator r=drewkimball a=andy-kimball

Add a new de-duplicator class for ChildKeys. Duplicate ChildKeys are going to become much more common with non-transactional fixups. Add an O(N) implementation that minimizes allocations by avoiding hashing KeyBytes values as strings. Instead, it hashes KeyBytes into a uint64 value, which is then hashed using a regular Go map. uint64 collisions are handled by rehashing.

Use the new de-duper in the index. In a later PR, it will be used directly in SearchSet as well to detect duplicates as early as possible.

Epic: CRDB-42943

Release note: None

Co-authored-by: Michael Butler <[email protected]>
Co-authored-by: Andrew Kimball <[email protected]>

Author: 3 people

pkg/crosscluster/logical/logical_replication_job.go

@@ -488,6 +488,7 @@ func (p *logicalReplicationPlanner) generatePlanImpl(
 		// During an offline initial scan, we need to replicate the whole table, not
 		// just the primary keys.
 		UseTableSpan: payload.CreateTable && progress.ReplicatedTime.IsEmpty(),
+		StreamID:     streampb.StreamID(payload.StreamID),
 	}
 	for _, pair := range payload.ReplicationPairs {
 		req.TableIDs = append(req.TableIDs, pair.SrcDescriptorID)

pkg/crosscluster/logical/logical_replication_job_test.go

@@ -498,10 +498,11 @@ func TestCreateTables(t *testing.T) {
 		gURL := replicationtestutils.GetExternalConnectionURI(t, srv, srv, serverutils.DBName("g"))
 
 		sqlG.Exec(t, "CREATE TABLE tab (pk int primary key, payload string)")
+		sqlG.Exec(t, "CREATE TABLE foo (pk int primary key, payload string)")
 
 		var jobID jobspb.JobID
 		// use create logically replicated table syntax
-		sqlG.QueryRow(t, "CREATE LOGICALLY REPLICATED TABLE tab2 FROM TABLE tab ON $1 WITH UNIDIRECTIONAL", gURL.String()).Scan(&jobID)
+		sqlG.QueryRow(t, "CREATE LOGICALLY REPLICATED TABLES (tab2, foo2) FROM TABLES (tab, foo) ON $1 WITH BIDIRECTIONAL ON $2", gURL.String(), gURL.String()).Scan(&jobID)
 		WaitUntilReplicatedTime(t, srv.Clock().Now(), sqlG, jobID)
 		// check that tab2 is empty
 		sqlG.CheckQueryResults(t, "SELECT * FROM tab2", [][]string{})
@@ -1581,6 +1582,32 @@ func WaitUntilReplicatedTime(
 	})
 }
 
+func GetReverseJobID(
+	ctx context.Context, t *testing.T, db *sqlutils.SQLRunner, parentID jobspb.JobID,
+) jobspb.JobID {
+	// get the created time of the parent job
+	var created time.Time
+	db.QueryRow(t, "SELECT created FROM system.jobs WHERE id = $1", parentID).Scan(&created)
+
+	var jobID jobspb.JobID
+	testutils.SucceedsSoon(t, func() error {
+		err := db.DB.QueryRowContext(ctx, `
+            SELECT id 
+            FROM system.jobs 
+            WHERE job_type = 'LOGICAL REPLICATION' 
+            AND id != $1
+            AND created > $2
+            ORDER BY created DESC 
+            LIMIT 1`,
+			parentID, created).Scan(&jobID)
+		if err != nil {
+			return errors.Wrapf(err, "reverse job not found for parent %d", parentID)
+		}
+		return nil
+	})
+	return jobID
+}
+
 type mockBatchHandler struct {
 	err error
 }
@@ -2028,10 +2055,11 @@ func TestUserPrivileges(t *testing.T) {
 		testuser.ExpectErr(t, "failed privilege check: table or system level REPLICATIONDEST privilege required: user testuser does not have REPLICATIONDEST privilege on relation tab", createStmt, dbBURL.String())
 		dbA.Exec(t, fmt.Sprintf("GRANT SYSTEM REPLICATIONDEST TO %s", username.TestUser))
 		testuser.QueryRow(t, createStmt, dbBURL.String()).Scan(&jobAID)
+		WaitUntilReplicatedTime(t, s.Clock().Now(), dbA, jobAID)
 		dbA.Exec(t, fmt.Sprintf("REVOKE SYSTEM REPLICATIONDEST FROM %s", username.TestUser))
 	})
 	t.Run("replication-src", func(t *testing.T) {
-		dbA.ExpectErr(t, "user testuser3 does not have REPLICATIONSOURCE system privilege", createStmt, dbBURL2.String())
+		dbA.ExpectErr(t, "user testuser3 does not have REPLICATIONSOURCE privilege on relation tab", createStmt, dbBURL2.String())
 		sourcePriv := "REPLICATIONSOURCE"
 		if rng.Intn(3) == 0 {
 			// Test deprecated privilege name.
@@ -2042,6 +2070,15 @@ func TestUserPrivileges(t *testing.T) {
 		dbA.QueryRow(t, createStmt, dbBURL2.String()).Scan(&jobAID)
 		dbB.Exec(t, fmt.Sprintf("REVOKE SYSTEM %s FROM %s", sourcePriv, username.TestUser+"3"))
 	})
+	t.Run("table-level-replication-src", func(t *testing.T) {
+		dbA.ExpectErr(t, "user testuser3 does not have REPLICATIONSOURCE privilege on relation tab", createStmt, dbBURL2.String())
+
+		dbB.Exec(t, fmt.Sprintf("GRANT REPLICATIONSOURCE ON TABLE tab TO %s", username.TestUser+"3"))
+		dbA.QueryRow(t, createStmt, dbBURL2.String()).Scan(&jobAID)
+		WaitUntilReplicatedTime(t, s.Clock().Now(), dbA, jobAID)
+
+		dbB.Exec(t, fmt.Sprintf("REVOKE REPLICATIONSOURCE ON TABLE tab FROM %s", username.TestUser+"3"))
+	})
 	t.Run("table-level-replication-dest", func(t *testing.T) {
 
 		dbA.Exec(t, `CREATE TABLE tab2 (x INT PRIMARY KEY)`)
@@ -2055,7 +2092,9 @@ func TestUserPrivileges(t *testing.T) {
 		testuser.ExpectErr(t, "failed privilege check: table or system level REPLICATIONDEST privilege required: user testuser does not have REPLICATIONDEST privilege on relation tab2", multiTableStmt, dbBURL.String())
 
 		dbA.Exec(t, fmt.Sprintf(`GRANT REPLICATIONDEST ON TABLE tab2 TO %s`, username.TestUser))
-		testuser.Exec(t, multiTableStmt, dbBURL.String())
+		testuser.QueryRow(t, multiTableStmt, dbBURL.String()).Scan(&jobAID)
+		WaitUntilReplicatedTime(t, s.Clock().Now(), dbA, jobAID)
+
 		dbA.Exec(t, fmt.Sprintf(`REVOKE REPLICATIONDEST ON TABLE tab FROM %s`, username.TestUser))
 		dbA.Exec(t, fmt.Sprintf(`REVOKE REPLICATIONDEST ON TABLE tab2 FROM %s`, username.TestUser))
 	})
@@ -2066,7 +2105,8 @@ func TestUserPrivileges(t *testing.T) {
 
 		// Grant CREATE privilege on destination database - should now succeed
 		dbA.Exec(t, `GRANT CREATE ON DATABASE a TO testuser`)
-		testuser.Exec(t, createStmt, dbBURL.String())
+		testuser.QueryRow(t, createStmt, dbBURL.String()).Scan(&jobAID)
+		WaitUntilReplicatedTime(t, s.Clock().Now(), dbA, jobAID)
 
 		dbAURL := replicationtestutils.GetExternalConnectionURI(t, s, s, serverutils.DBName("a"), serverutils.User(username.TestUser))
 
@@ -2075,8 +2115,14 @@ func TestUserPrivileges(t *testing.T) {
 		testuser.ExpectErr(t, " uri requires REPLICATIONDEST privilege for bidirectional replication: user testuser3 does not have REPLICATIONDEST privilege on relation tab", createStmtBidi, dbBURL2.String(), dbAURL.String())
 
 		dbB.Exec(t, fmt.Sprintf("GRANT SYSTEM REPLICATIONDEST TO %s", username.TestUser+"3"))
-		testuser.QueryRow(t, createStmtBidi, dbBURL2.String(), dbAURL.String()).Scan(&jobAID)
 
+		var jobAID2 jobspb.JobID
+		testuser.QueryRow(t, createStmtBidi, dbBURL2.String(), dbAURL.String()).Scan(&jobAID2)
+		WaitUntilReplicatedTime(t, s.Clock().Now(), dbA, jobAID2)
+
+		// Ensure the reverse job advances as well
+		reverseJobID := GetReverseJobID(ctx, t, dbA, jobAID2)
+		WaitUntilReplicatedTime(t, s.Clock().Now(), dbA, reverseJobID)
 	})
 }
 

pkg/crosscluster/producer/replication_manager.go

@@ -450,12 +450,30 @@ func (r *replicationStreamManagerImpl) AuthorizeViaJob(
 	return nil
 }
 
-func (r *replicationStreamManagerImpl) AuthorizeViaReplicationPriv(ctx context.Context) error {
-	err := r.evalCtx.SessionAccessor.CheckPrivilege(ctx,
-		syntheticprivilege.GlobalPrivilegeObject,
-		privilege.REPLICATIONSOURCE)
+// AuthorizeViaReplicationPriv ensures the user has the REPLICATIONSOUCE privilege. If tableNames is passed, then table level auth is tried.
+func (r *replicationStreamManagerImpl) AuthorizeViaReplicationPriv(
+	ctx context.Context, tableNames ...string,
+) (err error) {
+
+	authorize := func() error {
+		// First try fast path for system level priv.
+		err = r.evalCtx.SessionAccessor.CheckPrivilege(ctx,
+			syntheticprivilege.GlobalPrivilegeObject,
+			privilege.REPLICATIONSOURCE)
+		if err == nil {
+			return nil
+		} else if pgerror.GetPGCode(err) != pgcode.InsufficientPrivilege {
+			return err
+		}
+		if len(tableNames) != 0 {
+			err = replicationutils.AuthorizeTableLevelPriv(ctx, r.resolver, r.evalCtx.SessionAccessor, privilege.REPLICATIONSOURCE, tableNames)
+			if err == nil {
+				return nil
+			} else if pgerror.GetPGCode(err) != pgcode.InsufficientPrivilege {
+				return err
+			}
+		}
 
-	if pgerror.GetPGCode(err) == pgcode.InsufficientPrivilege {
 		// Fallback to legacy REPLICATION priv.
 		if fallbackErr := r.evalCtx.SessionAccessor.CheckPrivilege(ctx,
 			syntheticprivilege.GlobalPrivilegeObject,
@@ -465,10 +483,12 @@ func (r *replicationStreamManagerImpl) AuthorizeViaReplicationPriv(ctx context.C
 			// the deprecated REPLICATION priv.
 			return err
 		}
-	} else if err != nil {
-		return err
+		return nil
 	}
 
+	if err = authorize(); err != nil {
+		return err
+	}
 	r.authorized = true
 	return nil
 }

pkg/crosscluster/replicationutils/utils.go

@@ -347,6 +347,10 @@ func AuthorizeTableLevelPriv(
 			return err
 		}
 		lookupFlags := tree.ObjectLookupFlags{
+			// TODO(msbutler): for reasons beyond my paygrade, to grab offline
+			// descriptors, we need to also pass RequireMutable.
+			RequireMutable:       true,
+			IncludeOffline:       true,
 			Required:             true,
 			DesiredObjectKind:    tree.TableObject,
 			DesiredTableDescKind: tree.ResolveRequireTableDesc,

pkg/repstream/streampb/stream.proto

@@ -170,6 +170,7 @@ message LogicalReplicationPlanRequest {
     // PlanAsOf is the time as of which the plan should be produced.
     util.hlc.Timestamp plan_as_of = 2 [(gogoproto.nullable) = false];
     bool use_table_span = 3;
+    int64 stream_id = 4 [(gogoproto.customname) = "StreamID", (gogoproto.casttype) = "StreamID"];
 }
 
 // SourcePartition contains per partition information for a replication plan.

pkg/sql/sem/builtins/replication_builtins.go

@@ -511,7 +511,14 @@ var replicationBuiltins = map[string]builtinDefinition{
 				if err := protoutil.Unmarshal(reqBytes, &req); err != nil {
 					return nil, err
 				}
-				if err := mgr.AuthorizeViaReplicationPriv(ctx); err != nil {
+				if req.StreamID != 0 {
+					if err := mgr.AuthorizeViaJob(ctx, req.StreamID); err != nil {
+						return nil, err
+					}
+					// Auth via replication priv exists to ensure a user that planned their
+					// job pre 25.2, which will not send a stream id, can still plan their
+					// distsql flow.
+				} else if err := mgr.AuthorizeViaReplicationPriv(ctx); err != nil {
 					return nil, err
 				}
 
@@ -546,14 +553,14 @@ var replicationBuiltins = map[string]builtinDefinition{
 				if err != nil {
 					return nil, err
 				}
-				if err := mgr.AuthorizeViaReplicationPriv(ctx); err != nil {
-					return nil, err
-				}
 				reqBytes := []byte(tree.MustBeDBytes(args[0]))
 				req := streampb.ReplicationProducerRequest{}
 				if err := protoutil.Unmarshal(reqBytes, &req); err != nil {
 					return nil, err
 				}
+				if err := mgr.AuthorizeViaReplicationPriv(ctx, req.TableNames...); err != nil {
+					return nil, errors.Wrapf(err, "failed to auth")
+				}
 
 				spec, err := mgr.StartReplicationStreamForTables(ctx, req)
 				if err != nil {

pkg/sql/sem/eval/context.go

@@ -934,7 +934,7 @@ type ReplicationStreamManager interface {
 	) (streampb.ReplicationProducerSpec, error)
 
 	AuthorizeViaJob(ctx context.Context, streamID streampb.StreamID) error
-	AuthorizeViaReplicationPriv(ctx context.Context) error
+	AuthorizeViaReplicationPriv(ctx context.Context, tableNames ...string) error
 }
 
 // StreamIngestManager represents a collection of APIs that streaming replication supports

pkg/sql/vecindex/cspann/BUILD.bazel

@@ -11,6 +11,7 @@ filegroup(
 go_library(
     name = "cspann",
     srcs = [
+        "childkey_dedup.go",
         "cspannpb.go",
         "fixup_processor.go",
         "fixup_split.go",
@@ -49,6 +50,7 @@ go_library(
 go_test(
     name = "cspann_test",
     srcs = [
+        "childkey_dedup_test.go",
         "cspannpb_test.go",
         "fixup_processor_test.go",
         "fixup_worker_test.go",

pkg/sql/vecindex/cspann/childkey_dedup.go

@@ -0,0 +1,126 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package cspann
+
+import (
+	"bytes"
+	"hash/maphash"
+
+	"github.com/cockroachdb/errors"
+)
+
+// hashKeyFunc is a function type for hashing KeyBytes.
+type hashKeyFunc func(KeyBytes) uint64
+
+// childKeyDeDup provides de-duplication for ChildKey values. It supports both
+// PartitionKey and KeyBytes child keys efficiently without making unnecessary
+// allocations.
+type childKeyDeDup struct {
+	// initialCapacity is used to initialize the size of the data structures used
+	// by the de-duplicator.
+	initialCapacity int
+
+	// partitionKeys is used for PartitionKey deduplication.
+	partitionKeys map[PartitionKey]struct{}
+
+	// keyBytesMap maps from a KeyBytes hash to the actual KeyBytes for direct
+	// lookup and deduplication.
+	keyBytesMap map[uint64]KeyBytes
+
+	// seed pseudo-randomizes the hash function used by the de-duplicator.
+	seed maphash.Seed
+
+	// hashKeyBytes is the function used to hash KeyBytes. This is primarily used
+	// for testing to override the default hash function.
+	hashKeyBytes hashKeyFunc
+}
+
+// Init initializes the de-duplicator.
+func (dd *childKeyDeDup) Init(capacity int) {
+	dd.initialCapacity = capacity
+	dd.seed = maphash.MakeSeed()
+	dd.hashKeyBytes = dd.defaultHashKeyBytes
+	dd.Clear()
+}
+
+// TryAdd attempts to add a child key to the deduplication set. It returns true
+// if the key was added (wasn't a duplicate), or false if the key already exists
+// (is a duplicate).
+func (dd *childKeyDeDup) TryAdd(childKey ChildKey) bool {
+	// Handle PartitionKey case - simple map lookup.
+	if childKey.PartitionKey != 0 {
+		// Lazily initialize the partitionKeys map.
+		if dd.partitionKeys == nil {
+			dd.partitionKeys = make(map[PartitionKey]struct{}, dd.initialCapacity)
+		}
+
+		if _, exists := dd.partitionKeys[childKey.PartitionKey]; exists {
+			return false
+		}
+		dd.partitionKeys[childKey.PartitionKey] = struct{}{}
+		return true
+	}
+
+	// Handle KeyBytes case. Lazily initialize the KeyBytes structures.
+	if dd.keyBytesMap == nil {
+		dd.keyBytesMap = make(map[uint64]KeyBytes, dd.initialCapacity)
+	}
+
+	// Calculate original hash for the key bytes.
+	hash := dd.hashKeyBytes(childKey.KeyBytes)
+
+	// Check for the key, possibly having to look at multiple rehashed slots.
+	iterations := 0
+	for {
+		existingKey, exists := dd.keyBytesMap[hash]
+		if !exists {
+			// No collision, we can use this hash.
+			break
+		}
+
+		// Check if this is the same key.
+		if bytes.Equal(existingKey, childKey.KeyBytes) {
+			return false
+		}
+
+		// Hash collision, rehash to find a new slot.
+		hash = dd.rehash(hash)
+
+		iterations++
+		if iterations >= 100000 {
+			// We must have hit a cycle, which should never happen.
+			panic(errors.AssertionFailedf("rehash function cycled"))
+		}
+	}
+
+	// Add the key to the map.
+	dd.keyBytesMap[hash] = childKey.KeyBytes
+	return true
+}
+
+// Clear removes all entries from the deduplication set.
+func (dd *childKeyDeDup) Clear() {
+	// Reset all the data structures.
+	clear(dd.partitionKeys)
+	clear(dd.keyBytesMap)
+}
+
+// defaultHashKeyBytes is the default implementation of hashKeyBytes.
+func (dd *childKeyDeDup) defaultHashKeyBytes(key KeyBytes) uint64 {
+	return maphash.Bytes(dd.seed, key)
+}
+
+// rehash creates a new hash from an existing hash to resolve collisions.
+func (dd *childKeyDeDup) rehash(hash uint64) uint64 {
+	// These constants are large 64-bit primes.
+	hash ^= 0xc3a5c85c97cb3127
+	hash ^= hash >> 33
+	hash *= 0xff51afd7ed558ccd
+	hash ^= hash >> 33
+	hash *= 0xc4ceb9fe1a85ec53
+	hash ^= hash >> 33
+	return hash
+}