c2c: pass statement transaction when loading producer side job

This patch fixes a small bug where we previously used a new txn to read from the producer job instead of the txn associated with the crdb_internal builtin call. Epic: none Release note: none
cockroachdb · Aug 1, 2023 · ffea5af · ffea5af
1 parent 0e915fb
commit ffea5af
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 7 deletions.
diff --git a/pkg/ccl/logictestccl/testdata/logic_test/crdb_internal b/pkg/ccl/logictestccl/testdata/logic_test/crdb_internal
@@ -168,8 +168,6 @@ subtest replication-builtins
 
 user root
 
-query error pq: crdb_internal\.replication_stream_spec\(\): job.*is not a replication stream job
-SELECT crdb_internal.replication_stream_spec(crdb_internal.create_sql_schema_telemetry_job())
 
 query error pq: crdb_internal\.stream_ingestion_stats_json\(\): unimplemented
 SELECT crdb_internal.stream_ingestion_stats_json(1);

diff --git a/pkg/ccl/streamingccl/streamclient/partitioned_stream_client_test.go b/pkg/ccl/streamingccl/streamclient/partitioned_stream_client_test.go
@@ -160,6 +160,11 @@ INSERT INTO d.t2 VALUES (2);
 	_, err = client.Plan(ctx, 999)
 	require.True(t, testutils.IsError(err, fmt.Sprintf("job with ID %d does not exist", 999)), err)
 
+	var telemetryJobID int64
+	h.SysSQL.QueryRow(t, "SELECT crdb_internal.create_sql_schema_telemetry_job()").Scan(&telemetryJobID)
+	_, err = client.Plan(ctx, streampb.StreamID(telemetryJobID))
+	require.True(t, testutils.IsError(err, fmt.Sprintf("job with id %d is not a replication stream job", telemetryJobID)), err)
+
 	expectStreamState(streamID, jobs.StatusRunning)
 	status, err := client.Heartbeat(ctx, streamID, hlc.Timestamp{WallTime: timeutil.Now().UnixNano()})
 	require.NoError(t, err)

diff --git a/pkg/ccl/streamingccl/streamproducer/replication_manager.go b/pkg/ccl/streamingccl/streamproducer/replication_manager.go
@@ -55,7 +55,7 @@ func (r *replicationStreamManagerImpl) StreamPartition(
 func (r *replicationStreamManagerImpl) GetReplicationStreamSpec(
 	ctx context.Context, streamID streampb.StreamID,
 ) (*streampb.ReplicationStreamSpec, error) {
-	return getReplicationStreamSpec(ctx, r.evalCtx, streamID)
+	return getReplicationStreamSpec(ctx, r.evalCtx, r.txn, streamID)
 }
 
 // CompleteReplicationStream implements ReplicationStreamManager interface.

diff --git a/pkg/ccl/streamingccl/streamproducer/stream_lifetime.go b/pkg/ccl/streamingccl/streamproducer/stream_lifetime.go
@@ -188,7 +188,7 @@ func heartbeatReplicationStream(
 	// job progress.
 	if frontier == hlc.MaxTimestamp {
 		var status streampb.StreamReplicationStatus
-		pj, err := execConfig.JobRegistry.LoadJob(ctx, jobspb.JobID(streamID))
+		pj, err := execConfig.JobRegistry.LoadJobWithTxn(ctx, jobspb.JobID(streamID), txn)
 		if jobs.HasJobNotFoundError(err) || testutils.IsError(err, "not found in system.jobs table") {
 			status.StreamStatus = streampb.StreamReplicationStatus_STREAM_INACTIVE
 			return status, nil
@@ -219,13 +219,13 @@ func heartbeatReplicationStream(
 
 // getReplicationStreamSpec gets a replication stream specification for the specified stream.
 func getReplicationStreamSpec(
-	ctx context.Context, evalCtx *eval.Context, streamID streampb.StreamID,
+	ctx context.Context, evalCtx *eval.Context, txn isql.Txn, streamID streampb.StreamID,
 ) (*streampb.ReplicationStreamSpec, error) {
 	jobExecCtx := evalCtx.JobExecContext.(sql.JobExecContext)
 	// Returns error if the replication stream is not active
-	j, err := jobExecCtx.ExecCfg().JobRegistry.LoadJob(ctx, jobspb.JobID(streamID))
+	j, err := jobExecCtx.ExecCfg().JobRegistry.LoadJobWithTxn(ctx, jobspb.JobID(streamID), txn)
 	if err != nil {
-		return nil, errors.Wrapf(err, "replication stream %d has error", streamID)
+		return nil, errors.Wrapf(err, "could not load job for replication stream %d", streamID)
 	}
 	if j.Status() != jobs.StatusRunning {
 		return nil, errors.Errorf("replication stream %d is not running", streamID)