From de6dcbd8cd7a0e27a155b498c987fd36b38ae632 Mon Sep 17 00:00:00 2001
From: Sean Chittenden <sean@chittenden.org>
Date: Wed, 14 Dec 2016 01:04:35 -0800
Subject: [PATCH] Add the `owner` attribute to the `postgresql_schema`
 resource.

---
 .../postgresql/resource_postgresql_schema.go  | 43 +++++++++++++++++--
 .../resource_postgresql_schema_test.go        | 10 +++++
 .../r/postgresql_schema.html.markdown         | 14 +++---
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/builtin/providers/postgresql/resource_postgresql_schema.go b/builtin/providers/postgresql/resource_postgresql_schema.go
index 1c16004f4..a712a0838 100644
--- a/builtin/providers/postgresql/resource_postgresql_schema.go
+++ b/builtin/providers/postgresql/resource_postgresql_schema.go
@@ -13,7 +13,8 @@ import (
 )
 
 const (
-	schemaNameAttr = "name"
+	schemaNameAttr  = "name"
+	schemaOwnerAttr = "owner"
 )
 
 func resourcePostgreSQLSchema() *schema.Resource {
@@ -32,6 +33,12 @@ func resourcePostgreSQLSchema() *schema.Resource {
 				Required:    true,
 				Description: "The name of the schema",
 			},
+			schemaOwnerAttr: {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Computed:    true,
+				Description: "The ROLE name who owns the schema",
+			},
 		},
 	}
 }
@@ -48,6 +55,11 @@ func resourcePostgreSQLSchemaCreate(d *schema.ResourceData, meta interface{}) er
 	b := bytes.NewBufferString("CREATE SCHEMA ")
 	fmt.Fprintf(b, pq.QuoteIdentifier(schemaName))
 
+	switch v, ok := d.GetOk(schemaOwnerAttr); {
+	case ok:
+		fmt.Fprint(b, " AUTHORIZATION ", pq.QuoteIdentifier(v.(string)))
+	}
+
 	query := b.String()
 	_, err = conn.Query(query)
 	if err != nil {
@@ -88,8 +100,8 @@ func resourcePostgreSQLSchemaRead(d *schema.ResourceData, meta interface{}) erro
 	defer conn.Close()
 
 	schemaId := d.Id()
-	var schemaName string
-	err = conn.QueryRow("SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname=$1", schemaId).Scan(&schemaName)
+	var schemaName, schemaOwner string
+	err = conn.QueryRow("SELECT nspname, pg_catalog.pg_get_userbyid(nspowner) FROM pg_catalog.pg_namespace WHERE nspname=$1", schemaId).Scan(&schemaName, &schemaOwner)
 	switch {
 	case err == sql.ErrNoRows:
 		log.Printf("[WARN] PostgreSQL schema (%s) not found", schemaId)
@@ -99,6 +111,7 @@ func resourcePostgreSQLSchemaRead(d *schema.ResourceData, meta interface{}) erro
 		return errwrap.Wrapf("Error reading schema: {{err}}", err)
 	default:
 		d.Set(schemaNameAttr, schemaName)
+		d.Set(schemaOwnerAttr, schemaOwner)
 		d.SetId(schemaName)
 		return nil
 	}
@@ -116,6 +129,10 @@ func resourcePostgreSQLSchemaUpdate(d *schema.ResourceData, meta interface{}) er
 		return err
 	}
 
+	if err := setSchemaOwner(conn, d); err != nil {
+		return err
+	}
+
 	return resourcePostgreSQLSchemaRead(d, meta)
 }
 
@@ -139,3 +156,23 @@ func setSchemaName(conn *sql.DB, d *schema.ResourceData) error {
 
 	return nil
 }
+
+func setSchemaOwner(conn *sql.DB, d *schema.ResourceData) error {
+	if !d.HasChange(schemaOwnerAttr) {
+		return nil
+	}
+
+	oraw, nraw := d.GetChange(schemaOwnerAttr)
+	o := oraw.(string)
+	n := nraw.(string)
+	if n == "" {
+		return errors.New("Error setting schema owner to an empty string")
+	}
+
+	query := fmt.Sprintf("ALTER SCHEMA %s OWNER TO %s", pq.QuoteIdentifier(o), pq.QuoteIdentifier(n))
+	if _, err := conn.Query(query); err != nil {
+		return errwrap.Wrapf("Error updating schema OWNER: {{err}}", err)
+	}
+
+	return nil
+}
diff --git a/builtin/providers/postgresql/resource_postgresql_schema_test.go b/builtin/providers/postgresql/resource_postgresql_schema_test.go
index 98016d8a2..daba2d12f 100644
--- a/builtin/providers/postgresql/resource_postgresql_schema_test.go
+++ b/builtin/providers/postgresql/resource_postgresql_schema_test.go
@@ -26,6 +26,11 @@ func TestAccPostgresqlSchema_Basic(t *testing.T) {
 
 					resource.TestCheckResourceAttr(
 						"postgresql_schema.test1", "name", "foo"),
+
+					resource.TestCheckResourceAttr(
+						"postgresql_schema.test2", "name", "bar"),
+					resource.TestCheckResourceAttr(
+						"postgresql_schema.test2", "owner", "myrole3"),
 				),
 			},
 		},
@@ -112,4 +117,9 @@ resource "postgresql_role" "myrole3" {
 resource "postgresql_schema" "test1" {
   name = "foo"
 }
+
+resource "postgresql_schema" "test2" {
+  name = "bar"
+  owner = "${postgresql_role.myrole3.name}"
+}
 `
diff --git a/website/source/docs/providers/postgresql/r/postgresql_schema.html.markdown b/website/source/docs/providers/postgresql/r/postgresql_schema.html.markdown
index 3b0b1bd53..16694a5c0 100644
--- a/website/source/docs/providers/postgresql/r/postgresql_schema.html.markdown
+++ b/website/source/docs/providers/postgresql/r/postgresql_schema.html.markdown
@@ -16,7 +16,8 @@ PostgreSQL database.
 
 ```
 resource "postgresql_schema" "my_schema" {
-  name = "my_schema"
+  name  = "my_schema"
+  owner = "postgres"
 }
 ```
 
@@ -25,20 +26,21 @@ resource "postgresql_schema" "my_schema" {
 * `name` - (Required) The name of the schema. Must be unique in the PostgreSQL
   database instance where it is configured.
 
+* `owner` - (Optional) The ROLE who owns the schema.
+
 ## Import Example
 
 `postgresql_schema` supports importing resources.  Supposing the following
 Terraform:
 
 ```
-provider "postgresql" {
-  alias = "admindb"
+resource "postgresql_schema" "public" {
+  name  = "public"
 }
 
 resource "postgresql_schema" "schema_foo" {
-  provider = "postgresql.admindb"
-
-  name = "my_schema"
+  name  = "my_schema"
+  owner = "postgres"
 }
 ```