Merge pull request #1169 from dolthub/fulghum/udf-alias

PL/pgSQL: Add support for the `ALIAS` statement
dolthub · Feb 5, 2025 · 30130dc · 30130dc
2 parents 9c3900c + ec8bba1
commit 30130dc
Show file tree

Hide file tree

Showing 5 changed files with 108 additions and 2 deletions.
diff --git a/server/functions/framework/interpreter_logic.go b/server/functions/framework/interpreter_logic.go
@@ -48,6 +48,12 @@ func (iFunc InterpretedFunction) Call(ctx *sql.Context, runner analyzer.Statemen
 
 		operation := iFunc.Statements[counter]
 		switch operation.OpCode {
+		case OpCode_Alias:
+			iv := stack.GetVariable(operation.PrimaryData)
+			if iv == nil {
+				return nil, fmt.Errorf("variable `%s` could not be found", operation.PrimaryData)
+			}
+			stack.NewVariableAlias(operation.Target, iv)
 		case OpCode_Assign:
 			iv := stack.GetVariable(operation.Target)
 			if iv == nil {

diff --git a/server/functions/framework/interpreter_operation.go b/server/functions/framework/interpreter_operation.go
@@ -19,7 +19,8 @@ package framework
 type OpCode uint16
 
 const (
-	OpCode_Assign     OpCode = iota // https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-ASSIGNMENT
+	OpCode_Alias      OpCode = iota // https://www.postgresql.org/docs/15/plpgsql-declarations.html#PLPGSQL-DECLARATION-ALIAS
+	OpCode_Assign                   // https://www.postgresql.org/docs/15/plpgsql-statements.html#PLPGSQL-STATEMENTS-ASSIGNMENT
 	OpCode_Case                     // https://www.postgresql.org/docs/15/plpgsql-control-structures.html#PLPGSQL-CONDITIONALS
 	OpCode_Declare                  // https://www.postgresql.org/docs/15/plpgsql-declarations.html
 	OpCode_DeleteInto               // https://www.postgresql.org/docs/15/plpgsql-statements.html

diff --git a/server/functions/framework/interpreter_stack.go b/server/functions/framework/interpreter_stack.go
@@ -86,6 +86,12 @@ func (is *InterpreterStack) NewVariableWithValue(name string, typ *pgtypes.Doltg
 	}
 }
 
+// NewVariableAlias creates a new variable alias, named |alias|, in the current frame of this stack,
+// pointing to the specified |variable|.
+func (is *InterpreterStack) NewVariableAlias(alias string, variable *InterpreterVariable) {
+	is.stack.Peek().variables[alias] = variable
+}
+
 // PushScope creates a new scope.
 func (is *InterpreterStack) PushScope() {
 	is.stack.Push(&InterpreterScopeDetails{

diff --git a/server/functions/interpreted_examples.go b/server/functions/interpreted_examples.go
@@ -24,6 +24,7 @@ import (
 // interpreter functionality.
 func initInterpretedExamples() {
 	framework.RegisterFunction(interpretedAssignment)
+	framework.RegisterFunction(interpretedAlias)
 }
 
 // interpretedAssignment is roughly equivalent to the (expected) parsed output of the following function definition:
@@ -151,3 +152,86 @@ var interpretedAssignment = framework.InterpretedFunction{
 		},
 	},
 }
+
+// interpretedAlias is roughly equivalent to the (expected) parsed output of the following function definition:
+/*
+CREATE FUNCTION interpreted_alias(input TEXT)
+RETURNS TEXT AS $$
+DECLARE
+    var1 TEXT;
+	var2 TEXT;
+BEGIN
+    DECLARE
+		alias1 ALIAS FOR var1;
+		alias2 ALIAS FOR alias1;
+		alias3 ALIAS FOR input;
+	BEGIN
+        alias2 := alias3;
+	END;
+	RETURN var1;
+END;
+$$ LANGUAGE plpgsql;
+*/
+var interpretedAlias = framework.InterpretedFunction{
+	ID:                 id.NewFunction("pg_catalog", "interpreted_alias", pgtypes.Text.ID),
+	ReturnType:         pgtypes.Text,
+	ParameterNames:     []string{"input"},
+	ParameterTypes:     []*pgtypes.DoltgresType{pgtypes.Text},
+	Variadic:           false,
+	IsNonDeterministic: false,
+	Strict:             true,
+	Labels:             nil,
+	Statements: []framework.InterpreterOperation{
+		{ // 0
+			OpCode: framework.OpCode_ScopeBegin,
+		},
+		{ // 1
+			OpCode:      framework.OpCode_Declare,
+			Target:      `var1`,
+			PrimaryData: `text`,
+		},
+		{ // 2
+			OpCode:      framework.OpCode_Declare,
+			Target:      `var2`,
+			PrimaryData: `text`,
+		},
+		{ // 3
+			OpCode: framework.OpCode_ScopeBegin,
+		},
+		{ // 4
+			OpCode:      framework.OpCode_Alias,
+			Target:      `alias1`,
+			PrimaryData: `var1`,
+		},
+		{ // 5
+			OpCode:      framework.OpCode_Alias,
+			Target:      `alias2`,
+			PrimaryData: `alias1`,
+		},
+		{ // 6
+			OpCode:      framework.OpCode_Alias,
+			Target:      `alias3`,
+			PrimaryData: `input`,
+		},
+		{ // 7
+			OpCode: framework.OpCode_ScopeBegin,
+		},
+		{ // 8
+			OpCode:        framework.OpCode_Assign,
+			PrimaryData:   `SELECT $1;`,
+			SecondaryData: []string{`alias3`},
+			Target:        `alias2`,
+		},
+		{ // 9
+			OpCode: framework.OpCode_ScopeEnd,
+		},
+		{ // 10
+			OpCode:        framework.OpCode_Return,
+			PrimaryData:   `SELECT $1;`,
+			SecondaryData: []string{`var1`},
+		},
+		{ // 11
+			OpCode: framework.OpCode_ScopeEnd,
+		},
+	},
+}
diff --git a/testing/go/create_function_test.go b/testing/go/create_function_test.go
@@ -23,7 +23,7 @@ import (
 func TestCreateFunction(t *testing.T) {
 	RunScripts(t, []ScriptTest{
 		{
-			Name: "Interpreter Example",
+			Name: "Interpreter Assignment Example",
 			Assertions: []ScriptTestAssertion{
 				{
 					Query: "SELECT interpreted_assignment('Hello');",
@@ -51,5 +51,14 @@ func TestCreateFunction(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name: "Interpreter Alias Example",
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:    "SELECT interpreted_alias('123');",
+					Expected: []sql.Row{{"123"}},
+				},
+			},
+		},
 	})
 }