cyrilgdn · AndrewJackson2020 · May 20, 2024 · Dec 24, 2024
diff --git a/postgresql/model_pg_procedure.go b/postgresql/model_pg_procedure.go
@@ -0,0 +1,139 @@
+package postgresql
+
+import (
+	"strings"
+
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+)
+
+// PGProcedure is the model for the database procedure
+type PGProcedure struct {
+	Schema          string
+	Name            string
+	Language        string
+	Body            string
+	Args            []PGProcedureArg
+	SecurityDefiner bool
+}
+
+type PGProcedureArg struct {
+	Name    string
+	Type    string
+	Mode    string
+	Default string
+}
+
+func (pgProcedure *PGProcedure) FromResourceData(d *schema.ResourceData) error {
+
+	if v, ok := d.GetOk(funcSchemaAttr); ok {
+		pgProcedure.Schema = v.(string)
+	} else {
+		pgProcedure.Schema = "public"
+	}
+
+	pgProcedure.Name = d.Get(funcNameAttr).(string)
+	if v, ok := d.GetOk(funcLanguageAttr); ok {
+		pgProcedure.Language = v.(string)
+	} else {
+		pgProcedure.Language = "plpgsql"
+	}
+	pgProcedure.Body = normalizeFunctionBody(d.Get(funcBodyAttr).(string))
+	pgProcedure.Args = []PGProcedureArg{}
+
+	if v, ok := d.GetOk(funcSecurityDefinerAttr); ok {
+		pgProcedure.SecurityDefiner = v.(bool)
+	} else {
+		pgProcedure.SecurityDefiner = false
+	}
+
+	if args, ok := d.GetOk(funcArgAttr); ok {
+		args := args.([]interface{})
+
+		for _, arg := range args {
+			arg := arg.(map[string]interface{})
+
+			var pgArg PGProcedureArg
+
+			if v, ok := arg[funcArgModeAttr]; ok {
+				pgArg.Mode = v.(string)
+			}
+
+			if v, ok := arg[funcArgNameAttr]; ok {
+				pgArg.Name = v.(string)
+			}
+
+			pgArg.Type = arg[funcArgTypeAttr].(string)
+
+			if v, ok := arg[funcArgDefaultAttr]; ok {
+				pgArg.Default = v.(string)
+			}
+
+			pgProcedure.Args = append(pgProcedure.Args, pgArg)
+		}
+	}
+
+	return nil
+}
+
+func (pgProcedure *PGProcedure) Parse(functionDefinition string) error {
+
+	pgProcedureData := findStringSubmatchMap(
+		`(?si)CREATE\sOR\sREPLACE\sPROCEDURE\s(?P<Schema>[^.]+)\.(?P<Name>[^(]+)\((?P<Args>.*)\).*LANGUAGE\s(?P<Language>[^\n\s]+)\s*(?P<Security>(SECURITY DEFINER)?).*\$[a-zA-Z]*\$(?P<Body>.*)\$[a-zA-Z]*\$`,
+		functionDefinition,
+	)
+
+	argsData := pgProcedureData["Args"]
+
+	args := []PGProcedureArg{}
+
+	if argsData != "" {
+		rawArgs := strings.Split(argsData, ",")
+		for i := 0; i < len(rawArgs); i++ {
+			var arg PGProcedureArg
+			err := arg.Parse(rawArgs[i])
+			if err != nil {
+				continue
+			}
+			args = append(args, arg)
+		}
+	}
+
+	pgProcedure.Schema = pgProcedureData["Schema"]
+	pgProcedure.Name = pgProcedureData["Name"]
+	pgProcedure.Language = pgProcedureData["Language"]
+	pgProcedure.Body = pgProcedureData["Body"]
+	pgProcedure.Args = args
+	pgProcedure.SecurityDefiner = len(pgProcedureData["Security"]) > 0
+
+	return nil
+}
+
+func (pgProcedureArg *PGProcedureArg) Parse(ProcedureArgDefinition string) error {
+
+	// Check if default exists
+	argDefinitions := findStringSubmatchMap(`(?si)(?P<ArgData>.*)\sDEFAULT\s(?P<ArgDefault>.*)`, ProcedureArgDefinition)
+
+	argData := ProcedureArgDefinition
+	if len(argDefinitions) > 0 {
+		argData = argDefinitions["ArgData"]
+		pgProcedureArg.Default = argDefinitions["ArgDefault"]
+	}
+
+	pgProcedureArgData := findStringSubmatchMap(`(?si)((?P<Mode>IN|OUT|INOUT|VARIADIC)\s)?(?P<Name>[^\s]+)\s(?P<Type>.*)`, argData)
+
+	pgProcedureArg.Name = pgProcedureArgData["Name"]
+	pgProcedureArg.Type = pgProcedureArgData["Type"]
+	pgProcedureArg.Mode = pgProcedureArgData["Mode"]
+	if pgProcedureArg.Mode == "" {
+		pgProcedureArg.Mode = "IN"
+	}
+	return nil
+}
+
+func normalizeProcedureBody(body string) string {
+	newBodyMap := findStringSubmatchMap(`(?si).*\$[a-zA-Z]*\$\s(?P<Body>.*)\s\$[a-zA-Z]*\$.*`, body)
+	if newBody, ok := newBodyMap["Body"]; ok {
+		return newBody
+	}
+	return body
+}
diff --git a/postgresql/model_pg_procedure_test.go b/postgresql/model_pg_procedure_test.go
@@ -0,0 +1,221 @@
+package postgresql
+
+import (
+	"testing"
+
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestProcFromResourceData(t *testing.T) {
+	d := mockProcedureResourceData(t, PGProcedure{
+		Name: "increment",
+		Body: "BEGIN result = i + 1; END;",
+		Args: []PGProcedureArg{
+			{
+				Name: "i",
+				Type: "integer",
+			},
+		},
+	})
+
+	var pgProcedure PGProcedure
+
+	err := pgProcedure.FromResourceData(d)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, pgProcedure, PGProcedure{
+		Schema:          "public",
+		Name:            "increment",
+		Language:        "plpgsql",
+		SecurityDefiner: false,
+		Body:            "BEGIN result = i + 1; END;",
+		Args: []PGProcedureArg{
+			{
+				Name: "i",
+				Type: "integer",
+			},
+		},
+	})
+}
+
+func TestProcFromResourceDataWithArguments(t *testing.T) {
+	d := mockProcedureResourceData(t, PGProcedure{
+		Name: "increment",
+		Body: "BEGIN result = i + 1; END;",
+		Args: []PGProcedureArg{
+			{
+				Name: "i",
+				Type: "integer",
+			},
+		},
+		SecurityDefiner: true,
+	})
+
+	var pgFunction PGProcedure
+
+	err := pgFunction.FromResourceData(d)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, pgFunction, PGProcedure{
+		Schema:          "public",
+		Name:            "increment",
+		Language:        "plpgsql",
+		SecurityDefiner: true,
+		Body:            "BEGIN result = i + 1; END;",
+		Args: []PGProcedureArg{
+			{
+				Name: "i",
+				Type: "integer",
+			},
+		},
+	})
+}
+
+func TestPGProcedureParseWithArguments(t *testing.T) {
+
+	var functionDefinition = `
+CREATE OR REPLACE PROCEDURE public.pg_func_test(showtext boolean, default_null integer DEFAULT NULL::integer, simple_default integer DEFAULT 42, long_default character varying DEFAULT 'foo'::character varying)
+LANGUAGE sql
+SECURITY DEFINER
+AS $function$SELECT 1;$function$
+`
+
+	var pgFunction PGProcedure
+
+	err := pgFunction.Parse(functionDefinition)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, pgFunction, PGProcedure{
+		Name:            "pg_func_test",
+		Schema:          "public",
+		Language:        "sql",
+		SecurityDefiner: true,
+		Body:            "SELECT 1;",
+		Args: []PGProcedureArg{
+			{
+				Mode: "IN",
+				Name: "showtext",
+				Type: "boolean",
+			},
+			{
+				Mode:    "IN",
+				Name:    "default_null",
+				Type:    "integer",
+				Default: "NULL::integer",
+			},
+			{
+				Mode:    "IN",
+				Name:    "simple_default",
+				Type:    "integer",
+				Default: "42",
+			},
+			{
+				Mode:    "IN",
+				Name:    "long_default",
+				Type:    "character varying",
+				Default: "'foo'::character varying",
+			},
+		},
+	})
+}
+
+func TestPGProcedureParseWithoutArguments(t *testing.T) {
+
+	var functionDefinition = `
+CREATE OR REPLACE PROCEDURE public.pg_func_test()
+LANGUAGE plpgsql
+AS $function$
+MultiLine Function
+$function$
+	`
+
+	var pgFunction PGProcedure
+
+	err := pgFunction.Parse(functionDefinition)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, pgFunction, PGProcedure{
+		Name:            "pg_func_test",
+		Schema:          "public",
+		Language:        "plpgsql",
+		SecurityDefiner: false,
+		Body: `
+MultiLine Function
+`,
+		Args: []PGProcedureArg{},
+	})
+}
+
+func TestPGProcedureArgParseWithDefault(t *testing.T) {
+
+	var functionArgDefinition = `default_null integer DEFAULT NULL::integer`
+
+	var pgProcedureArg PGProcedureArg
+
+	err := pgProcedureArg.Parse(functionArgDefinition)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, pgProcedureArg, PGProcedureArg{
+		Mode:    "IN",
+		Name:    "default_null",
+		Type:    "integer",
+		Default: "NULL::integer",
+	})
+}
+
+func TestPGProcedureArgParseWithoutDefault(t *testing.T) {
+
+	var functionArgDefinition = `num integer`
+
+	var pgFunctionArg PGProcedureArg
+
+	err := pgFunctionArg.Parse(functionArgDefinition)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, pgFunctionArg, PGProcedureArg{
+		Mode: "IN",
+		Name: "num",
+		Type: "integer",
+	})
+}
+
+func mockProcedureResourceData(t *testing.T, obj PGProcedure) *schema.ResourceData {
+
+	state := terraform.InstanceState{}
+
+	state.ID = ""
+	// Build the attribute map from ForemanModel
+	attributes := make(map[string]interface{})
+
+	attributes["name"] = obj.Name
+	attributes["language"] = obj.Language
+	attributes["body"] = obj.Body
+	attributes["security_definer"] = obj.SecurityDefiner
+
+	var args []interface{}
+
+	for _, a := range obj.Args {
+		args = append(args, map[string]interface{}{
+			"type":    a.Type,
+			"name":    a.Name,
+			"mode":    a.Mode,
+			"default": a.Default,
+		})
+	}
+
+	attributes["arg"] = args
+
+	return schema.TestResourceDataRaw(t, resourcePostgreSQLProcedure().Schema, attributes)
+}
diff --git a/postgresql/provider.go b/postgresql/provider.go
@@ -213,6 +213,7 @@ func Provider() *schema.Provider {
 			"postgresql_physical_replication_slot": resourcePostgreSQLPhysicalReplicationSlot(),
 			"postgresql_schema":                    resourcePostgreSQLSchema(),
 			"postgresql_role":                      resourcePostgreSQLRole(),
+			"postgresql_procedure":                 resourcePostgreSQLProcedure(),
 			"postgresql_function":                  resourcePostgreSQLFunction(),
 			"postgresql_server":                    resourcePostgreSQLServer(),
 			"postgresql_user_mapping":              resourcePostgreSQLUserMapping(),