sql: support partial index schema changes in LDR tables using the SQL writer

msbutler · msbutler · commit c7cc4d9ce3b9 · 2025-04-17T08:42:19.000-04:00
If logical_replication.consumer.immediate_mode_writer is set to legacy-kv, we
disallow schema changes related to partial indexes on table with LDR jobs
running on them.

Note that if a user runs LDR in the undocumented validated mode, which always
uses the sql writer, and but has this setting equal 'legacy-kv', partial index
schema changes will still fail. This is a papercuit i'm willing to deal with.

Epic: none

Release note (ops change): partial index schema changes are supported in
replicating tables when logical_replication.consumer.immediate_mode_writer is
not set to legacy-kv
diff --git a/pkg/crosscluster/logical/logical_replication_job_test.go b/pkg/crosscluster/logical/logical_replication_job_test.go
@@ -2180,7 +2180,6 @@ func TestLogicalReplicationSchemaChanges(t *testing.T) {
 		// adding unsuppored indexes disallowed
 		{"add virtual column index", "CREATE INDEX virtual_col_idx ON tab(virtual_col)", false},
 		{"add hash index", "CREATE INDEX hash_idx ON tab(pk) USING HASH WITH (bucket_count = 4)", false},
-		{"add partial index", "CREATE INDEX partial_idx ON tab(composite_col) WHERE pk > 0", false},
 		{"add unique index", "CREATE UNIQUE INDEX unique_idx ON tab(composite_col)", false},
 
 		// Drop table is blocked
@@ -2421,9 +2420,17 @@ func TestPartialIndexes(t *testing.T) {
 	dbA.CheckQueryResultsRetry(t, "SELECT * FROM foo WHERE pi = 2", [][]string{{"1", "2", "hello"}})
 	dbA.CheckQueryResultsRetry(t, "SELECT * FROM foo WHERE pi = -2", [][]string{{"-1", "-2", "world"}})
 
+	// Check that dbA and dbB can add an additional partial index, enabled when using the sql writer.
+	dbA.Exec(t, "CREATE INDEX idx2 ON a.foo (pi) WHERE pk = 0")
+	dbB.Exec(t, "CREATE INDEX idx2 ON b.foo (pi) WHERE pk = 0")
+
 	// Check for partial indexes when using legacy kv writer.
 	dbA.Exec(t, "SET CLUSTER SETTING logical_replication.consumer.immediate_mode_writer = 'legacy-kv'")
 	dbA.ExpectErr(t, "cannot create logical replication stream: table foo has a partial index idx", "CREATE LOGICAL REPLICATION STREAM FROM TABLE foo ON $1 INTO TABLE foo", dbBURL.String())
+
+	// With the kv writer set, we can no longer create partial indexes on the replicating tables.
+	dbA.ExpectErr(t, " this schema change is disallowed on table foo because it is referenced by one or more logical replication jobs", "CREATE INDEX idx3 ON a.foo (pi) WHERE pk = 0")
+	dbB.ExpectErr(t, " this schema change is disallowed on table foo because it is referenced by one or more logical replication jobs", "CREATE INDEX idx3 ON b.foo (pi) WHERE pk = 0")
 }
 
 // TestLogicalReplicationCreationChecks verifies that we check that the table
diff --git a/pkg/sql/alter_index.go b/pkg/sql/alter_index.go
@@ -46,7 +46,7 @@ func (p *planner) AlterIndex(ctx context.Context, n *tree.AlterIndex) (planNode,
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/alter_index_visible.go b/pkg/sql/alter_index_visible.go
@@ -66,7 +66,7 @@ func (p *planner) AlterIndexVisible(
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go
@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/auditlogging/auditevents"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
@@ -37,6 +38,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/semenumpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlclustersettings"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
@@ -101,7 +103,7 @@ func (p *planner) AlterTable(ctx context.Context, n *tree.AlterTable) (planNode,
 
 	// Disallow schema changes if this table's schema is locked, unless it is to
 	// set/reset the "schema_locked" storage parameter.
-	if err = checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err = checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
@@ -460,7 +462,7 @@ func (n *alterTableNode) startExec(params runParams) error {
 				for _, updated := range affected {
 					// Disallow schema change if the FK references a table whose schema is
 					// locked.
-					if err := checkSchemaChangeIsAllowed(updated, n.n); err != nil {
+					if err := checkSchemaChangeIsAllowed(updated, n.n, params.ExecCfg().Settings); err != nil {
 						return err
 					}
 					if err := params.p.writeSchemaChange(
@@ -2344,7 +2346,9 @@ func (p *planner) tryRemoveFKBackReferences(
 //     not modifying the value of schema_locked.
 //   - The table is referenced by logical data replication jobs, and the statement
 //     is not in the allow list of LDR schema changes.
-func checkSchemaChangeIsAllowed(desc catalog.TableDescriptor, n tree.Statement) (ret error) {
+func checkSchemaChangeIsAllowed(
+	desc catalog.TableDescriptor, n tree.Statement, settings *cluster.Settings,
+) (ret error) {
 	if desc == nil {
 		return nil
 	}
@@ -2358,9 +2362,9 @@ func checkSchemaChangeIsAllowed(desc catalog.TableDescriptor, n tree.Statement)
 				virtualColNames = append(virtualColNames, col.GetName())
 			}
 		}
-		if !tree.IsAllowedLDRSchemaChange(n, virtualColNames) {
+		kvWriterEnabled := sqlclustersettings.LDRWriterType(sqlclustersettings.LDRImmediateModeWriter.Get(&settings.SV))
+		if !tree.IsAllowedLDRSchemaChange(n, virtualColNames, kvWriterEnabled == sqlclustersettings.LDRWriterTypeLegacyKV) {
 			return sqlerrors.NewDisallowedSchemaChangeOnLDRTableErr(desc.GetName(), desc.TableDesc().LDRJobIDs)
-
 		}
 	}
 	return nil
diff --git a/pkg/sql/alter_table_locality.go b/pkg/sql/alter_table_locality.go
@@ -81,7 +81,7 @@ func (p *planner) AlterTableLocality(
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/alter_table_set_schema.go b/pkg/sql/alter_table_set_schema.go
@@ -85,7 +85,7 @@ func (p *planner) AlterTableSetSchema(
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/create_index.go b/pkg/sql/create_index.go
@@ -97,7 +97,7 @@ func (p *planner) CreateIndex(ctx context.Context, n *tree.CreateIndex) (planNod
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/drop_index.go b/pkg/sql/drop_index.go
@@ -66,7 +66,7 @@ func (p *planner) DropIndex(ctx context.Context, n *tree.DropIndex) (planNode, e
 		}
 
 		// Disallow schema changes if this table's schema is locked.
-		if err = checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+		if err = checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 			return nil, err
 		}
 
diff --git a/pkg/sql/drop_table.go b/pkg/sql/drop_table.go
@@ -73,7 +73,7 @@ func (p *planner) DropTable(ctx context.Context, n *tree.DropTable) (planNode, e
 	for _, toDel := range td {
 		droppedDesc := toDel.desc
 		// Disallow the DROP if this table's schema is locked.
-		if err := checkSchemaChangeIsAllowed(droppedDesc, n); err != nil {
+		if err := checkSchemaChangeIsAllowed(droppedDesc, n, p.ExecCfg().Settings); err != nil {
 			return nil, err
 		}
 		for _, fk := range droppedDesc.InboundForeignKeys() {
diff --git a/pkg/sql/rename_column.go b/pkg/sql/rename_column.go
@@ -53,7 +53,7 @@ func (p *planner) RenameColumn(ctx context.Context, n *tree.RenameColumn) (planN
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/rename_index.go b/pkg/sql/rename_index.go
@@ -64,7 +64,7 @@ func (p *planner) RenameIndex(ctx context.Context, n *tree.RenameIndex) (planNod
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/rename_table.go b/pkg/sql/rename_table.go
@@ -88,7 +88,7 @@ func (p *planner) RenameTable(ctx context.Context, n *tree.RenameTable) (planNod
 	}
 
 	// Disallow schema changes if this table's schema is locked.
-	if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {
+	if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {
 		return nil, err
 	}
 
diff --git a/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go b/pkg/sql/schemachanger/scbuild/internal/scbuildstmt/helpers.go
@@ -1006,7 +1006,9 @@ func checkTableSchemaChangePrerequisites(
 			}).MustGetOneElement()
 			virtualColNames = append(virtualColNames, col.Name)
 		})
-		if !tree.IsAllowedLDRSchemaChange(n, virtualColNames) {
+
+		kvWriterEnabled := sqlclustersettings.LDRWriterType(sqlclustersettings.LDRImmediateModeWriter.Get(&b.ClusterSettings().SV))
+		if !tree.IsAllowedLDRSchemaChange(n, virtualColNames, kvWriterEnabled == sqlclustersettings.LDRWriterTypeLegacyKV) {
 			_, _, ns := scpb.FindNamespace(tableElements)
 			if ns == nil {
 				panic(errors.AssertionFailedf("programming error: Namespace element not found"))
diff --git a/pkg/sql/sem/tree/schema_helpers.go b/pkg/sql/sem/tree/schema_helpers.go
@@ -32,7 +32,7 @@ func IsSetOrResetSchemaLocked(n Statement) bool {
 // IsAllowedLDRSchemaChange returns true if the schema change statement is
 // allowed to occur while the table is being referenced by a logical data
 // replication job as a destination table.
-func IsAllowedLDRSchemaChange(n Statement, virtualColNames []string) bool {
+func IsAllowedLDRSchemaChange(n Statement, virtualColNames []string, kvWriterEnabled bool) bool {
 	switch s := n.(type) {
 	case *CreateIndex:
 		// Don't allow creating an index on a virtual column.
@@ -41,6 +41,10 @@ func IsAllowedLDRSchemaChange(n Statement, virtualColNames []string) bool {
 				return false
 			}
 		}
+		if s.Predicate != nil && !kvWriterEnabled {
+			// The sql writer supports partial indexes.
+			return true
+		}
 		// Disallow unique, partial, or hash-sharded indexes. Having these indexes
 		// on a destination table could cause inserts to fail.
 		// NB: hash-sharded indexes are disallowed since they create an index on a
diff --git a/pkg/sql/sem/tree/schema_helpers_test.go b/pkg/sql/sem/tree/schema_helpers_test.go
@@ -118,7 +118,7 @@ func TestIsAllowedLDRSchemaChange(t *testing.T) {
 			}
 			// Tests for virtual column checks are in
 			// TestLogicalReplicationCreationChecks.
-			if got := tree.IsAllowedLDRSchemaChange(stmt.AST, nil /* virtualColNames */); got != tc.isAllowed {
+			if got := tree.IsAllowedLDRSchemaChange(stmt.AST, nil /* virtualColNames */, true); got != tc.isAllowed {
 				t.Errorf("expected %v, got %v", tc.isAllowed, got)
 			}
 		})
diff --git a/pkg/sql/set_zone_config.go b/pkg/sql/set_zone_config.go
@@ -347,7 +347,7 @@ func (n *setZoneConfigNode) startExec(params runParams) error {
 	}
 
 	// Disallow schema changes if it's a table and its schema is locked.
-	if err = checkSchemaChangeIsAllowed(table, n.stmt); err != nil {
+	if err = checkSchemaChangeIsAllowed(table, n.stmt, params.p.ExecCfg().Settings); err != nil {
 		return err
 	}
 

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ func (p planner) AlterIndex(ctx context.Context, n tree.AlterIndex) (planNode,`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`// Disallow schema changes if this table's schema is locked.`
`49`		`- if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`49`	`+ if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`50`	`50`	`return nil, err`
`51`	`51`	`}`
`52`	`52`
Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ func (p *planner) AlterIndexVisible(`
`66`	`66`	`}`
`67`	`67`
`68`	`68`	`// Disallow schema changes if this table's schema is locked.`
`69`		`- if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`69`	`+ if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`70`	`70`	`return nil, err`
`71`	`71`	`}`
`72`	`72`
Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ func (p *planner) AlterTableLocality(`
`81`	`81`	`}`
`82`	`82`
`83`	`83`	`// Disallow schema changes if this table's schema is locked.`
`84`		`- if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`84`	`+ if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`85`	`85`	`return nil, err`
`86`	`86`	`}`
`87`	`87`
Original file line number	Diff line number	Diff line change
`@@ -85,7 +85,7 @@ func (p *planner) AlterTableSetSchema(`
`85`	`85`	`}`
`86`	`86`
`87`	`87`	`// Disallow schema changes if this table's schema is locked.`
`88`		`- if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`88`	`+ if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`89`	`89`	`return nil, err`
`90`	`90`	`}`
`91`	`91`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ func (p planner) CreateIndex(ctx context.Context, n tree.CreateIndex) (planNod`
`97`	`97`	`}`
`98`	`98`
`99`	`99`	`// Disallow schema changes if this table's schema is locked.`
`100`		`- if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`100`	`+ if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`101`	`101`	`return nil, err`
`102`	`102`	`}`
`103`	`103`
Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ func (p planner) DropIndex(ctx context.Context, n tree.DropIndex) (planNode, e`
`66`	`66`	`}`
`67`	`67`
`68`	`68`	`// Disallow schema changes if this table's schema is locked.`
`69`		`- if err = checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`69`	`+ if err = checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`70`	`70`	`return nil, err`
`71`	`71`	`}`
`72`	`72`
Original file line number	Diff line number	Diff line change
`@@ -73,7 +73,7 @@ func (p planner) DropTable(ctx context.Context, n tree.DropTable) (planNode, e`
`73`	`73`	`for _, toDel := range td {`
`74`	`74`	`droppedDesc := toDel.desc`
`75`	`75`	`// Disallow the DROP if this table's schema is locked.`
`76`		`- if err := checkSchemaChangeIsAllowed(droppedDesc, n); err != nil {`
	`76`	`+ if err := checkSchemaChangeIsAllowed(droppedDesc, n, p.ExecCfg().Settings); err != nil {`
`77`	`77`	`return nil, err`
`78`	`78`	`}`
`79`	`79`	`for _, fk := range droppedDesc.InboundForeignKeys() {`
Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,7 @@ func (p planner) RenameColumn(ctx context.Context, n tree.RenameColumn) (planN`
`53`	`53`	`}`
`54`	`54`
`55`	`55`	`// Disallow schema changes if this table's schema is locked.`
`56`		`- if err := checkSchemaChangeIsAllowed(tableDesc, n); err != nil {`
	`56`	`+ if err := checkSchemaChangeIsAllowed(tableDesc, n, p.ExecCfg().Settings); err != nil {`
`57`	`57`	`return nil, err`
`58`	`58`	`}`
`59`	`59`