support CREATE|ALTER|DROP PROCEDURE and ALTER FUNCTION

postgres/parser/sem/tree/alter_function.go

@@ -0,0 +1,77 @@
+// 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.
+
+// Copyright 2020 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package tree
+
+var _ Statement = &AlterFunction{}
+
+// AlterFunction represents a ALTER FUNCTION statement.
+type AlterFunction struct {
+	Name      *UnresolvedObjectName
+	Args      RoutineArgs
+	Options   []RoutineOption
+	Restrict  bool
+	Rename    *UnresolvedObjectName
+	Owner     string
+	Schema    string
+	No        bool
+	Extension string
+}
+
+// Format implements the NodeFormatter interface.
+func (node *AlterFunction) Format(ctx *FmtCtx) {
+	ctx.WriteString("ALTER FUNCTION ")
+	ctx.FormatNode(node.Name)
+	if node.Args != nil {
+		ctx.WriteString(" ( ")
+		ctx.FormatNode(node.Name)
+		ctx.WriteString(" )")
+	}
+	if node.Options != nil {
+		for i, option := range node.Options {
+			if i != 0 {
+				ctx.WriteByte(' ')
+			}
+			ctx.FormatNode(option)
+		}
+		if node.Restrict {
+			ctx.WriteString(" RESTRICT")
+		}
+	} else if node.Rename != nil {
+		ctx.WriteString(" RENAME TO ")
+		ctx.FormatNode(node.Rename)
+	} else if node.Owner != "" {
+		ctx.WriteString(" OWNER TO ")
+		ctx.WriteString(node.Owner)
+	} else if node.Schema != "" {
+		ctx.WriteString(" SET SCHEMA ")
+		ctx.WriteString(node.Schema)
+	} else {
+		if node.No {
+			ctx.WriteString(" NO")
+		}
+		ctx.WriteString(" DEPENDS ON EXTENSION ")
+		ctx.WriteString(node.Extension)
+	}
+}

postgres/parser/sem/tree/alter_procedure.go

@@ -0,0 +1,77 @@
+// 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.
+
+// Copyright 2020 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package tree
+
+var _ Statement = &AlterProcedure{}
+
+// AlterProcedure represents a ALTER PROCEDURE statement.
+type AlterProcedure struct {
+	Name      *UnresolvedObjectName
+	Args      RoutineArgs
+	Options   []RoutineOption
+	Restrict  bool
+	Rename    *UnresolvedObjectName
+	Owner     string
+	Schema    string
+	No        bool
+	Extension string
+}
+
+// Format implements the NodeFormatter interface.
+func (node *AlterProcedure) Format(ctx *FmtCtx) {
+	ctx.WriteString("ALTER PROCEDURE ")
+	ctx.FormatNode(node.Name)
+	if node.Args != nil {
+		ctx.WriteString(" ( ")
+		ctx.FormatNode(node.Name)
+		ctx.WriteString(" )")
+	}
+	if node.Options != nil {
+		for i, option := range node.Options {
+			if i != 0 {
+				ctx.WriteByte(' ')
+			}
+			ctx.FormatNode(option)
+		}
+		if node.Restrict {
+			ctx.WriteString(" RESTRICT")
+		}
+	} else if node.Rename != nil {
+		ctx.WriteString(" RENAME TO ")
+		ctx.FormatNode(node.Rename)
+	} else if node.Owner != "" {
+		ctx.WriteString(" OWNER TO ")
+		ctx.WriteString(node.Owner)
+	} else if node.Schema != "" {
+		ctx.WriteString(" SET SCHEMA ")
+		ctx.WriteString(node.Schema)
+	} else {
+		if node.No {
+			ctx.WriteString(" NO")
+		}
+		ctx.WriteString(" DEPENDS ON EXTENSION ")
+		ctx.WriteString(node.Extension)
+	}
+}

postgres/parser/sem/tree/create_function.go

@@ -46,8 +46,9 @@ func (node *CreateFunction) Format(ctx *FmtCtx) {
 	ctx.WriteString("FUNCTION ")
 	ctx.FormatNode(node.Name)
 	if len(node.Args) != 0 {
-		ctx.WriteByte(' ')
+		ctx.WriteString(" (")
 		ctx.FormatNode(node.Args)
+		ctx.WriteString(" )")
 	}
 	if node.RetType != nil {
 		if len(node.RetType) == 1 && node.RetType[0].Name == "" {
@@ -124,6 +125,7 @@ const (
 	OptionAs1
 	OptionAs2
 	OptionSqlBody
+	OptionReset // For ALTER { FUNCTION | PROCEDURE } use only
 )
 
 type RoutineOption struct {
@@ -145,6 +147,10 @@ type RoutineOption struct {
 	ObjFile        string
 	LinkSymbol     string
 	SqlBody        Statement
+
+	// For ALTER { FUNCTION | PROCEDURE } use only
+	ResetParam string
+	ResetAll   bool
 }
 
 // Format implements the NodeFormatter interface.
@@ -216,6 +222,13 @@ func (node RoutineOption) Format(ctx *FmtCtx) {
 		ctx.WriteString(node.LinkSymbol)
 	case OptionSqlBody:
 		ctx.FormatNode(node.SqlBody)
+	case OptionReset:
+		ctx.WriteString("RESET ")
+		if node.ResetAll {
+			ctx.WriteString("ALL")
+		} else {
+			ctx.WriteString(node.ResetParam)
+		}
 	}
 }
 

postgres/parser/sem/tree/create_procedure.go

@@ -0,0 +1,65 @@
+// 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.
+
+// Copyright 2012, Google Inc. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in licenses/BSD-vitess.txt.
+
+// Portions of this file are additionally subject to the following
+// license and copyright.
+//
+// Copyright 2015 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+// This code was derived from https://github.com/youtube/vitess.
+
+package tree
+
+var _ Statement = &CreateProcedure{}
+
+// CreateProcedure represents a CREATE PROCEDURE statement.
+type CreateProcedure struct {
+	Name    *UnresolvedObjectName
+	Replace bool
+	Args    RoutineArgs
+	Options []RoutineOption
+}
+
+// Format implements the NodeFormatter interface.
+func (node *CreateProcedure) Format(ctx *FmtCtx) {
+	ctx.WriteString("CREATE ")
+	if node.Replace {
+		ctx.WriteString("OR REPLACE ")
+	}
+	ctx.WriteString("PROCEDURE ")
+	ctx.FormatNode(node.Name)
+	if len(node.Args) != 0 {
+		ctx.WriteString(" (")
+		ctx.FormatNode(node.Args)
+		ctx.WriteString(" )")
+	}
+	for i, option := range node.Options {
+		if i != 0 {
+			ctx.WriteByte(' ')
+		}
+		ctx.FormatNode(option)
+	}
+}

postgres/parser/sem/tree/drop.go

@@ -102,7 +102,7 @@ var _ Statement = &DropFunction{}
 
 // DropFunction represents a DROP FUNCTION statement.
 type DropFunction struct {
-	Functions    []FunctionWithArgs
+	Functions    []RoutineWithArgs
 	IfExists     bool
 	DropBehavior DropBehavior
 }
@@ -127,13 +127,13 @@ func (node *DropFunction) Format(ctx *FmtCtx) {
 	}
 }
 
-// FunctionWithArgs represents the function name and its arguments, if any, for DROP FUNCTION statement.
-type FunctionWithArgs struct {
+// RoutineWithArgs represents the routine name and its arguments, if any, for DROP { FUNCTION | PROCEDURE } statement.
+type RoutineWithArgs struct {
 	Name *UnresolvedObjectName
 	Args RoutineArgs
 }
 
-func (node *FunctionWithArgs) Format(ctx *FmtCtx) {
+func (node *RoutineWithArgs) Format(ctx *FmtCtx) {
 	ctx.FormatNode(node.Name)
 	ctx.WriteString(" (")
 	if len(node.Args) != 0 {
@@ -168,6 +168,35 @@ func (node *DropIndex) Format(ctx *FmtCtx) {
 	}
 }
 
+var _ Statement = &DropProcedure{}
+
+// DropProcedure represents a DROP PROCEDURE statement.
+type DropProcedure struct {
+	Procedures   []RoutineWithArgs
+	IfExists     bool
+	DropBehavior DropBehavior
+}
+
+// Format implements the NodeFormatter interface.
+func (node *DropProcedure) Format(ctx *FmtCtx) {
+	ctx.WriteString("DROP PROCEDURE ")
+	if node.IfExists {
+		ctx.WriteString("IF EXISTS ")
+	}
+	for i, f := range node.Procedures {
+		if i != 0 {
+			ctx.WriteString(", ")
+		}
+		ctx.FormatNode(&f)
+	}
+	switch node.DropBehavior {
+	case DropDefault:
+	default:
+		ctx.WriteByte(' ')
+		ctx.WriteString(dropBehaviorName[node.DropBehavior])
+	}
+}
+
 var _ Statement = &DropTable{}
 
 // DropTable represents a DROP TABLE statement.