Integrate object-level locks for concurrent grants (PR #595)

README.md

@@ -8,7 +8,7 @@ This intent of this fork is to maintain a working version of the provider while
 This repository aims to be up to date with the original one, and also to add some features that are needed for our use cases.
 We re-integrated the following PRs that were open on the original repository:
 
-- [Use object-level locks for concurrent grants to improve parallelism](https://github.com/cyrilgdn/terraform-provider-postgresql/pull/595)
+- [Use object-level locks for concurrent grants to improve parallelism](https://github.com/cyrilgdn/terraform-provider-postgresql/pull/595) ✅
 - [Support role configuration parameters](https://github.com/cyrilgdn/terraform-provider-postgresql/pull/305)
 
 ---

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

@@ -586,14 +586,62 @@ func pgLockRole(txn *sql.Tx, role string) error {
 	return nil
 }
 
-// Lock a database and all his members to avoid concurrent updates on some resources
-func pgLockDatabase(txn *sql.Tx, database string) error {
-	// Disable statement timeout for this connection otherwise the lock could fail
+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 {
 	if _, err := txn.Exec("SET statement_timeout = 0"); err != nil {
 		return fmt.Errorf("could not disable statement_timeout: %w", err)
 	}
-	if _, err := txn.Exec("SELECT pg_advisory_xact_lock(oid::bigint) FROM pg_database WHERE datname = $1", database); err != nil {
-		return fmt.Errorf("could not get advisory lock for database %s: %w", database, err)
+
+	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