Repository: chanzuckerberg/terraform-provider-snowflake
Pull request #798 (2021-12-26 08:31:01)

feat: Support create function with Java language
I implemented `snowflake_function` for Java UDF.

## References
* https://docs.snowflake.com/en/sql-reference/sql/create-function.html

alldoami: /ok-to-test sha=4bd7cd4
github-actions[bot]: Integration tests failure for [4bd7cd4](https://github.com/chanzuckerberg/terraform-provider-snowflake/actions/runs/1650864569)
alldoami: /ok-to-test sha=6474a7c
github-actions[bot]: Integration tests success for [6474a7c](https://github.com/chanzuckerberg/terraform-provider-snowflake/actions/runs/1650876460)

Patch:
diff --git a/docs/resources/function.md b/docs/resources/function.md
index bc504884..16cfef9f 100644
--- a/docs/resources/function.md
+++ b/docs/resources/function.md
@@ -27,10 +27,13 @@ description: |-
 
 - **arguments** (Block List) List of the arguments for the function (see [below for nested schema](#nestedblock--arguments))
 - **comment** (String) Specifies a comment for the function.
+- **handler** (String) the handler method for Java function.
 - **id** (String) The ID of this resource.
+- **imports** (List of String) jar files to import for Java function.
 - **language** (String) The language of the statement
 - **null_input_behavior** (String) Specifies the behavior of the function when called with null inputs.
 - **return_behavior** (String) Specifies the behavior of the function when returning results
+- **target_path** (String) the target path for compiled jar file for Java function.
 
 <a id="nestedblock--arguments"></a>
 ### Nested Schema for `arguments`
diff --git a/pkg/resources/function.go b/pkg/resources/function.go
index 4cad34fa..597bc21a 100644
--- a/pkg/resources/function.go
+++ b/pkg/resources/function.go
@@ -102,6 +102,27 @@ var functionSchema = map[string]*schema.Schema{
 		Default:     "user-defined function",
 		Description: "Specifies a comment for the function.",
 	},
+	"imports": {
+		Type: schema.TypeList,
+		Elem: &schema.Schema{
+			Type: schema.TypeString,
+		},
+		Optional:    true,
+		ForceNew:    true,
+		Description: "jar files to import for Java function.",
+	},
+	"handler": {
+		Type:        schema.TypeString,
+		Optional:    true,
+		ForceNew:    true,
+		Description: "the handler method for Java function.",
+	},
+	"target_path": {
+		Type:        schema.TypeString,
+		Optional:    true,
+		ForceNew:    true,
+		Description: "the target path for compiled jar file for Java function.",
+	},
 }
 
 // Function returns a pointer to the resource representing a stored function
@@ -162,6 +183,25 @@ func CreateFunction(d *schema.ResourceData, meta interface{}) error {
 		builder.WithComment(v.(string))
 	}
 
+	// Set optionals, imports for Java
+	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
+	if v, ok := d.GetOk("handler"); ok {
+		builder.WithHandler(v.(string))
+	}
+
+	// target path for Java
+	if v, ok := d.GetOk("target_path"); ok {
+		builder.WithTargetPath(v.(string))
+	}
+
 	q, err := builder.Create()
 	if err != nil {
 		return err
@@ -263,6 +303,24 @@ func ReadFunction(d *schema.ResourceData, meta interface{}) error {
 					return err
 				}
 			}
+		case "imports":
+			importsString := 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
+			}
+		case "target_path":
+			if err = d.Set("target_path", desc.Value.String); err != nil {
+				return err
+			}
+		case "runtime_version":
+			// runtime version for Java function. currently not used.
 		default:
 			log.Printf("[WARN] unexpected function property %v returned from Snowflake", desc.Property.String)
 		}
diff --git a/pkg/snowflake/function.go b/pkg/snowflake/function.go
index b46018b8..cebe2111 100644
--- a/pkg/snowflake/function.go
+++ b/pkg/snowflake/function.go
@@ -22,6 +22,9 @@ type FunctionBuilder struct {
 	nullInputBehavior string // "CALLED ON NULL INPUT" or "RETURNS NULL ON NULL INPUT"
 	returnType        string
 	language          string
+	imports           []string // for Java imports
+	handler           string   // for Java handler
+	targetPath        string   // for Java target path for compiled jar file
 	comment           string
 	statement         string
 }
@@ -83,6 +86,24 @@ func (pb *FunctionBuilder) WithLanguage(s string) *FunctionBuilder {
 	return pb
 }
 
+// WithImports adds jar files to import for Java function
+func (pb *FunctionBuilder) WithImports(s []string) *FunctionBuilder {
+	pb.imports = s
+	return pb
+}
+
+// WithHandler sets the handler method for Java function
+func (pb *FunctionBuilder) WithHandler(s string) *FunctionBuilder {
+	pb.handler = s
+	return pb
+}
+
+// WithTargetPath sets the target path for compiled jar file for Java function
+func (pb *FunctionBuilder) WithTargetPath(s string) *FunctionBuilder {
+	pb.targetPath = s
+	return pb
+}
+
 // WithComment adds a comment to the FunctionBuilder
 func (pb *FunctionBuilder) WithComment(c string) *FunctionBuilder {
 	pb.comment = c
@@ -153,6 +174,21 @@ func (pb *FunctionBuilder) Create() (string, error) {
 	if pb.comment != "" {
 		q.WriteString(fmt.Sprintf(" COMMENT = '%v'", EscapeString(pb.comment)))
 	}
+	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.targetPath != "" {
+		q.WriteString(fmt.Sprintf(" TARGET_PATH = '%v'", pb.targetPath))
+	}
 	q.WriteString(fmt.Sprintf(" AS $$%v$$", pb.statement))
 	return q.String(), nil
 }


Test patch:
diff --git a/pkg/resources/function_acceptance_test.go b/pkg/resources/function_acceptance_test.go
index 9ac23a0a..77d17dc2 100644
--- a/pkg/resources/function_acceptance_test.go
+++ b/pkg/resources/function_acceptance_test.go
@@ -21,6 +21,7 @@ func TestAcc_Function(t *testing.T) {
 	expBody1 := "3.141592654::FLOAT"
 	expBody2 := "var X=3\nreturn X"
 	expBody3 := "select 1, 2\nunion all\nselect 3, 4\n"
+	expBody4 := `class CoolFunc {public static String test(int n) {return "hello!";}}`
 
 	resource.Test(t, resource.TestCase{
 		Providers: providers(),
@@ -45,6 +46,13 @@ func TestAcc_Function(t *testing.T) {
 					resource.TestCheckResourceAttr("snowflake_function.test_funct_complex", "arguments.#", "2"),
 					resource.TestCheckResourceAttr("snowflake_function.test_funct_complex", "arguments.1.name", "ARG2"),
 					resource.TestCheckResourceAttr("snowflake_function.test_funct_complex", "arguments.1.type", "DATE"),
+
+					resource.TestCheckResourceAttr("snowflake_function.test_funct_java", "name", functName),
+					resource.TestCheckResourceAttr("snowflake_function.test_funct_java", "comment", "Terraform acceptance test for java"),
+					resource.TestCheckResourceAttr("snowflake_function.test_funct_java", "statement", expBody4),
+					resource.TestCheckResourceAttr("snowflake_function.test_funct_java", "arguments.#", "1"),
+					resource.TestCheckResourceAttr("snowflake_function.test_funct_java", "arguments.0.name", "ARG1"),
+					resource.TestCheckResourceAttr("snowflake_function.test_funct_java", "arguments.0.type", "NUMBER"),
 				),
 			},
 		},
@@ -86,6 +94,21 @@ func functionConfig(db, schema, name string) string {
 		statement = "var X=3\nreturn X"
 	}
 
+	resource "snowflake_function" "test_funct_java" {
+		name = "%s"
+		database = snowflake_database.test_database.name
+		schema   = snowflake_schema.test_schema.name
+		arguments {
+			name = "arg1"
+			type = "number"
+		}
+		comment = "Terraform acceptance test for java"
+		return_type = "varchar"
+		language = "java"
+		handler = "CoolFunc.test"
+		statement = "class CoolFunc {public static String test(int n) {return \"hello!\";}}"
+	}
+
 	resource "snowflake_function" "test_funct_complex" {
 		name = "%s"
 		database = snowflake_database.test_database.name
@@ -106,5 +129,5 @@ union all
 select 3, 4
 EOT
 	}
-	`, db, schema, name, name, name)
+	`, db, schema, name, name, name, name)
 }
diff --git a/pkg/snowflake/function_test.go b/pkg/snowflake/function_test.go
index f9263d7b..14216506 100644
--- a/pkg/snowflake/function_test.go
+++ b/pkg/snowflake/function_test.go
@@ -6,7 +6,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func getFunction(withArgs bool) *FunctionBuilder {
+func getJavaScriptFuction(withArgs bool) *FunctionBuilder {
 	s := Function("test_db", "test_schema", "test_func", []string{})
 	s.WithStatement(`var message = "Hi"` + "\n" + `return message`)
 	s.WithReturnType("varchar")
@@ -18,9 +18,27 @@ func getFunction(withArgs bool) *FunctionBuilder {
 	return s
 }
 
+const javafunc = `class CoolFunc {` + "\n" +
+	`  public static String test(String u, int c) {` + "\n" +
+	`    return u;` + "\n" +
+	`  }` + "\n" +
+	`}`
+
+func getJavaFuction(withArgs bool) *FunctionBuilder {
+	s := Function("test_db", "test_schema", "test_func", []string{})
+	s.WithReturnType("varchar")
+	s.WithStatement(javafunc)
+	if withArgs {
+		s.WithArgs([]map[string]string{
+			{"name": "user", "type": "varchar"},
+			{"name": "count", "type": "number"}})
+	}
+	return s
+}
+
 func TestFunctionQualifiedName(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(true)
+	s := getJavaScriptFuction(true)
 	qn, _ := s.QualifiedName()
 	r.Equal(`"test_db"."test_schema"."test_func"(VARCHAR, DATE)`, qn)
 	qna, _ := s.QualifiedNameWithoutArguments()
@@ -29,7 +47,7 @@ func TestFunctionQualifiedName(t *testing.T) {
 
 func TestFunctionCreate(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(true)
+	s := getJavaScriptFuction(true)
 
 	r.Equal([]string{"VARCHAR", "DATE"}, s.ArgTypes())
 	createStmnt, _ := s.Create()
@@ -39,9 +57,9 @@ func TestFunctionCreate(t *testing.T) {
 	r.Equal(expected, createStmnt)
 }
 
-func TestFunctionCreateWithOptionalParams(t *testing.T) {
+func TestFunctionCreateWithJavaScriptFunction(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(true)
+	s := getJavaScriptFuction(true)
 	s.WithNullInputBehavior("RETURNS NULL ON NULL INPUT")
 	s.WithReturnBehavior("IMMUTABLE")
 	s.WithComment("this is cool func!")
@@ -54,30 +72,82 @@ func TestFunctionCreateWithOptionalParams(t *testing.T) {
 	r.Equal(expected, createStmnt)
 }
 
+func TestFunctionCreateWithJavaFunction(t *testing.T) {
+	r := require.New(t)
+	s := getJavaFuction(true)
+	s.WithNullInputBehavior("RETURNS NULL ON NULL INPUT")
+	s.WithReturnBehavior("IMMUTABLE")
+	s.WithComment("this is cool func!")
+	s.WithLanguage("JAVA")
+	s.WithHandler("CoolFunc.test")
+	createStmnt, _ := s.Create()
+	expected := `CREATE OR REPLACE FUNCTION "test_db"."test_schema"."test_func"` +
+		`(user VARCHAR, count NUMBER) RETURNS VARCHAR` +
+		` LANGUAGE JAVA RETURNS NULL ON NULL INPUT IMMUTABLE COMMENT = 'this is cool func!'` +
+		` HANDLER = 'CoolFunc.test' AS $$` + javafunc + `$$`
+	r.Equal(expected, createStmnt)
+}
+
+func TestFunctionCreateWithJavaFunctionWithImports(t *testing.T) {
+	r := require.New(t)
+	s := getJavaFuction(true)
+	s.WithNullInputBehavior("RETURNS NULL ON NULL INPUT")
+	s.WithReturnBehavior("IMMUTABLE")
+	s.WithComment("this is cool func!")
+	s.WithLanguage("JAVA")
+	s.WithImports([]string{"@~/stage/myudf1.jar", "@~/stage/myudf2.jar"})
+	s.WithHandler("CoolFunc.test")
+	createStmnt, _ := s.Create()
+	expected := `CREATE OR REPLACE FUNCTION "test_db"."test_schema"."test_func"` +
+		`(user VARCHAR, count NUMBER) RETURNS VARCHAR` +
+		` LANGUAGE JAVA RETURNS NULL ON NULL INPUT IMMUTABLE COMMENT = 'this is cool func!'` +
+		` IMPORTS = ('@~/stage/myudf1.jar', '@~/stage/myudf2.jar') HANDLER = 'CoolFunc.test'` +
+		` AS $$` + javafunc + `$$`
+	r.Equal(expected, createStmnt)
+}
+
+func TestFunctionCreateWithJavaFunctionWithTargetPath(t *testing.T) {
+	r := require.New(t)
+	s := getJavaFuction(true)
+	s.WithNullInputBehavior("RETURNS NULL ON NULL INPUT")
+	s.WithReturnBehavior("IMMUTABLE")
+	s.WithComment("this is cool func!")
+	s.WithLanguage("JAVA")
+	s.WithTargetPath("@~/stage/myudf1.jar")
+	s.WithHandler("CoolFunc.test")
+	createStmnt, _ := s.Create()
+	expected := `CREATE OR REPLACE FUNCTION "test_db"."test_schema"."test_func"` +
+		`(user VARCHAR, count NUMBER) RETURNS VARCHAR` +
+		` LANGUAGE JAVA RETURNS NULL ON NULL INPUT IMMUTABLE COMMENT = 'this is cool func!'` +
+		` HANDLER = 'CoolFunc.test' TARGET_PATH = '@~/stage/myudf1.jar'` +
+		` AS $$` + javafunc + `$$`
+	r.Equal(expected, createStmnt)
+}
+
 func TestFunctionDrop(t *testing.T) {
 	r := require.New(t)
 
 	// Without arg
-	s := getFunction(false)
+	s := getJavaScriptFuction(false)
 	stmnt, _ := s.Drop()
 	r.Equal(stmnt, `DROP FUNCTION "test_db"."test_schema"."test_func"()`)
 
 	// With arg
-	ss := getFunction(true)
+	ss := getJavaScriptFuction(true)
 	stmnt, _ = ss.Drop()
 	r.Equal(`DROP FUNCTION "test_db"."test_schema"."test_func"(VARCHAR, DATE)`, stmnt)
 }
 
 func TestFunctionShow(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(false)
+	s := getJavaScriptFuction(false)
 	stmnt := s.Show()
 	r.Equal(stmnt, `SHOW USER FUNCTIONS LIKE 'test_func' IN SCHEMA "test_db"."test_schema"`)
 }
 
 func TestFunctionRename(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(false)
+	s := getJavaScriptFuction(false)
 
 	stmnt, _ := s.Rename("new_func")
 	expected := `ALTER FUNCTION "test_db"."test_schema"."test_func"() RENAME TO "test_db"."test_schema"."new_func"`
@@ -86,7 +156,7 @@ func TestFunctionRename(t *testing.T) {
 
 func TestFunctionChangeComment(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(true)
+	s := getJavaScriptFuction(true)
 
 	stmnt, _ := s.ChangeComment("not used")
 	expected := `ALTER FUNCTION "test_db"."test_schema"."test_func"(VARCHAR, DATE) SET COMMENT = 'not used'`
@@ -95,7 +165,7 @@ func TestFunctionChangeComment(t *testing.T) {
 
 func TestFunctionRemoveComment(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(false)
+	s := getJavaScriptFuction(false)
 
 	stmnt, _ := s.RemoveComment()
 	expected := `ALTER FUNCTION "test_db"."test_schema"."test_func"() UNSET COMMENT`
@@ -104,7 +174,7 @@ func TestFunctionRemoveComment(t *testing.T) {
 
 func TestFunctionDescribe(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(false)
+	s := getJavaScriptFuction(false)
 
 	stmnt, _ := s.Describe()
 	expected := `DESCRIBE FUNCTION "test_db"."test_schema"."test_func"()`
@@ -113,10 +183,10 @@ func TestFunctionDescribe(t *testing.T) {
 
 func TestFunctionArgumentsSignature(t *testing.T) {
 	r := require.New(t)
-	s := getFunction(false)
+	s := getJavaScriptFuction(false)
 	sign, _ := s.ArgumentsSignature()
 	r.Equal("test_func() RETURN VARCHAR", sign)
-	s = getFunction(true)
+	s = getJavaScriptFuction(true)
 	sign, _ = s.ArgumentsSignature()
 	r.Equal("test_func(VARCHAR, DATE) RETURN VARCHAR", sign)
 }