fix: read Name, Database and Schema during Procedure import (#819)

Snowflake-Labs · Jan 18, 2022 · d17656f · d17656f
1 parent 5c94e8e
commit d17656f
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 16 deletions.
diff --git a/pkg/resources/procedure.go b/pkg/resources/procedure.go
@@ -96,7 +96,7 @@ var procedureSchema = map[string]*schema.Schema{
 	"comment": {
 		Type:        schema.TypeString,
 		Optional:    true,
-		Default:     "user-defined function",
+		Default:     "user-defined procedure",
 		Description: "Specifies a comment for the procedure.",
 	},
 }
@@ -287,7 +287,20 @@ func ReadProcedure(d *schema.ResourceData, meta interface{}) error {
 	argSig, _ := proc.ArgumentsSignature()
 
 	for _, v := range foundProcedures {
-		if v.Arguments.String == argSig {
+		showArgs := strings.Split(v.Arguments.String, " RETURN ")
+		if showArgs[0] == argSig {
+			err = d.Set("name", v.Name.String)
+			if err != nil {
+				return err
+			}
+			err = d.Set("database", v.DatabaseName.String)
+			if err != nil {
+				return err
+			}
+			err = d.Set("schema", v.SchemaName.String)
+			if err != nil {
+				return err
+			}
 			err = d.Set("comment", v.Comment.String)
 			if err != nil {
 				return err

diff --git a/pkg/resources/procedure_acceptance_test.go b/pkg/resources/procedure_acceptance_test.go
@@ -37,7 +37,7 @@ func TestAcc_Procedure(t *testing.T) {
 					resource.TestCheckResourceAttr("snowflake_procedure.test_proc", "execute_as", "OWNER"),
 
 					resource.TestCheckResourceAttr("snowflake_procedure.test_proc_simple", "name", procName),
-					resource.TestCheckResourceAttr("snowflake_procedure.test_proc_simple", "comment", "user-defined function"),
+					resource.TestCheckResourceAttr("snowflake_procedure.test_proc_simple", "comment", "user-defined procedure"),
 					resource.TestCheckResourceAttr("snowflake_procedure.test_proc_simple", "statement", expBody1),
 
 					resource.TestCheckResourceAttr("snowflake_procedure.test_proc_complex", "name", procName),
@@ -51,6 +51,11 @@ func TestAcc_Procedure(t *testing.T) {
 					resource.TestCheckResourceAttr("snowflake_procedure.test_proc_complex", "null_input_behavior", "RETURNS NULL ON NULL INPUT"),
 				),
 			},
+			{
+				ResourceName:      "snowflake_procedure.test_proc_complex",
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
 		},
 	})
 }

diff --git a/pkg/resources/procedure_test.go b/pkg/resources/procedure_test.go
@@ -24,6 +24,7 @@ func prepDummyProcedureResource(t *testing.T) *schema.ResourceData {
 		"schema":          "my_schema",
 		"arguments":       []interface{}{argument1, argument2},
 		"return_type":     "varchar",
+		"comment":         "mock comment",
 		"return_behavior": "IMMUTABLE",
 		"statement":       procedureBody, //var message = DATA + DATA;return message
 	}
@@ -42,19 +43,19 @@ func TestProcedureCreate(t *testing.T) {
 	d := prepDummyProcedureResource(t)
 
 	WithMockDb(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
-		mock.ExpectExec(`CREATE OR REPLACE PROCEDURE "my_db"."my_schema"."my_proc"\(data VARCHAR, event_dt DATE\) RETURNS VARCHAR LANGUAGE javascript CALLED ON NULL INPUT IMMUTABLE COMMENT = 'user-defined function' EXECUTE AS OWNER AS \$\$hi\$\$`).WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectExec(`CREATE OR REPLACE PROCEDURE "my_db"."my_schema"."my_proc"\(data VARCHAR, event_dt DATE\) RETURNS VARCHAR LANGUAGE javascript CALLED ON NULL INPUT IMMUTABLE COMMENT = 'mock comment' EXECUTE AS OWNER AS \$\$hi\$\$`).WillReturnResult(sqlmock.NewResult(1, 1))
 		expectProcedureRead(mock)
 		err := resources.CreateProcedure(d, db)
 		r.NoError(err)
-		r.Equal("my_proc", d.Get("name").(string))
+		r.Equal("MY_PROC", d.Get("name").(string))
 		r.Equal("VARCHAR", d.Get("return_type").(string))
-		r.Equal("user-defined function", d.Get("comment").(string))
+		r.Equal("mock comment", d.Get("comment").(string))
 	})
 }
 
 func expectProcedureRead(mock sqlmock.Sqlmock) {
 	rows := sqlmock.NewRows([]string{"created_on", "name", "schema_name", "is_builtin", "is_aggregate", "is_ansi", "min_num_arguments", "max_num_arguments", "arguments", "description", "catalog_name", "is_table_function", "valid_for_clustering", "is_secure"}).
-		AddRow("now", "my_proc", "my_schema", "N", "N", "N", "1", "1", "MY_TEST_FUNCTION(VARCHAR) RETURN VARCHAR", "mock comment", "my_db", "N", "N", "N")
+		AddRow("now", "MY_PROC", "MY_SCHEMA", "N", "N", "N", "1", "1", "MY_PROC(VARCHAR, DATE) RETURN VARCHAR", "mock comment", "MY_DB", "N", "N", "N")
 	mock.ExpectQuery(`SHOW PROCEDURES LIKE 'my_proc' IN SCHEMA "my_db"."my_schema"`).WillReturnRows(rows)
 
 	describeRows := sqlmock.NewRows([]string{"property", "value"}).
@@ -72,16 +73,19 @@ func expectProcedureRead(mock sqlmock.Sqlmock) {
 func TestProcedureRead(t *testing.T) {
 	r := require.New(t)
 
-	d := prepDummyProcedureResource(t)
+	d := procedure(t, "my_db|my_schema|my_proc|VARCHAR-DATE", map[string]interface{}{})
 
 	WithMockDb(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		expectProcedureRead(mock)
 
 		err := resources.ReadProcedure(d, db)
 		r.NoError(err)
-		r.Equal("my_proc", d.Get("name").(string))
-		r.Equal("user-defined function", d.Get("comment").(string))
+		r.Equal("MY_PROC", d.Get("name").(string))
+		r.Equal("MY_DB", d.Get("database").(string))
+		r.Equal("MY_SCHEMA", d.Get("schema").(string))
+		r.Equal("mock comment", d.Get("comment").(string))
 		r.Equal("VARCHAR", d.Get("return_type").(string))
+		r.Equal("IMMUTABLE", d.Get("return_behavior").(string))
 		r.Equal(procedureBody, d.Get("statement").(string))
 
 		args := d.Get("arguments").([]interface{})
@@ -100,7 +104,7 @@ func TestProcedureRead(t *testing.T) {
 func TestProcedureDelete(t *testing.T) {
 	r := require.New(t)
 
-	d := prepDummyProcedureResource(t)
+	d := procedure(t, "my_db|my_schema|my_proc|VARCHAR-DATE", map[string]interface{}{})
 
 	WithMockDb(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		mock.ExpectExec(`DROP PROCEDURE "my_db"."my_schema"."my_proc"\(VARCHAR, DATE\)`).WillReturnResult(sqlmock.NewResult(1, 1))

diff --git a/pkg/snowflake/procedure.go b/pkg/snowflake/procedure.go
@@ -42,9 +42,9 @@ func (pb *ProcedureBuilder) QualifiedNameWithoutArguments() (string, error) {
 	return fmt.Sprintf(`"%v"."%v"."%v"`, pb.db, pb.schema, pb.name), nil
 }
 
-// Returns the arguments signature of the procedure in a form <procedure>(<type>, <type>, ..) RETURN <type>
+// Returns the arguments signature of the procedure in a form <PROCEDURE>(<TYPE>, <TYPE>, ..)
 func (pb *ProcedureBuilder) ArgumentsSignature() (string, error) {
-	return fmt.Sprintf(`%v(%v) RETURN %v`, pb.name, strings.Join(pb.argumentTypes, ", "), pb.returnType), nil
+	return fmt.Sprintf(`%v(%v)`, strings.ToUpper(pb.name), strings.ToUpper(strings.Join(pb.argumentTypes, ", "))), nil
 }
 
 // WithArgs sets the args and argumentTypes on the ProcedureBuilder
@@ -230,7 +230,7 @@ type procedure struct {
 	Name         sql.NullString `db:"name"`
 	SchemaName   sql.NullString `db:"schema_name"`
 	Text         sql.NullString `db:"text"`
-	DatabaseName sql.NullString `db:"database_name"`
+	DatabaseName sql.NullString `db:"catalog_name"`
 	Arguments    sql.NullString `db:"arguments"`
 }
 

diff --git a/pkg/snowflake/procedure_test.go b/pkg/snowflake/procedure_test.go
@@ -124,8 +124,8 @@ func TestProcedureArgumentsSignature(t *testing.T) {
 	r := require.New(t)
 	s := getProcedure(false)
 	sign, _ := s.ArgumentsSignature()
-	r.Equal("test_proc() RETURN VARCHAR", sign)
+	r.Equal("TEST_PROC()", sign)
 	s = getProcedure(true)
 	sign, _ = s.ArgumentsSignature()
-	r.Equal("test_proc(VARCHAR, DATE) RETURN VARCHAR", sign)
+	r.Equal("TEST_PROC(VARCHAR, DATE)", sign)
 }