From 444aee2bf903f3ad2de181405a90ca737b2d03f6 Mon Sep 17 00:00:00 2001
From: Jille Timmermans <jille@snoozethis.com>
Date: Wed, 25 Jan 2023 15:04:17 +0100
Subject: [PATCH] feat: Support for privileges on procedural languages

---
 README.md               |  2 +-
 dump.go                 |  9 +++++++++
 privileges.go           | 40 ++++++++++++++++++++++++++++++++++++++++
 schema.go               |  2 +-
 testdata/languages.yaml | 10 ++++++++++
 5 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 testdata/languages.yaml

diff --git a/README.md b/README.md
index 6322148..0bcaa97 100644
--- a/README.md
+++ b/README.md
@@ -142,7 +142,7 @@ Types work similarly as the others. For the purposes of pgperms you should consi
 
 We'll happily accept your contributions! There's still a lot of things not supported:
 
-- Permissions on columns, foreign data wrappers, foreign servers, routines, languages, large objects or tablespaces.
+- Permissions on columns, foreign data wrappers, foreign servers, routines, large objects or tablespaces.
 - Set up default privileges so that newly created tables already have the correct permissions without having to run pgperms?
 - A config setting to automatically manage all users (and thus delete any unlisted users without needing to tombstone them).
 - More test cases
diff --git a/dump.go b/dump.go
index f123f9d..9e1e23f 100644
--- a/dump.go
+++ b/dump.go
@@ -21,6 +21,7 @@ func Dump(ctx context.Context, conns *Connections) (string, error) {
 	c.DatabasePrivileges = mergePrivileges(c.DatabasePrivileges)
 	c.SchemaPrivileges = mergePrivileges(c.SchemaPrivileges)
 	c.TypePrivileges = mergePrivileges(c.TypePrivileges)
+	c.LanguagePrivileges = mergePrivileges(c.LanguagePrivileges)
 	b, err := yaml.Marshal(c)
 	if err != nil {
 		return "", err
@@ -83,6 +84,12 @@ func Gather(ctx context.Context, conns *Connections, interestingRoles, interesti
 		}
 		ret.TypePrivileges = append(ret.TypePrivileges, typPrivs...)
 
+		langPrivs, err := fetchLanguagePrivileges(ctx, dbconn, dbname, interestingRoles)
+		if err != nil {
+			return nil, err
+		}
+		ret.LanguagePrivileges = append(ret.LanguagePrivileges, langPrivs...)
+
 		derefNow(true)
 	}
 	return &ret, nil
@@ -131,5 +138,7 @@ func Sync(ctx context.Context, conns *Connections, desired []byte, ss SyncSink)
 	SyncPrivileges(ss, d.Databases, actual.TablePrivileges, d.TablePrivileges)
 	ss.AddBarrier()
 	SyncPrivileges(ss, d.Databases, actual.SequencePrivileges, d.SequencePrivileges)
+	ss.AddBarrier()
+	SyncPrivileges(ss, d.Databases, actual.LanguagePrivileges, d.LanguagePrivileges)
 	return nil
 }
diff --git a/privileges.go b/privileges.go
index 3e49d65..4247537 100644
--- a/privileges.go
+++ b/privileges.go
@@ -334,6 +334,46 @@ func fetchTypePrivileges(ctx context.Context, conn *pgx.Conn, database string, i
 	return privs, nil
 }
 
+func fetchLanguagePrivileges(ctx context.Context, conn *pgx.Conn, database string, interestingUsers []string) ([]GenericPrivilege, error) {
+	rows, err := conn.Query(ctx, "SELECT lanname, pg_get_userbyid(grantee) AS grantee, privilege_type, is_grantable FROM pg_catalog.pg_language, aclexplode(lanacl) WHERE pg_get_userbyid(grantee) = ANY($1)", interestingUsers)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	grouped := map[string]map[string]map[bool]privilegeSet{}
+	for rows.Next() {
+		var lang, grantee, privilege string
+		var grantable bool
+		if err := rows.Scan(&lang, &grantee, &privilege, &grantable); err != nil {
+			return nil, err
+		}
+		fqln := joinSchemaName(database, lang)
+		if grouped[grantee][fqln] == nil {
+			if grouped[grantee] == nil {
+				grouped[grantee] = map[string]map[bool]privilegeSet{}
+			}
+			grouped[grantee][fqln] = map[bool]privilegeSet{}
+		}
+		ps := grouped[grantee][fqln][grantable]
+		ps.Add(privilege)
+		grouped[grantee][fqln][grantable] = ps
+	}
+	var privs []GenericPrivilege
+	for grantee, tmp1 := range grouped {
+		for fqln, tmp2 := range tmp1 {
+			for grantable, ps := range tmp2 {
+				privs = append(privs, GenericPrivilege{
+					Roles:      []string{grantee},
+					Languages:  []string{fqln},
+					Privileges: ps.ListOrAll("languages"),
+					Grantable:  grantable,
+				})
+			}
+		}
+	}
+	return privs, nil
+}
+
 type privilegeSet int
 
 func (ps *privilegeSet) Add(priv string) {
diff --git a/schema.go b/schema.go
index 8adce93..f11e31c 100644
--- a/schema.go
+++ b/schema.go
@@ -19,7 +19,7 @@ type Config struct {
 	// ForeignDataWrapperPrivileges []GenericPrivilege `yaml:"foreign_data_wrapper_privileges,omitempty"`
 	// ForeignServerPrivileges      []GenericPrivilege `yaml:"foreign_server_privileges,omitempty"`
 	// RoutinePrivileges            []GenericPrivilege `yaml:"routine_privileges,omitempty"`
-	// LanguagePrivileges           []GenericPrivilege `yaml:"language_privileges,omitempty"`
+	LanguagePrivileges []GenericPrivilege `yaml:"language_privileges,omitempty"`
 	// LargeObjectPrivileges        []GenericPrivilege `yaml:"large_object_privileges,omitempty"`
 	// TablespacePrivileges         []GenericPrivilege `yaml:"tablespace_privileges,omitempty"`
 	TypePrivileges []GenericPrivilege `yaml:"type_privileges,omitempty"`
diff --git a/testdata/languages.yaml b/testdata/languages.yaml
new file mode 100644
index 0000000..815821e
--- /dev/null
+++ b/testdata/languages.yaml
@@ -0,0 +1,10 @@
+preparation:
+  - CREATE USER someone
+  - GRANT USAGE ON LANGUAGE plpgsql TO someone
+config:
+  roles:
+    someone:
+  databases:
+    - postgres
+expected:
+- "/*                 postgres */ REVOKE USAGE ON LANGUAGE plpgsql FROM someone"