From 9277d24f655b81adbeb42f35fe982aa7d9680a32 Mon Sep 17 00:00:00 2001 From: Jason Fulghum Date: Wed, 14 May 2025 16:37:17 -0700 Subject: [PATCH] Adding a dolt_update_column_tag() stored procedure --- go/cmd/dolt/cli/arg_parser_helpers.go | 8 ++ go/cmd/dolt/commands/schcmds/update-tag.go | 8 +- .../dprocedures/dolt_update_column_tag.go | 129 ++++++++++++++++++ .../doltcore/sqle/dprocedures/init.go | 1 + .../bats/sql-update-column-tag.bats | 68 +++++++++ 5 files changed, 209 insertions(+), 5 deletions(-) create mode 100644 go/libraries/doltcore/sqle/dprocedures/dolt_update_column_tag.go create mode 100644 integration-tests/bats/sql-update-column-tag.bats diff --git a/go/cmd/dolt/cli/arg_parser_helpers.go b/go/cmd/dolt/cli/arg_parser_helpers.go index ee62b1fe70..b02fffc6f7 100644 --- a/go/cmd/dolt/cli/arg_parser_helpers.go +++ b/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 { diff --git a/go/cmd/dolt/commands/schcmds/update-tag.go b/go/cmd/dolt/commands/schcmds/update-tag.go index a4b8d72d08..ad1dc8cf39 100644 --- a/go/cmd/dolt/commands/schcmds/update-tag.go +++ b/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 { diff --git a/go/libraries/doltcore/sqle/dprocedures/dolt_update_column_tag.go b/go/libraries/doltcore/sqle/dprocedures/dolt_update_column_tag.go new file mode 100644 index 0000000000..2775590949 --- /dev/null +++ b/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 ") + } + + 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 +} diff --git a/go/libraries/doltcore/sqle/dprocedures/init.go b/go/libraries/doltcore/sqle/dprocedures/init.go index bc01e257b6..61aa99c562 100644 --- a/go/libraries/doltcore/sqle/dprocedures/init.go +++ b/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}, diff --git a/integration-tests/bats/sql-update-column-tag.bats b/integration-tests/bats/sql-update-column-tag.bats new file mode 100644 index 0000000000..f28170cec1 --- /dev/null +++ b/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 +} \ No newline at end of file