support AGGREGATE statement syntaxes (#152)

postgres/parser/sem/tree/create_aggregate.go

@@ -0,0 +1,201 @@
+// 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 tree
+
+import "strings"
+
+var _ Statement = &CreateAggregate{}
+
+// CreateAggregate represents a CREATE AGGREGATE statement.
+type CreateAggregate struct {
+	Name        Name
+	Replace     bool
+	Args        RoutineArgs
+	SFunc       string
+	SType       ResolvableTypeReference
+	AggOptions  CreateAggOptions
+	OrderByArgs RoutineArgs
+	BaseType    ResolvableTypeReference
+}
+
+// Format implements the NodeFormatter interface.
+func (node *CreateAggregate) Format(ctx *FmtCtx) {
+	ctx.WriteString("CREATE ")
+	if node.Replace {
+		ctx.WriteString("OR REPLACE ")
+	}
+	ctx.WriteString("AGGREGATE ")
+	ctx.FormatNode(&node.Name)
+	ctx.WriteString(" ( ")
+	if node.OrderByArgs != nil {
+		if node.Args != nil {
+			ctx.FormatNode(node.Args)
+		}
+		ctx.WriteString(" ORDER BY ")
+		ctx.FormatNode(node.OrderByArgs)
+		ctx.WriteString(" ) ( ")
+	} else if node.Args != nil {
+		ctx.FormatNode(node.Args)
+		ctx.WriteString(" ) ( ")
+	} else {
+		ctx.WriteString("BASETYPE = ")
+		ctx.WriteString(node.BaseType.SQLString())
+		ctx.WriteString(" , ")
+	}
+	ctx.WriteString("SFUNC = ")
+	ctx.WriteString(node.SFunc)
+	ctx.WriteString(" , STYPE = ")
+	ctx.WriteString(node.BaseType.SQLString())
+	if node.AggOptions != nil {
+		ctx.FormatNode(&node.AggOptions)
+	}
+}
+
+type FinalFuncModifyType string
+
+const (
+	FinalFuncModifyReadOnly  FinalFuncModifyType = "READ_ONLY"
+	FinalFuncModifyShareable FinalFuncModifyType = "SHAREABLE"
+	FinalFuncModifyReadWrite FinalFuncModifyType = "READ_WRITE"
+)
+
+type CreateAggOptions []CreateAggOption
+
+// Format implements the NodeFormatter interface.
+func (node *CreateAggOptions) Format(ctx *FmtCtx) {
+	for _, option := range *node {
+		ctx.WriteString(" , ")
+		ctx.FormatNode(&option)
+	}
+}
+
+type CreateAggOption struct {
+	Option CreateOptionType
+	// IntVal is used for SSpace and MSSpace
+	IntVal Expr
+	// CondVal is used for InitCond and MInitCond
+	CondVal Expr
+	// BoolVal is used for FinalFuncExtra and MFinalFuncExtra
+	BoolVal bool
+	// StrVal is used for FinalFunc, CombineFunc, SerialFunc,
+	// DeserialFunc, MSFunc, MInvFunc and MFinalFunc
+	StrVal string
+	// FinalFuncModify is used for FinalFuncModify and MFinalFuncModify
+	FinalFuncModify FinalFuncModifyType
+	// TypeVal is used for MSType
+	TypeVal  ResolvableTypeReference
+	Parallel Parallel
+	SortOp   Operator
+	// Hypothetical does not define any stored value.
+}
+
+// Format implements the NodeFormatter interface.
+func (node *CreateAggOption) Format(ctx *FmtCtx) {
+	switch node.Option {
+	case AggOptTypeSSpace:
+		ctx.WriteString("SSPACE = ")
+		ctx.FormatNode(node.IntVal)
+	case AggOptTypeFinalFunc:
+		ctx.WriteString("FINALFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeFinalFuncExtra:
+		if node.BoolVal {
+			ctx.WriteString("FINALFUNC_EXTRA = TRUE")
+		} else {
+			ctx.WriteString("FINALFUNC_EXTRA = FALSE")
+		}
+	case AggOptTypeFinalFuncModify:
+		ctx.WriteString("FINALFUNC_MODIFY")
+		ctx.WriteString(string(node.FinalFuncModify))
+	case AggOptTypeCombineFunc:
+		ctx.WriteString("COMBINEFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeSerialFunc:
+		ctx.WriteString("SERIALFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeDeserialFunc:
+		ctx.WriteString("DESERIALFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeInitCond:
+		ctx.WriteString("INITCOND = ")
+		ctx.FormatNode(node.CondVal)
+	case AggOptTypeMSFunc:
+		ctx.WriteString("MSFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeMInvFunc:
+		ctx.WriteString("MINVFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeMSType:
+		ctx.WriteString("MSTYPE = ")
+		ctx.WriteString(node.TypeVal.SQLString())
+	case AggOptTypeMSSpace:
+		ctx.WriteString("MSSPACE = ")
+		ctx.FormatNode(node.IntVal)
+	case AggOptTypeMFinalFunc:
+		ctx.WriteString("MFINALFUNC = ")
+		ctx.WriteString(node.StrVal)
+	case AggOptTypeMFinalFuncExtra:
+		if node.BoolVal {
+			ctx.WriteString("MFINALFUNC_EXTRA = TRUE")
+		} else {
+			ctx.WriteString("MFINALFUNC_EXTRA = FALSE")
+		}
+	case AggOptTypeMFinalFuncModify:
+		ctx.WriteString("MFINALFUNC_MODIFY")
+		ctx.WriteString(string(node.FinalFuncModify))
+	case AggOptTypeMInitCond:
+		ctx.WriteString("MINITCOND = ")
+		ctx.FormatNode(node.CondVal)
+	case AggOptTypeSortOp:
+		ctx.WriteString("SORTOP = ")
+		switch op := node.SortOp.(type) {
+		case UnaryOperator:
+			ctx.WriteString(op.String())
+		case BinaryOperator:
+			ctx.WriteString(op.String())
+		case ComparisonOperator:
+			ctx.WriteString(op.String())
+		}
+	case AggOptTypeParallel:
+		ctx.WriteString("PARALLEL = ")
+		ctx.WriteString(strings.ToUpper(string(node.Parallel)))
+	case AggOptTypeHypothetical:
+		ctx.WriteString("HYPOTHETICAL")
+	}
+}
+
+type CreateOptionType int
+
+const (
+	AggOptTypeSSpace CreateOptionType = iota
+	AggOptTypeFinalFunc
+	AggOptTypeFinalFuncExtra
+	AggOptTypeFinalFuncModify
+	AggOptTypeCombineFunc
+	AggOptTypeSerialFunc
+	AggOptTypeDeserialFunc
+	AggOptTypeInitCond
+	AggOptTypeMSFunc
+	AggOptTypeMInvFunc
+	AggOptTypeMSType
+	AggOptTypeMSSpace
+	AggOptTypeMFinalFunc
+	AggOptTypeMFinalFuncExtra
+	AggOptTypeMFinalFuncModify
+	AggOptTypeMInitCond
+	AggOptTypeSortOp
+	AggOptTypeParallel
+	AggOptTypeHypothetical
+)

postgres/parser/sem/tree/drop.go

@@ -53,6 +53,43 @@ func (d DropBehavior) String() string {
 	return dropBehaviorName[d]
 }
 
+var _ Statement = &DropAggregate{}
+
+// DropAggregate represents a DROP AGGREGATE statement.
+type DropAggregate struct {
+	Aggregates   []AggregateToDrop
+	IfExists     bool
+	DropBehavior DropBehavior
+}
+
+type AggregateToDrop struct {
+	Name   Name
+	AggSig *AggregateSignature
+}
+
+// Format implements the NodeFormatter interface.
+func (node *DropAggregate) Format(ctx *FmtCtx) {
+	ctx.WriteString("DROP AGGREGATE ")
+	if node.IfExists {
+		ctx.WriteString("IF EXISTS ")
+	}
+	for i, agg := range node.Aggregates {
+		if i != 0 {
+			ctx.WriteString(" , ")
+		}
+		ctx.FormatNode(&agg.Name)
+		ctx.WriteString(" ( ")
+		ctx.FormatNode(agg.AggSig)
+		ctx.WriteString(" ) ")
+	}
+	switch node.DropBehavior {
+	case DropDefault:
+	default:
+		ctx.WriteByte(' ')
+		ctx.WriteString(dropBehaviorName[node.DropBehavior])
+	}
+}
+
 var _ Statement = &DropDatabase{}
 
 // DropDatabase represents a DROP DATABASE statement.

postgres/parser/sem/tree/stmt.go

@@ -482,6 +482,12 @@ func (*CopyFrom) StatementType() StatementType { return CopyIn }
 // StatementTag returns a short string identifying the type of statement.
 func (*CopyFrom) StatementTag() string { return "COPY" }
 
+// StatementType implements the Statement interface.
+func (*CreateAggregate) StatementType() StatementType { return DDL }
+
+// StatementTag returns a short string identifying the type of statement.
+func (*CreateAggregate) StatementTag() string { return "CREATE AGGREGATE" }
+
 // StatementType implements the Statement interface.
 func (*CreateChangefeed) StatementType() StatementType { return Rows }
 
@@ -634,6 +640,12 @@ func (n *Delete) StatementType() StatementType { return n.Returning.statementTyp
 // StatementTag returns a short string identifying the type of statement.
 func (*Delete) StatementTag() string { return "DELETE" }
 
+// StatementType implements the Statement interface.
+func (*DropAggregate) StatementType() StatementType { return DDL }
+
+// StatementTag returns a short string identifying the type of statement.
+func (*DropAggregate) StatementTag() string { return "DROP AGGREGATE" }
+
 // StatementType implements the Statement interface.
 func (*DropDatabase) StatementType() StatementType { return DDL }
 
@@ -1232,6 +1244,7 @@ func (n *CommentOnIndex) String() string            { return AsString(n) }
 func (n *CommentOnTable) String() string            { return AsString(n) }
 func (n *CommitTransaction) String() string         { return AsString(n) }
 func (n *CopyFrom) String() string                  { return AsString(n) }
+func (n *CreateAggregate) String() string           { return AsString(n) }
 func (n *CreateChangefeed) String() string          { return AsString(n) }
 func (n *CreateDatabase) String() string            { return AsString(n) }
 func (n *CreateDomain) String() string              { return AsString(n) }
@@ -1251,6 +1264,7 @@ func (n *CreateStats) String() string               { return AsString(n) }
 func (n *CreateView) String() string                { return AsString(n) }
 func (n *Deallocate) String() string                { return AsString(n) }
 func (n *Delete) String() string                    { return AsString(n) }
+func (n *DropAggregate) String() string             { return AsString(n) }
 func (n *DropDatabase) String() string              { return AsString(n) }
 func (n *DropDomain) String() string                { return AsString(n) }
 func (n *DropExtension) String() string             { return AsString(n) }

testing/generation/command_docs/custom_variables.go

@@ -53,6 +53,7 @@ var GlobalCustomVariables = map[string]utils.StatementGenerator{
 	"loid":                customDefinition(`99999`),
 	"maxvalue":            customDefinition(`1`),
 	"minvalue":            customDefinition(`1`),
+	"mstate_data_size":    customDefinition(`16`),
 	"neighbor_enum_value": customDefinition(`'1'`),
 	"new_enum_value":      customDefinition(`'1'`),
 	"numeric_literal":     customDefinition(`1`),
@@ -65,10 +66,12 @@ var GlobalCustomVariables = map[string]utils.StatementGenerator{
 	"select":              customDefinition(`SELECT 1`),
 	"sequence_options":    customDefinition(`NO MINVALUE`),
 	"snapshot_id":         customDefinition(`'snapshot_id'`),
+	"sort_operator":       customDefinition(`>`),
 	"source_encoding":     customDefinition(`'UTF8'`),
 	"source_query":        customDefinition(`SELECT 1`),
 	"sql_body":            customDefinition(`BEGIN ATOMIC END | RETURN 1`),
 	"start":               customDefinition(`0`),
+	"state_data_size":     customDefinition(`16`),
 	"storage_parameter":   customDefinition(`fillfactor`),
 	"strategy_number":     customDefinition(`3`),
 	"string_literal":      customDefinition(`'str'`),