Use object-level locks for concurrent grants to improve parallelism

examples/issues/178/dev.tfrc

@@ -0,0 +1,6 @@
+provider_installation {
+  dev_overrides {
+    "cyrilgdn/postgresql" = "../../../"
+  }
+  direct {}
+}

examples/issues/178/locals.tf

@@ -0,0 +1,4 @@
+locals {
+  read_only_users  = toset([for i in range(var.user_ro_count) : "user_ro_${i}"])
+  read_write_users = toset([for i in range(var.user_rw_count) : "user_rw_${i}"])
+}

examples/issues/178/main.tf

@@ -0,0 +1,191 @@
+terraform {
+  required_providers {
+    docker = {
+      source  = "kreuzwerker/docker"
+      version = ">= 3.0.2"
+    }
+    postgresql = {
+      source  = "cyrilgdn/postgresql"
+      version = ">= 1.25"
+    }
+  }
+}
+
+provider "docker" {
+  host = var.docker_host
+}
+
+resource "docker_image" "postgres" {
+  name         = var.postgres_image
+  keep_locally = var.keep_image
+}
+
+resource "docker_container" "postgres" {
+  image = docker_image.postgres.image_id
+  name  = "postgres"
+  wait  = true
+  ports {
+    internal = var.POSTGRES_PORT
+    external = var.POSTGRES_PORT
+  }
+  env = [
+    "POSTGRES_PASSWORD=${var.POSTGRES_PASSWORD}"
+  ]
+  healthcheck {
+    test         = ["CMD-SHELL", "pg_isready"]
+    interval     = "5s"
+    timeout      = "5s"
+    retries      = 5
+    start_period = "2s"
+  }
+  upload {
+    file    = "/docker-entrypoint-initdb.d/mock-tables.sql"
+    content = <<EOS
+      CREATE DATABASE "test" OWNER "${var.POSTGRES_DBNAME}";
+      \connect ${var.POSTGRES_DBNAME}
+
+      DO $$
+      DECLARE
+        table_count int := ${var.table_count};
+      BEGIN
+        FOR count IN 0..table_count LOOP
+          EXECUTE format('CREATE TABLE table_%s (test int)', count);
+        END LOOP;
+      END $$;
+    EOS
+  }
+}
+
+provider "postgresql" {
+  scheme      = "postgres"
+  host        = var.POSTGRES_HOST
+  port        = docker_container.postgres.ports[0].external
+  database    = var.POSTGRES_PASSWORD
+  username    = var.POSTGRES_PASSWORD
+  password    = var.POSTGRES_PASSWORD
+  sslmode     = "disable"
+  superuser   = false
+  lock_grants = true
+}
+
+resource "postgresql_role" "readonly_role" {
+  name             = "readonly"
+  login            = false
+  superuser        = false
+  create_database  = false
+  create_role      = false
+  inherit          = false
+  replication      = false
+  connection_limit = -1
+}
+
+resource "postgresql_role" "readwrite_role" {
+  name             = "readwrite"
+  login            = false
+  superuser        = false
+  create_database  = false
+  create_role      = false
+  inherit          = false
+  replication      = false
+  connection_limit = -1
+}
+
+resource "postgresql_grant" "readonly_role" {
+  database          = var.POSTGRES_DBNAME
+  role              = postgresql_role.readonly_role.name
+  object_type       = "table"
+  schema            = "public"
+  privileges        = ["SELECT"]
+  with_grant_option = false
+}
+
+resource "postgresql_grant" "readwrite_role" {
+  database          = var.POSTGRES_DBNAME
+  role              = postgresql_role.readwrite_role.name
+  object_type       = "table"
+  schema            = "public"
+  privileges        = ["SELECT", "INSERT", "UPDATE", "DELETE"]
+  with_grant_option = false
+}
+
+resource "postgresql_role" "readonly_users" {
+  for_each         = local.read_only_users
+  name             = each.value
+  roles            = [postgresql_role.readonly_role.name]
+  login            = true
+  superuser        = false
+  create_database  = false
+  create_role      = false
+  inherit          = true
+  replication      = false
+  connection_limit = -1
+}
+
+resource "postgresql_role" "readwrite_users" {
+  for_each         = local.read_write_users
+  name             = each.value
+  roles            = [postgresql_role.readonly_role.name]
+  login            = true
+  superuser        = false
+  create_database  = false
+  create_role      = false
+  inherit          = true
+  replication      = false
+  connection_limit = -1
+}
+
+resource "postgresql_grant" "connect_db_readonly_role" {
+  for_each    = postgresql_role.readonly_users
+  database    = var.POSTGRES_DBNAME
+  object_type = "database"
+  privileges  = ["CREATE", "CONNECT"]
+  role        = each.value.name
+}
+
+resource "postgresql_grant" "connect_db_readwrite_role" {
+  for_each    = postgresql_role.readwrite_users
+  database    = var.POSTGRES_DBNAME
+  object_type = "database"
+  privileges  = ["CREATE", "CONNECT"]
+  role        = each.value.name
+}
+
+resource "postgresql_grant" "usage_readonly_role" {
+  for_each          = postgresql_role.readonly_users
+  database          = var.POSTGRES_DBNAME
+  role              = each.value.name
+  object_type       = "schema"
+  schema            = "public"
+  privileges        = ["USAGE"]
+  with_grant_option = false
+}
+
+resource "postgresql_grant" "usage_readwrite_role" {
+  for_each          = postgresql_role.readwrite_users
+  database          = var.POSTGRES_DBNAME
+  role              = each.value.name
+  object_type       = "schema"
+  schema            = "public"
+  privileges        = ["USAGE"]
+  with_grant_option = false
+}
+
+resource "postgresql_grant" "select_readonly_role" {
+  for_each          = postgresql_role.readonly_users
+  database          = var.POSTGRES_DBNAME
+  role              = each.value.name
+  object_type       = "table"
+  schema            = "public"
+  privileges        = ["SELECT"]
+  with_grant_option = false
+}
+
+resource "postgresql_grant" "crud_readwrite_role" {
+  for_each          = postgresql_role.readwrite_users
+  database          = var.POSTGRES_DBNAME
+  role              = each.value.name
+  object_type       = "table"
+  schema            = "public"
+  privileges        = ["SELECT", "UPDATE", "INSERT", "DELETE"]
+  with_grant_option = false
+}

examples/issues/178/test.tftest.hcl

@@ -0,0 +1,3 @@
+run "concurrent_grants" {
+  command = apply
+}

examples/issues/178/variables.tf

@@ -0,0 +1,73 @@
+variable "postgres_image" {
+  description = "Which postgres docker image to use"
+  default     = "postgres:17"
+  type        = string
+  sensitive   = false
+}
+
+variable "docker_host" {
+  description = "Socket path to docker host to use for testing"
+  default     = "unix:///var/run/docker.sock"
+  type        = string
+  sensitive   = false
+}
+
+variable "table_count" {
+  description = "Number of mock tables to create"
+  default     = 300
+  type        = number
+  sensitive   = false
+}
+
+variable "user_ro_count" {
+  description = "Number of mock RO users to create"
+  default     = 30
+  type        = number
+  sensitive   = false
+}
+
+variable "user_rw_count" {
+  description = "Number of mock RW users to create"
+  default     = 30
+  type        = number
+  sensitive   = false
+}
+
+variable "POSTGRES_DBNAME" {
+  default   = "postgres"
+  type      = string
+  sensitive = false
+}
+
+variable "POSTGRES_USER" {
+  default   = "postgres"
+  type      = string
+  sensitive = false
+}
+
+variable "POSTGRES_PASSWORD" {
+  description = "Password for docker POSTGRES_USER"
+  default     = "postgres"
+  type        = string
+  sensitive   = false
+}
+
+variable "POSTGRES_HOST" {
+  default   = "127.0.0.1"
+  type      = string
+  sensitive = false
+}
+
+variable "POSTGRES_PORT" {
+  description = "Which port postgres should listen on."
+  default     = 5432
+  type        = number
+  sensitive   = false
+}
+
+variable "keep_image" {
+  description = "If true, then the Docker image won't be deleted on destroy operation. If this is false, it will delete the image from the docker local storage on destroy operation."
+  default     = true
+  type        = bool
+  sensitive   = false
+}

postgresql/config.go

@@ -159,6 +159,10 @@ func (db *DBConnection) isSuperuser() (bool, error) {
 	return superuser, nil
 }
 
+func (db *DBConnection) IsLockGrants() bool {
+	return db.client.IsLockGrants()
+}
+
 type ClientCertificateConfig struct {
 	CertificatePath string
 	KeyPath         string
@@ -183,6 +187,7 @@ type Config struct {
 	SSLClientCert                   *ClientCertificateConfig
 	SSLRootCertPath                 string
 	GCPIAMImpersonateServiceAccount string
+	LockGrants                      bool
 }
 
 // Client struct holding connection string
@@ -333,6 +338,10 @@ func (c *Client) Connect() (*DBConnection, error) {
 	return conn, nil
 }
 
+func (c *Client) IsLockGrants() bool {
+	return c.config.LockGrants
+}
+
 // fingerprintCapabilities queries PostgreSQL to populate a local catalog of
 // capabilities.  This is only run once per Client.
 func fingerprintCapabilities(db *sql.DB) (*semver.Version, error) {

postgresql/helpers.go

@@ -599,6 +599,63 @@ func pgLockDatabase(txn *sql.Tx, database string) error {
 	return nil
 }
 
+func generateGrantLockID(database, objectType, schema, objectName string) string {
+	switch objectType {
+	case "database":
+		return fmt.Sprintf("grant:db:%s", database)
+
+	case "schema":
+		return fmt.Sprintf("grant:schema:%s.%s", database, schema)
+
+	case "foreign_data_wrapper":
+		return fmt.Sprintf("grant:fdw:%s.%s", database, objectName)
+
+	case "foreign_server":
+		return fmt.Sprintf("grant:srv:%s.%s", database, objectName)
+
+	case "table", "sequence", "column":
+		if objectName == "" {
+			return fmt.Sprintf("grant:schema:%s.%s", database, schema)
+		}
+		return fmt.Sprintf("grant:%s:%s.%s.%s", objectType, database, schema, objectName)
+
+	case "function", "procedure", "routine":
+		if objectName == "" {
+			return fmt.Sprintf("grant:schema:%s.%s", database, schema)
+		}
+		funcName := strings.Split(objectName, "(")[0]
+		return fmt.Sprintf("grant:%s:%s.%s.%s", objectType, database, schema, funcName)
+
+	default:
+		return fmt.Sprintf("grant:db:%s", database)
+	}
+}
+
+func pgLockGrantTarget(txn *sql.Tx, d *schema.ResourceData) error {
+	database := d.Get("database").(string)
+	objectType := d.Get("object_type").(string)
+	schemaName := d.Get("schema").(string)
+	objects := d.Get("objects").(*schema.Set)
+
+	if objects.Len() == 0 || objectType == "database" || objectType == "schema" {
+		lockID := generateGrantLockID(database, objectType, schemaName, "")
+		if _, err := txn.Exec("SELECT pg_advisory_xact_lock(hashtext($1)::bigint)", lockID); err != nil {
+			return fmt.Errorf("could not acquire advisory lock for %s: %w", lockID, err)
+		}
+		return nil
+	}
+
+	for _, obj := range objects.List() {
+		objectName := obj.(string)
+		lockID := generateGrantLockID(database, objectType, schemaName, objectName)
+		if _, err := txn.Exec("SELECT pg_advisory_xact_lock(hashtext($1)::bigint)", lockID); err != nil {
+			return fmt.Errorf("could not acquire advisory lock for %s: %w", lockID, err)
+		}
+	}
+
+	return nil
+}
+
 func arrayDifference(a, b []any) (diff []any) {
 	m := make(map[any]bool)