dolthub--dolt/go/libraries/doltcore/sqle/expranalysis/expranalysis.go

// Copyright 2024 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 expranalysis

import (
	"fmt"

	"github.com/dolthub/go-mysql-server/memory"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/analyzer"
	"github.com/dolthub/go-mysql-server/sql/expression"
	"github.com/dolthub/go-mysql-server/sql/plan"
	"github.com/dolthub/go-mysql-server/sql/planbuilder"
	"github.com/dolthub/go-mysql-server/sql/transform"

	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/overrides"
	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/sqlfmt"
)

// ResolveDefaultExpression returns a sql.Expression for the column default or generated expression for the
// column provided
func ResolveDefaultExpression(ctx *sql.Context, tableName string, sch schema.Schema, col schema.Column) (sql.Expression, error) {
	ct, err := parseCreateTable(ctx, tableName, sch)
	if err != nil {
		return nil, err
	}

	colIdx := ct.PkSchema().Schema.IndexOfColName(col.Name)
	if colIdx < 0 {
		return nil, fmt.Errorf("unable to find column %s in analyzed query", col.Name)
	}

	sqlCol := ct.PkSchema().Schema[colIdx]
	expr := sqlCol.Default
	if expr == nil || expr.Expr == nil {
		expr = sqlCol.Generated
	}

	if expr == nil || expr.Expr == nil {
		return nil, fmt.Errorf("unable to find default or generated expression")
	}

	return expr, nil
}

// ResolveCheckExpression returns a sql.Expression for the check provided
func ResolveCheckExpression(ctx *sql.Context, tableName string, sch schema.Schema, checkExpr string) (sql.Expression, error) {
	ct, err := parseCreateTable(ctx, tableName, sch)
	if err != nil {
		return nil, err
	}

	formatter := overrides.SchemaFormatterFromContext(ctx)
	for _, check := range ct.Checks() {
		// Check definitions created before v1.55.3 may not have backquotes around identifiers
		quotedExpr := stripTableNamesFromExpression(ctx, formatter, check.Expr, true).String()
		unquotedExpr := stripTableNamesFromExpression(ctx, formatter, check.Expr, false).String()
		if quotedExpr == checkExpr || unquotedExpr == checkExpr {
			return check.Expr, nil
		}
	}

	return nil, fmt.Errorf("unable to find check expression")
}

// ResolveExpression compiles a predicate or check constraint string into a sql.Expression resolved.
// Used to re-hydrate partial index predicates and check expressions from their string representation.
// It uses SELECT statement on the expression FROM given table.
func ResolveExpression(ctx *sql.Context, tableName string, predicateStr string) (sql.Expression, error) {
	// TODO: added * in SELECT stmt to avoid pruning columns during analyzer.
	parsed, err := parseQuery(ctx, fmt.Sprintf("SELECT *, %s FROM %s", predicateStr, tableName))
	if err != nil {
		return nil, err
	}

	sess, ok := ctx.Session.(SessionDbProvider)
	if ok {
		// field index in GetField is incorrect, it needs `fix_exec_indexes` analyzer rule
		a := analyzer.NewDefault(sess.GenericProvider())
		parsed, err = a.Analyze(ctx, parsed, nil, nil)
		if err != nil {
			return nil, err
		}
		// comes out in ContextRootFinalizer, so unwrap it
		parsed = parsed.Children()[0]
	}

	proj, ok := parsed.(*plan.Project)
	if !ok {
		return nil, fmt.Errorf("expected *plan.Project parsing predicate, got %T", parsed)
	}
	expr := proj.Projections[len(proj.Projections)-1]
	if alias, ok := expr.(*expression.Alias); ok {
		expr = alias.Child
	}
	return expr, nil
}

func stripTableNamesFromExpression(ctx *sql.Context, formatter sql.SchemaFormatter, expr sql.Expression, quoted bool) sql.Expression {
	e, _, _ := transform.Expr(ctx, expr, func(ctx *sql.Context, e sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
		if col, ok := e.(*expression.GetField); ok {
			return col.WithTable("").WithQuotedNames(formatter, quoted), transform.NewTree, nil
		}
		return e, transform.SameTree, nil
	})
	return e
}

// SessionDbProvider is a way to access the Dolt database provider without import cycles.
type SessionDbProvider interface {
	GenericProvider() sql.MutableDatabaseProvider
}

func parseCreateTable(ctx *sql.Context, tableName string, sch schema.Schema) (*plan.CreateTable, error) {
	createTable, err := sqlfmt.GenerateCreateTableStatement(ctx, tableName, sch, nil, nil)
	if err != nil {
		return nil, err
	}

	pseudoAnalyzedQuery, err := parseQuery(ctx, createTable)
	if err != nil {
		return nil, err
	}

	ct, ok := pseudoAnalyzedQuery.(*plan.CreateTable)
	if !ok {
		return nil, fmt.Errorf("expected a *plan.CreateTable node, but got %T", pseudoAnalyzedQuery)
	}
	return ct, nil
}

func parseQuery(ctx *sql.Context, query string) (sql.Node, error) {
	// Doltgres must use the existing provider due to custom functions, so we split the logic here depending on whether
	// this is being called from a query context (which will have a valid Dolt session) or some other context (which
	// will need to construct a new provider).
	var b *planbuilder.Builder
	engOverrides := overrides.EngineOverridesFromContext(ctx)
	sess, ok := ctx.Session.(SessionDbProvider)
	if ok {
		catalog := analyzer.NewCatalog(sess.GenericProvider(), engOverrides)
		catalog.AuthHandler = sql.NoopAuthorizationHandler{}
		b = planbuilder.New(ctx, catalog, nil)
	} else {
		mockDatabase := memory.NewDatabase("mydb")
		mockProvider := memory.NewDBProvider(mockDatabase)
		catalog := analyzer.NewCatalog(mockProvider, engOverrides)
		catalog.AuthHandler = sql.NoopAuthorizationHandler{}
		// We need a new context for this operation
		parseCtx := sql.NewEmptyContext()
		parseCtx.SetCurrentDatabase("mydb")

		b = planbuilder.New(parseCtx, catalog, nil)
	}

	pseudoAnalyzedQuery, _, _, _, err := b.Parse(query, nil, false)
	if err != nil {
		return nil, err
	}
	return pseudoAnalyzedQuery, nil
}