github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/libraries/doltcore/sqle/schema_table.go

// Copyright 2020 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 sqle

import (
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/dolthub/go-mysql-server/sql"
	gmstypes "github.com/dolthub/go-mysql-server/sql/types"
	"github.com/dolthub/vitess/go/vt/proto/query"

	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
	"github.com/dolthub/dolt/go/libraries/doltcore/schema/typeinfo"
)

const (
	viewFragment    = "view"
	triggerFragment = "trigger"
	eventFragment   = "event"
)

type Extra struct {
	CreatedAt int64
}

func mustNewColWithTypeInfo(name string, tag uint64, typeInfo typeinfo.TypeInfo, partOfPK bool, defaultVal string, autoIncrement bool, comment string, constraints ...schema.ColConstraint) schema.Column {
	col, err := schema.NewColumnWithTypeInfo(name, tag, typeInfo, partOfPK, defaultVal, autoIncrement, comment, constraints...)
	if err != nil {
		panic(err)
	}
	return col
}

func mustCreateStringType(baseType query.Type, length int64, collation sql.CollationID) sql.StringType {
	ti, err := gmstypes.CreateString(baseType, length, collation)
	if err != nil {
		panic(err)
	}
	return ti
}

// dolt_schemas columns
var schemasTableCols = schema.NewColCollection(
	mustNewColWithTypeInfo(doltdb.SchemasTablesTypeCol, schema.DoltSchemasTypeTag, typeinfo.CreateVarStringTypeFromSqlType(mustCreateStringType(query.Type_VARCHAR, 64, sql.Collation_utf8mb4_0900_ai_ci)), true, "", false, ""),
	mustNewColWithTypeInfo(doltdb.SchemasTablesNameCol, schema.DoltSchemasNameTag, typeinfo.CreateVarStringTypeFromSqlType(mustCreateStringType(query.Type_VARCHAR, 64, sql.Collation_utf8mb4_0900_ai_ci)), true, "", false, ""),
	mustNewColWithTypeInfo(doltdb.SchemasTablesFragmentCol, schema.DoltSchemasFragmentTag, typeinfo.CreateVarStringTypeFromSqlType(gmstypes.LongText), false, "", false, ""),
	mustNewColWithTypeInfo(doltdb.SchemasTablesExtraCol, schema.DoltSchemasExtraTag, typeinfo.JSONType, false, "", false, ""),
	mustNewColWithTypeInfo(doltdb.SchemasTablesSqlModeCol, schema.DoltSchemasSqlModeTag, typeinfo.CreateVarStringTypeFromSqlType(mustCreateStringType(query.Type_VARCHAR, 256, sql.Collation_utf8mb4_0900_ai_ci)), false, "", false, ""),
)

var schemaTableSchema = schema.MustSchemaFromCols(schemasTableCols)

// getOrCreateDoltSchemasTable returns the `dolt_schemas` table in `db`, creating it if it does not already exist.
// Also migrates data to the correct format if necessary.
func getOrCreateDoltSchemasTable(ctx *sql.Context, db Database) (retTbl *WritableDoltTable, retErr error) {
	tbl, found, err := db.GetTableInsensitive(ctx, doltdb.SchemasTableName)
	if err != nil {
		return nil, err
	}

	if found {
		schemasTable := tbl.(*WritableDoltTable)
		// Old schemas table contains the `id` column or is missing an `extra` column.
		if tbl.Schema().Contains(doltdb.SchemasTablesIdCol, doltdb.SchemasTableName) || !tbl.Schema().Contains(doltdb.SchemasTablesExtraCol, doltdb.SchemasTableName) {
			return migrateOldSchemasTableToNew(ctx, db, schemasTable)
		} else {
			return schemasTable, nil
		}
	}

	root, err := db.GetRoot(ctx)
	if err != nil {
		return nil, err
	}

	// Create new empty table
	err = db.createDoltTable(ctx, doltdb.SchemasTableName, doltdb.DefaultSchemaName, root, schemaTableSchema)
	if err != nil {
		return nil, err
	}
	tbl, found, err = db.GetTableInsensitive(ctx, doltdb.SchemasTableName)
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, sql.ErrTableNotFound.New(doltdb.SchemasTableName)
	}

	return tbl.(*WritableDoltTable), nil
}

func migrateOldSchemasTableToNew(ctx *sql.Context, db Database, schemasTable *WritableDoltTable) (newTable *WritableDoltTable, rerr error) {
	// Copy all of the old data over and add an index column and an extra column
	iter, err := SqlTableToRowIter(ctx, schemasTable.DoltTable, nil)
	if err != nil {
		return nil, err
	}

	// The dolt_schemas table has undergone various changes over time and multiple possible schemas for it exist, so we
	// need to get the column indexes from the current schema
	nameIdx := schemasTable.sqlSchema().IndexOfColName(doltdb.SchemasTablesNameCol)
	typeIdx := schemasTable.sqlSchema().IndexOfColName(doltdb.SchemasTablesTypeCol)
	fragmentIdx := schemasTable.sqlSchema().IndexOfColName(doltdb.SchemasTablesFragmentCol)
	extraIdx := schemasTable.sqlSchema().IndexOfColName(doltdb.SchemasTablesExtraCol)
	sqlModeIdx := schemasTable.sqlSchema().IndexOfColName(doltdb.SchemasTablesSqlModeCol)

	defer func(iter sql.RowIter, ctx *sql.Context) {
		err := iter.Close(ctx)
		if err != nil && rerr == nil {
			rerr = err
		}
	}(iter, ctx)

	var newRows []sql.Row
	for {
		sqlRow, err := iter.Next(ctx)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		newRow := make(sql.Row, schemasTableCols.Size())
		newRow[0] = sqlRow[typeIdx]
		newRow[1] = sqlRow[nameIdx]
		newRow[2] = sqlRow[fragmentIdx]
		if extraIdx >= 0 {
			newRow[3] = sqlRow[extraIdx]
		}
		if sqlModeIdx >= 0 {
			newRow[4] = sqlRow[sqlModeIdx]
		}

		newRows = append(newRows, newRow)
	}

	err = db.dropTable(ctx, doltdb.SchemasTableName)
	if err != nil {
		return nil, err
	}

	root, err := db.GetRoot(ctx)
	if err != nil {
		return nil, err
	}

	err = db.createDoltTable(ctx, doltdb.SchemasTableName, doltdb.DefaultSchemaName, root, schemaTableSchema)
	if err != nil {
		return nil, err
	}

	tbl, found, err := db.GetTableInsensitive(ctx, doltdb.SchemasTableName)
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, sql.ErrTableNotFound.New(doltdb.SchemasTableName)
	}

	inserter := tbl.(*WritableDoltTable).Inserter(ctx)
	for _, row := range newRows {
		err = inserter.Insert(ctx, row)
		if err != nil {
			return nil, err
		}
	}

	err = inserter.Close(ctx)
	if err != nil {
		return nil, err
	}

	return tbl.(*WritableDoltTable), nil
}

// fragFromSchemasTable returns the row with the given schema fragment if it exists.
func fragFromSchemasTable(ctx *sql.Context, tbl *WritableDoltTable, fragType string, name string) (r sql.Row, found bool, rerr error) {
	fragType, name = strings.ToLower(fragType), strings.ToLower(name)

	// This performs a full table scan in the worst case, but it's only used when adding or dropping a trigger or view
	iter, err := SqlTableToRowIter(ctx, tbl.DoltTable, nil)
	if err != nil {
		return nil, false, err
	}

	defer func(iter sql.RowIter, ctx *sql.Context) {
		err := iter.Close(ctx)
		if err != nil && rerr == nil {
			rerr = err
		}
	}(iter, ctx)

	// The dolt_schemas table has undergone various changes over time and multiple possible schemas for it exist, so we
	// need to get the column indexes from the current schema
	nameIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesNameCol)
	typeIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesTypeCol)

	for {
		sqlRow, err := iter.Next(ctx)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, false, err
		}

		// These columns are case insensitive, make sure to do a case-insensitive comparison
		if strings.ToLower(sqlRow[typeIdx].(string)) == fragType && strings.ToLower(sqlRow[nameIdx].(string)) == name {
			return sqlRow, true, nil
		}
	}

	return nil, false, nil
}

type schemaFragment struct {
	name     string
	fragment string
	created  time.Time
	// sqlMode indicates the SQL_MODE that was used when this schema fragment was initially parsed. SQL_MODE settings
	// such as ANSI_QUOTES control customized parsing behavior needed for some schema fragments.
	sqlMode string
}

func getSchemaFragmentsOfType(ctx *sql.Context, tbl *WritableDoltTable, fragType string) (sf []schemaFragment, rerr error) {
	iter, err := SqlTableToRowIter(ctx, tbl.DoltTable, nil)
	if err != nil {
		return nil, err
	}

	// The dolt_schemas table has undergone various changes over time and multiple possible schemas for it exist, so we
	// need to get the column indexes from the current schema
	nameIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesNameCol)
	typeIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesTypeCol)
	fragmentIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesFragmentCol)
	extraIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesExtraCol)
	sqlModeIdx := tbl.sqlSchema().IndexOfColName(doltdb.SchemasTablesSqlModeCol)

	defer func(iter sql.RowIter, ctx *sql.Context) {
		err := iter.Close(ctx)
		if err != nil && rerr == nil {
			rerr = err
		}
	}(iter, ctx)

	var frags []schemaFragment
	for {
		sqlRow, err := iter.Next(ctx)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		if sqlRow[typeIdx] != fragType {
			continue
		}

		sqlModeString := ""
		if sqlModeIdx >= 0 {
			if s, ok := sqlRow[sqlModeIdx].(string); ok {
				sqlModeString = s
			}
		} else {
			defaultSqlMode, err := loadDefaultSqlMode()
			if err != nil {
				return nil, err
			}
			sqlModeString = defaultSqlMode
		}

		// For older tables, use 1 as the trigger creation time
		if extraIdx < 0 || sqlRow[extraIdx] == nil {
			frags = append(frags, schemaFragment{
				name:     sqlRow[nameIdx].(string),
				fragment: sqlRow[fragmentIdx].(string),
				created:  time.Unix(1, 0).UTC(), // TablePlus editor thinks 0 is out of range
				sqlMode:  sqlModeString,
			})
			continue
		}

		// Extract Created Time from JSON column
		createdTime, err := getCreatedTime(ctx, sqlRow[extraIdx].(sql.JSONWrapper))
		if err != nil {
			return nil, err
		}

		frags = append(frags, schemaFragment{
			name:     sqlRow[nameIdx].(string),
			fragment: sqlRow[fragmentIdx].(string),
			created:  time.Unix(createdTime, 0).UTC(),
			sqlMode:  sqlModeString,
		})
	}

	return frags, nil
}

// loadDefaultSqlMode loads the default value for the @@SQL_MODE system variable and returns it, along
// with any unexpected errors encountered while reading the default value.
func loadDefaultSqlMode() (string, error) {
	global, _, ok := sql.SystemVariables.GetGlobal("SQL_MODE")
	if !ok {
		return "", fmt.Errorf("unable to load default @@SQL_MODE")
	}
	s, ok := global.GetDefault().(string)
	if !ok {
		return "", fmt.Errorf("unexpected type for @@SQL_MODE default value: %T", global.GetDefault())
	}
	return s, nil
}

func getCreatedTime(ctx *sql.Context, extraCol sql.JSONWrapper) (int64, error) {
	doc, err := extraCol.ToInterface()
	if err != nil {
		return 0, err
	}

	err = fmt.Errorf("value %v does not contain creation time", doc)

	obj, ok := doc.(map[string]interface{})
	if !ok {
		return 0, err
	}

	v, ok := obj["CreatedAt"]
	if !ok {
		return 0, err
	}

	f, ok := v.(float64)
	if !ok {
		return 0, err
	}
	return int64(f), nil
}