Adding a dolt_update_column_tag() stored procedure

go/cmd/dolt/cli/arg_parser_helpers.go

@@ -80,6 +80,14 @@ func CreateConflictsResolveArgParser() *argparser.ArgParser {
 	return ap
 }
 
+func CreateUpdateTagArgParser() *argparser.ArgParser {
+	ap := argparser.NewArgParserWithMaxArgs("update-tag", 3)
+	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"table", "The name of the table"})
+	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"column", "The name of the column"})
+	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"tag", "The new tag value"})
+	return ap
+}
+
 func CreateMergeArgParser() *argparser.ArgParser {
 	ap := argparser.NewArgParserWithMaxArgs("merge", 1)
 	ap.TooManyArgsErrorFunc = func(receivedArgs []string) error {

go/cmd/dolt/commands/schcmds/update-tag.go

@@ -29,6 +29,8 @@ import (
 	"github.com/dolthub/dolt/go/store/types"
 )
 
+// TODO: Update tag should be migrated to call the new dolt_update_column_tag() stored procedure
+
 var updateTagDocs = cli.CommandDocumentationContent{
 	ShortDesc: "Update the tag of the specified column",
 	LongDesc: `{{.EmphasisLeft}}dolt schema update-tag{{.EmphasisRight}}
@@ -59,11 +61,7 @@ func (cmd UpdateTagCmd) Docs() *cli.CommandDocumentation {
 }
 
 func (cmd UpdateTagCmd) ArgParser() *argparser.ArgParser {
-	ap := argparser.NewArgParserWithMaxArgs(cmd.Name(), 3)
-	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"table", "The name of the table"})
-	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"column", "The name of the column"})
-	ap.ArgListHelp = append(ap.ArgListHelp, [2]string{"tag", "The new tag value"})
-	return ap
+	return cli.CreateUpdateTagArgParser()
 }
 
 func (cmd UpdateTagCmd) Exec(ctx context.Context, commandStr string, args []string, dEnv *env.DoltEnv, cliCtx cli.CliContext) int {

go/libraries/doltcore/sqle/dprocedures/dolt_update_column_tag.go

@@ -0,0 +1,129 @@
+// Copyright 2025 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dprocedures
+
+import (
+	"fmt"
+	"strconv"
+
+	"github.com/dolthub/go-mysql-server/sql"
+
+	"github.com/dolthub/dolt/go/cmd/dolt/cli"
+	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
+	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
+)
+
+// doltUpdateColumnTag updates the tag for a specified column, leaving the change in the working set to later be
+// committed.
+func doltUpdateColumnTag(ctx *sql.Context, args ...string) (sql.RowIter, error) {
+	tableName, columnName, tag, err := parseUpdateColumnTagArgs(args...)
+	if err != nil {
+		return nil, err
+	}
+
+	doltSession := dsess.DSessFromSess(ctx.Session)
+	roots, ok := doltSession.GetRoots(ctx, ctx.GetCurrentDatabase())
+	if !ok {
+		return nil, fmt.Errorf("unable to load roots")
+	}
+	root := roots.Working
+
+	tbl, tName, ok, err := doltdb.GetTableInsensitive(ctx, root, doltdb.TableName{Name: tableName})
+	if err != nil {
+		return nil, err
+	} else if !ok {
+		return nil, fmt.Errorf("table %s does not exist", tableName)
+	}
+
+	sch, err := tbl.GetSchema(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	newSch, err := updateColumnTag(sch, columnName, tag)
+	if err != nil {
+		return nil, fmt.Errorf("failed to update column tag: %w", err)
+	}
+
+	tbl, err = tbl.UpdateSchema(ctx, newSch)
+	if err != nil {
+		return nil, fmt.Errorf("failed to update table schema: %w", err)
+	}
+
+	root, err = root.PutTable(ctx, doltdb.TableName{Name: tName}, tbl)
+	if err != nil {
+		return nil, fmt.Errorf("failed to put table in root: %w", err)
+	}
+
+	if err = doltSession.SetWorkingRoot(ctx, ctx.GetCurrentDatabase(), root); err != nil {
+		return nil, err
+	}
+
+	return rowToIter(int64(0)), nil
+}
+
+// parseUpdateColumnTagArgs parses |args| and returns the tableName, columnName, and tag specified, otherwise
+// returns an error if there were any problems.
+func parseUpdateColumnTagArgs(args ...string) (tableName, columnName string, tag uint64, err error) {
+	apr, err := cli.ParseArgs(cli.CreateUpdateTagArgParser(), args, nil)
+	if err != nil {
+		return "", "", 0, err
+	}
+	if len(apr.Args) != 3 {
+		return "", "", 0,
+			fmt.Errorf("incorrect number of arguments: must provide <table> <column> <tag>")
+	}
+
+	tableName, columnName, tagStr := apr.Args[0], apr.Args[1], apr.Args[2]
+
+	tag, err = strconv.ParseUint(tagStr, 10, 64)
+	if err != nil {
+		return "", "", 0, fmt.Errorf("failed to parse tag %s: %w", tagStr, err)
+	}
+
+	return tableName, columnName, tag, nil
+}
+
+// updateColumnTag updates |sch| by setting the tag for the column named |name| to |tag|.
+func updateColumnTag(sch schema.Schema, name string, tag uint64) (schema.Schema, error) {
+	var found bool
+	columns := sch.GetAllCols().GetColumns()
+	// Find column and update its tag
+	for i, col := range columns {
+		if col.Name == name {
+			col.Tag = tag
+			columns[i] = col
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		return nil, fmt.Errorf("column %s does not exist", name)
+	}
+
+	newSch, err := schema.SchemaFromCols(schema.NewColCollection(columns...))
+	if err != nil {
+		return nil, err
+	}
+
+	if err = newSch.SetPkOrdinals(sch.GetPkOrdinals()); err != nil {
+		return nil, err
+	}
+	newSch.SetCollation(sch.GetCollation())
+
+	return newSch, nil
+}

go/libraries/doltcore/sqle/dprocedures/init.go

@@ -33,6 +33,7 @@ var DoltProcedures = []sql.ExternalStoredProcedureDetails{
 	{Name: "dolt_count_commits", Schema: int64Schema("ahead", "behind"), Function: doltCountCommits, ReadOnly: true},
 	{Name: "dolt_fetch", Schema: int64Schema("status"), Function: doltFetch, AdminOnly: true},
 	{Name: "dolt_undrop", Schema: int64Schema("status"), Function: doltUndrop, AdminOnly: true},
+	{Name: "dolt_update_column_tag", Schema: int64Schema("status"), Function: doltUpdateColumnTag, AdminOnly: true},
 	{Name: "dolt_purge_dropped_databases", Schema: int64Schema("status"), Function: doltPurgeDroppedDatabases, AdminOnly: true},
 	{Name: "dolt_rebase", Schema: doltRebaseProcedureSchema, Function: doltRebase},
 

integration-tests/bats/sql-update-column-tag.bats

@@ -0,0 +1,68 @@
+#!/usr/bin/env bats
+load $BATS_TEST_DIRNAME/helper/common.bash
+
+setup() {
+    setup_common
+}
+
+teardown() {
+    teardown_common
+}
+
+# Tests the basic functionality of the dolt_update_column_tag stored procedure.
+#
+# Note that we use BATS to test this, since reading column tags is not supported
+# via a SQL interface, only from the `dolt schema tags` command currently,
+# otherwise we'd prefer enginetests in go.
+@test "sql-update-column-tag: update column tag" {
+    dolt sql -q "create table t1 (pk int primary key, c1 int);"
+
+    run dolt schema tags
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ "t1" ]] || false
+    [[ "$output" =~ "pk" ]] || false
+    [[ "$output" =~ "c1" ]] || false
+    [[ ! "$output" =~ " t1    | pk     | 42 " ]] || false
+    [[ ! "$output" =~ " t1    | c1     | 420 " ]] || false
+
+    dolt sql -q "call dolt_update_column_tag('t1', 'pk', 42);"
+    dolt sql -q "call dolt_update_column_tag('t1', 'c1', 420);"
+
+    run dolt schema tags
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ " t1    | pk     | 42 " ]] || false
+    [[ "$output" =~ " t1    | c1     | 420 " ]] || false
+}
+
+# Tests error cases for the dolt_update_column_tag stored procedure.
+@test "sql-update-column-tag: error cases" {
+    dolt sql -q "create table t1 (pk int primary key, c1 int);"
+
+    # invalid arg count
+    run dolt sql -q "call dolt_update_column_tag();"
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "incorrect number of arguments" ]] || false
+
+    run dolt sql -q "call dolt_update_column_tag('t1', 'pk');"
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "incorrect number of arguments" ]] || false
+
+    run dolt sql -q "call dolt_update_column_tag('t1', 'pk', 42, 'zzz');"
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "Expected at most 3" ]] || false
+
+    # invalid table
+    run dolt sql -q "call dolt_update_column_tag('doesnotexist', 'pk', 42);"
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "does not exist" ]] || false
+
+    # invalid column
+    run dolt sql -q "call dolt_update_column_tag('t1', 'doesnotexist', 42);"
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "does not exist" ]] || false
+
+    # invalid tag
+    run dolt sql -q "call dolt_update_column_tag('t1', 'pk', 'not an integer');"
+    [ "$status" -eq 1 ]
+    [[ "$output" =~ "failed to parse tag" ]] || false
+}