feat: Add support for packages, imports, handler and runtimeVersion to snowflake_procedure resource

pkg/resources/procedure.go

@@ -14,7 +14,7 @@ import (
 	"golang.org/x/exp/slices"
 )
 
-var procedureLanguages = []string{"javascript", "java", "scala", "SQL"}
+var procedureLanguages = []string{"javascript", "java", "scala", "SQL", "python"}
 
 var procedureSchema = map[string]*schema.Schema{
 	"name": {
@@ -89,7 +89,7 @@ var procedureSchema = map[string]*schema.Schema{
 	"statement": {
 		Type:             schema.TypeString,
 		Required:         true,
-		Description:      "Specifies the javascript code used to create the procedure.",
+		Description:      "Specifies the code used to create the procedure.",
 		ForceNew:         true,
 		DiffSuppressFunc: DiffSuppressStatement,
 	},
@@ -132,6 +132,36 @@ var procedureSchema = map[string]*schema.Schema{
 		Default:     "user-defined procedure",
 		Description: "Specifies a comment for the procedure.",
 	},
+	"runtime_version": {
+		Type:        schema.TypeString,
+		Optional:    true,
+		ForceNew:    true,
+		Description: "Required for Python procedures. Specifies Python runtime version.",
+	},
+	"packages": {
+		Type: schema.TypeList,
+		Elem: &schema.Schema{
+			Type: schema.TypeString,
+		},
+		Optional:    true,
+		ForceNew:    true,
+		Description: "List of package imports to use for Java / Python procedures. For Java, package imports should be of the form: package_name:version_number, where package_name is snowflake_domain:package. For Python use it should be: ('numpy','pandas','xgboost==1.5.0').",
+	},
+	"imports": {
+		Type: schema.TypeList,
+		Elem: &schema.Schema{
+			Type: schema.TypeString,
+		},
+		Optional:    true,
+		ForceNew:    true,
+		Description: "Imports for Java / Python procedures. For Java this a list of jar files, for Python this is a list of Python files.",
+	},
+	"handler": {
+		Type:        schema.TypeString,
+		Optional:    true,
+		ForceNew:    true,
+		Description: "The handler method for Java / Python procedures.",
+	},
 }
 
 func DiffTypes(k, old, new string, d *schema.ResourceData) bool {
@@ -197,10 +227,38 @@ func CreateProcedure(d *schema.ResourceData, meta interface{}) error {
 		builder.WithLanguage(strings.ToUpper(v.(string)))
 	}
 
+	// Set optionals, runtime version for Python
+	if v, ok := d.GetOk("runtime_version"); ok {
+		builder.WithRuntimeVersion(v.(string))
+	}
+
 	if v, ok := d.GetOk("comment"); ok {
 		builder.WithComment(v.(string))
 	}
 
+	// Set optionals, packages for Java / Python
+	if _, ok := d.GetOk("packages"); ok {
+		packages := []string{}
+		for _, pack := range d.Get("packages").([]interface{}) {
+			packages = append(packages, pack.(string))
+		}
+		builder.WithPackages(packages)
+	}
+
+	// Set optionals, imports for Java / Python
+	if _, ok := d.GetOk("imports"); ok {
+		imports := []string{}
+		for _, imp := range d.Get("imports").([]interface{}) {
+			imports = append(imports, imp.(string))
+		}
+		builder.WithImports(imports)
+	}
+
+	// handler for Java / Python
+	if v, ok := d.GetOk("handler"); ok {
+		builder.WithHandler(v.(string))
+	}
+
 	q, err := builder.Create()
 	if err != nil {
 		return err
@@ -297,6 +355,30 @@ func ReadProcedure(d *schema.ResourceData, meta interface{}) error {
 			if err := d.Set("language", desc.Value.String); err != nil {
 				return err
 			}
+		case "runtime_version":
+			if err := d.Set("runtime_version", desc.Value.String); err != nil {
+				return err
+			}
+		case "packages":
+			packagesString := strings.ReplaceAll(strings.ReplaceAll(strings.ReplaceAll(desc.Value.String, "[", ""), "]", ""), "'", "")
+			if packagesString != "" { // Do nothing for Java / Python functions without packages
+				packages := strings.Split(packagesString, ",")
+				if err := d.Set("packages", packages); err != nil {
+					return err
+				}
+			}
+		case "imports":
+			importsString := strings.ReplaceAll(strings.ReplaceAll(strings.ReplaceAll(desc.Value.String, "[", ""), "]", ""), "'", "")
+			if importsString != "" { // Do nothing for Java functions without imports
+				imports := strings.Split(importsString, ",")
+				if err := d.Set("imports", imports); err != nil {
+					return err
+				}
+			}
+		case "handler":
+			if err := d.Set("handler", desc.Value.String); err != nil {
+				return err
+			}
 
 		default:
 			log.Printf("[WARN] unexpected procedure property %v returned from Snowflake", desc.Property.String)

pkg/snowflake/procedure.go

@@ -19,11 +19,15 @@ type ProcedureBuilder struct {
 	args              []map[string]string
 	returnBehavior    string // VOLATILE, IMMUTABLE
 	nullInputBehavior string // "CALLED ON NULL INPUT" or "RETURNS NULL ON NULL INPUT"
-	language          string // SQL, JAVASCRIPT, JAVA, SCALA
 	returnType        string
+	language          string // SQL, JAVASCRIPT, JAVA, SCALA
+	packages          []string
+	imports           []string // for Java / Python imports
+	handler           string   // for Java / Python handler
 	executeAs         string
 	comment           string
 	statement         string
+	runtimeVersion    string // for Python runtime version
 }
 
 // QualifiedName prepends the db and schema and appends argument types.
@@ -89,6 +93,30 @@ func (pb *ProcedureBuilder) WithLanguage(s string) *ProcedureBuilder {
 	return pb
 }
 
+// WithRuntimeVersion.
+func (pb *ProcedureBuilder) WithRuntimeVersion(r string) *ProcedureBuilder {
+	pb.runtimeVersion = r
+	return pb
+}
+
+// WithPackages.
+func (pb *ProcedureBuilder) WithPackages(s []string) *ProcedureBuilder {
+	pb.packages = s
+	return pb
+}
+
+// WithImports adds jar files to import for Java function or Python file for Python function.
+func (pb *ProcedureBuilder) WithImports(s []string) *ProcedureBuilder {
+	pb.imports = s
+	return pb
+}
+
+// WithHandler sets the handler method for Java / Python function.
+func (pb *ProcedureBuilder) WithHandler(s string) *ProcedureBuilder {
+	pb.handler = s
+	return pb
+}
+
 // WithComment adds a comment to the ProcedureBuilder.
 func (pb *ProcedureBuilder) WithComment(c string) *ProcedureBuilder {
 	pb.comment = c
@@ -156,6 +184,30 @@ func (pb *ProcedureBuilder) Create() (string, error) {
 	if pb.returnBehavior != "" {
 		q.WriteString(fmt.Sprintf(` %v`, EscapeString(pb.returnBehavior)))
 	}
+	if pb.runtimeVersion != "" {
+		q.WriteString(fmt.Sprintf(" RUNTIME_VERSION = '%v'", EscapeString(pb.runtimeVersion)))
+	}
+	if len(pb.packages) > 0 {
+		q.WriteString(` PACKAGES = (`)
+		packages := []string{}
+		for _, pack := range pb.packages {
+			packages = append(packages, fmt.Sprintf(`'%v'`, pack))
+		}
+		q.WriteString(strings.Join(packages, ", "))
+		q.WriteString(`)`)
+	}
+	if len(pb.imports) > 0 {
+		q.WriteString(` IMPORTS = (`)
+		imports := []string{}
+		for _, imp := range pb.imports {
+			imports = append(imports, fmt.Sprintf(`'%v'`, imp))
+		}
+		q.WriteString(strings.Join(imports, ", "))
+		q.WriteString(`)`)
+	}
+	if pb.handler != "" {
+		q.WriteString(fmt.Sprintf(" HANDLER = '%v'", pb.handler))
+	}
 	if pb.comment != "" {
 		q.WriteString(fmt.Sprintf(" COMMENT = '%v'", EscapeString(pb.comment)))
 	}

pkg/snowflake/procedure_test.go

@@ -46,13 +46,18 @@ func TestProcedureCreateWithOptionalParams(t *testing.T) {
 	s := getProcedure(true)
 	s.WithNullInputBehavior("RETURNS NULL ON NULL INPUT")
 	s.WithReturnBehavior("IMMUTABLE")
-	s.WithLanguage("JAVASCRIPT")
+	s.WithLanguage("PYTHON")
+	s.WithRuntimeVersion("3.8")
+	s.WithPackages([]string{"snowflake-snowpark-python", "pandas"})
+	s.WithImports([]string{"@\"test_db\".\"test_schema\".\"test_stage\"/handler.py"})
+	s.WithHandler("handler.test")
 	s.WithComment("this is cool proc!")
 	createStmnt, _ := s.Create()
 	expected := `CREATE OR REPLACE PROCEDURE "test_db"."test_schema"."test_proc"` +
-		`(user VARCHAR, eventdt DATE) RETURNS VARCHAR LANGUAGE JAVASCRIPT RETURNS NULL ON NULL INPUT` +
-		` IMMUTABLE COMMENT = 'this is cool proc!' EXECUTE AS CALLER AS $$` +
-		`var message = "Hi"` + "\nreturn message$$"
+		`(user VARCHAR, eventdt DATE) RETURNS VARCHAR LANGUAGE PYTHON RETURNS NULL ON NULL INPUT ` +
+		`IMMUTABLE RUNTIME_VERSION = '3.8' PACKAGES = ('snowflake-snowpark-python', 'pandas') ` +
+		`IMPORTS = ('@"test_db"."test_schema"."test_stage"/handler.py') HANDLER = 'handler.test' ` +
+		`COMMENT = 'this is cool proc!' EXECUTE AS CALLER AS $$var message = "Hi"` + "\nreturn message$$"
 	r.Equal(expected, createStmnt)
 }