Add database parameter on function resource (#216)

* Add database parameter on function resource

* update documentation

* add skipIfNotAcc

Author: LucasBoisserie

postgresql/resource_postgresql_function.go

@@ -16,6 +16,7 @@ const (
 	funcArgAttr         = "arg"
 	funcReturnsAttr     = "returns"
 	funcDropCascadeAttr = "drop_cascade"
+	funcDatabaseAttr    = "database"
 
 	funcArgTypeAttr    = "type"
 	funcArgNameAttr    = "name"
@@ -99,6 +100,13 @@ func resourcePostgreSQLFunction() *schema.Resource {
 				Optional:    true,
 				Default:     false,
 			},
+			funcDatabaseAttr: {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Computed:    true,
+				ForceNew:    true,
+				Description: "The database where the function is located. If not specified, the provider default database is used.",
+			},
 		},
 	}
 }
@@ -129,9 +137,22 @@ func resourcePostgreSQLFunctionExists(db *DBConnection, d *schema.ResourceData)
 	signature := getFunctionSignature(d)
 	functionExists := false
 
+	txn, err := startTransaction(db.client, d.Get(funcDatabaseAttr).(string))
+	if err != nil {
+		return false, err
+	}
+	defer deferredRollback(txn)
+
 	query := fmt.Sprintf("SELECT to_regprocedure('%s') IS NOT NULL", signature)
-	err := db.QueryRow(query).Scan(&functionExists)
-	return functionExists, err
+	if err := txn.QueryRow(query).Scan(&functionExists); err != nil {
+		return false, err
+	}
+
+	if err := txn.Commit(); err != nil {
+		return false, err
+	}
+
+	return functionExists, nil
 }
 
 func resourcePostgreSQLFunctionRead(db *DBConnection, d *schema.ResourceData) error {
@@ -154,7 +175,14 @@ func resourcePostgreSQLFunctionReadImpl(db *DBConnection, d *schema.ResourceData
 		`FROM pg_proc p ` +
 		`LEFT JOIN pg_namespace n ON p.pronamespace = n.oid ` +
 		`WHERE p.oid = to_regprocedure($1)`
-	err := db.QueryRow(query, signature).Scan(&funcSchema, &funcName)
+
+	txn, err := startTransaction(db.client, d.Get(funcDatabaseAttr).(string))
+	if err != nil {
+		return err
+	}
+	defer deferredRollback(txn)
+
+	err = txn.QueryRow(query, signature).Scan(&funcSchema, &funcName)
 	switch {
 	case err == sql.ErrNoRows:
 		log.Printf("[WARN] PostgreSQL function: %s", signature)
@@ -164,6 +192,16 @@ func resourcePostgreSQLFunctionReadImpl(db *DBConnection, d *schema.ResourceData
 		return fmt.Errorf("Error reading function: %w", err)
 	}
 
+	if err := txn.Commit(); err != nil {
+		return err
+	}
+
+	if v, ok := d.GetOk(funcDatabaseAttr); ok {
+		d.Set(funcDatabaseAttr, v)
+	} else {
+		d.Set(funcDatabaseAttr, db.client.databaseName)
+	}
+
 	d.Set(funcNameAttr, funcName)
 	d.Set(funcSchemaAttr, funcSchema)
 	d.SetId(signature)
@@ -187,7 +225,18 @@ func resourcePostgreSQLFunctionDelete(db *DBConnection, d *schema.ResourceData)
 	}
 
 	sql := fmt.Sprintf("DROP FUNCTION IF EXISTS %s %s", signature, dropMode)
-	if _, err := db.Exec(sql); err != nil {
+
+	txn, err := startTransaction(db.client, d.Get(funcDatabaseAttr).(string))
+	if err != nil {
+		return err
+	}
+	defer deferredRollback(txn)
+
+	if _, err := txn.Exec(sql); err != nil {
+		return err
+	}
+
+	if err := txn.Commit(); err != nil {
 		return err
 	}
 
@@ -271,7 +320,18 @@ func createFunction(db *DBConnection, d *schema.ResourceData, replace bool) erro
 	fmt.Fprint(b, "\n", d.Get(funcBodyAttr).(string))
 
 	sql := b.String()
-	if _, err := db.Exec(sql); err != nil {
+
+	txn, err := startTransaction(db.client, d.Get(funcDatabaseAttr).(string))
+	if err != nil {
+		return err
+	}
+	defer deferredRollback(txn)
+
+	if _, err := txn.Exec(sql); err != nil {
+		return err
+	}
+
+	if err := txn.Commit(); err != nil {
 		return err
 	}
 

postgresql/resource_postgresql_function_test.go

@@ -35,7 +35,7 @@ resource "postgresql_function" "basic_function" {
 			{
 				Config: config,
 				Check: resource.ComposeTestCheckFunc(
-					testAccCheckPostgresqlFunctionExists("postgresql_function.basic_function"),
+					testAccCheckPostgresqlFunctionExists("postgresql_function.basic_function", ""),
 					resource.TestCheckResourceAttr(
 						"postgresql_function.basic_function", "name", "basic_function"),
 					resource.TestCheckResourceAttr(
@@ -46,6 +46,53 @@ resource "postgresql_function" "basic_function" {
 	})
 }
 
+func TestAccPostgresqlFunction_SpecificDatabase(t *testing.T) {
+	skipIfNotAcc(t)
+
+	dbSuffix, teardown := setupTestDatabase(t, true, true)
+	defer teardown()
+
+	dbName, _ := getTestDBNames(dbSuffix)
+
+	config := `
+resource "postgresql_function" "basic_function" {
+    name = "basic_function"
+	database = "%s"
+	returns = "integer"
+    body = <<-EOF
+        AS $$
+        BEGIN
+            RETURN 1;
+        END;
+        $$ LANGUAGE plpgsql;
+    EOF
+}
+`
+
+	resource.Test(t, resource.TestCase{
+		PreCheck: func() {
+			testAccPreCheck(t)
+			testCheckCompatibleVersion(t, featureFunction)
+		},
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckPostgresqlFunctionDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: fmt.Sprintf(config, dbName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckPostgresqlFunctionExists("postgresql_function.basic_function", dbName),
+					resource.TestCheckResourceAttr(
+						"postgresql_function.basic_function", "name", "basic_function"),
+					resource.TestCheckResourceAttr(
+						"postgresql_function.basic_function", "database", dbName),
+					resource.TestCheckResourceAttr(
+						"postgresql_function.basic_function", "schema", "public"),
+				),
+			},
+		},
+	})
+}
+
 func TestAccPostgresqlFunction_MultipleArgs(t *testing.T) {
 	config := `
 resource "postgresql_schema" "test" {
@@ -86,7 +133,7 @@ resource "postgresql_function" "increment" {
 			{
 				Config: config,
 				Check: resource.ComposeTestCheckFunc(
-					testAccCheckPostgresqlFunctionExists("postgresql_function.increment"),
+					testAccCheckPostgresqlFunctionExists("postgresql_function.increment", ""),
 					resource.TestCheckResourceAttr(
 						"postgresql_function.increment", "name", "increment"),
 					resource.TestCheckResourceAttr(
@@ -136,7 +183,7 @@ resource "postgresql_function" "func" {
 			{
 				Config: configCreate,
 				Check: resource.ComposeTestCheckFunc(
-					testAccCheckPostgresqlFunctionExists("postgresql_function.func"),
+					testAccCheckPostgresqlFunctionExists("postgresql_function.func", ""),
 					resource.TestCheckResourceAttr(
 						"postgresql_function.func", "name", "func"),
 					resource.TestCheckResourceAttr(
@@ -146,7 +193,7 @@ resource "postgresql_function" "func" {
 			{
 				Config: configUpdate,
 				Check: resource.ComposeTestCheckFunc(
-					testAccCheckPostgresqlFunctionExists("postgresql_function.func"),
+					testAccCheckPostgresqlFunctionExists("postgresql_function.func", ""),
 					resource.TestCheckResourceAttr(
 						"postgresql_function.func", "name", "func"),
 					resource.TestCheckResourceAttr(
@@ -157,7 +204,7 @@ resource "postgresql_function" "func" {
 	})
 }
 
-func testAccCheckPostgresqlFunctionExists(n string) resource.TestCheckFunc {
+func testAccCheckPostgresqlFunctionExists(n string, database string) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		rs, ok := s.RootModule().Resources[n]
 		if !ok {
@@ -171,7 +218,7 @@ func testAccCheckPostgresqlFunctionExists(n string) resource.TestCheckFunc {
 		signature := rs.Primary.ID
 
 		client := testAccProvider.Meta().(*Client)
-		txn, err := startTransaction(client, "")
+		txn, err := startTransaction(client, database)
 		if err != nil {
 			return err
 		}

website/docs/r/postgresql_function.html.markdown

@@ -38,6 +38,9 @@ resource "postgresql_function" "increment" {
 * `schema` - (Optional) The schema where the function is located.
   If not specified, the function is created in the current schema.
 
+* `database` - (Optional) The database where the function is located.
+  If not specified, the function is created in the current database.
+
 * `arg` - (Optional) List of arguments for the function.
   * `type` - (Required) The type of the argument.
   * `name` - (Optional) The name of the argument.