Support setting per-role config with ALTER ROLE SET (#4)

* Support setting per-role config with ALTER ROLE SET

* fix Q_GET_ALL_CURRENT_NONDEFAULTS missing partition tables

* Bump to version 0.5.0

Co-authored-by: Johan Fleury <jfleury@arcaik.net>
Co-authored-by: Dennis Zhang <dennis@gorgias.com>

CHANGELOG.md

@@ -7,10 +7,17 @@ and this project tries to adhere to [Semantic Versioning](http://semver.org/spec
 ## [Unreleased]
 _this space intentionally left blank_
 
+## [0.5.0] - 2022-09-30
+
+### In Code
+
+- BREAKING CHANGE: Add support for setting per-role config with ALTER ROLE SET (@johanfleury)
+pgbedrock will now remove any existing config that’s not defined in its configuration.
+
 ## [0.4.4] - 2022-09-06
 
 ### In Code
-- Remove Python 2 support
+- Remove Python 2 support (@johanfleury)
 - Add support for SCRAM-SHA-256 passwords (passwords can now be passed plain text or hashed) (@johanfleury)
 
 ## [0.4.3] - 2022-09-01

README.rst

@@ -40,6 +40,8 @@ As an example, the definition for the ``jdoe`` role in the spec might look like
         is_superuser: no
         attributes:
             - PASSWORD "{{ env['JDOE_PASSWORD'] }}"
+        configs:
+          statement_timeout: 42s
         member_of:
             - analyst
         owns:
@@ -75,6 +77,7 @@ When pgbedrock is run, it would make sure that:
     * ``jdoe`` is not a superuser
     * ``jdoe``'s password is the same as what is in the ``$JDOE_PASSWORD`` environment variable
     * All other role attributes for ``jdoe`` are the Postgres defaults (as defined by `pg_authid`_).
+    * ``jdoe``’s session config ``statement_timeout`` is set to ``42s``
     * ``jdoe`` is a member of the ``analyst`` role
     * ``jdoe`` is a member of no other roles
     * ``jdoe`` owns the ``finance_reports`` schema

docs/index.rst

@@ -25,6 +25,8 @@ As an example, the definition for the ``jdoe`` role in the spec might look like
         is_superuser: no
         attributes:
             - PASSWORD "{{ env['JDOE_PASSWORD'] }}"
+        configs:
+          statement_timeout: 42s
         member_of:
             - analyst
         owns:
@@ -58,6 +60,7 @@ When pgbedrock is run, it would make sure that:
     * ``jdoe`` is not a superuser
     * ``jdoe``'s password is the same as what is in the ``$JDOE_PASSWORD`` environment variable
     * All other role attributes for ``jdoe`` are the Postgres defaults (as defined by `pg_authid`_).
+    * ``jdoe``’s session config ``statement_timeout`` is set to ``42s``
     * ``jdoe`` is a member of the ``analyst`` role
     * ``jdoe`` is a member of no other roles
     * ``jdoe`` owns the ``finance_reports`` schema

pgbedrock/__init__.py

@@ -1,2 +1,2 @@
-__version__ = '0.4.3'
+__version__ = '0.5.0'
 LOG_FORMAT = '%(levelname)s:%(filename)s:%(funcName)s:%(lineno)s - %(message)s'

pgbedrock/attributes.py

@@ -4,6 +4,8 @@
 import hmac
 import base64
 import logging
+import re
+import math
 
 import click
 import psycopg2
@@ -20,7 +22,9 @@
 Q_ALTER_CONN_LIMIT = 'ALTER ROLE "{}" WITH CONNECTION LIMIT {}; -- Previous value: {}'
 Q_ALTER_PASSWORD = "ALTER ROLE \"{}\" WITH ENCRYPTED PASSWORD '{}';"
 Q_REMOVE_PASSWORD = "ALTER ROLE \"{}\" WITH PASSWORD NULL;"
-Q_ALTER_ROLE = 'ALTER ROLE "{}" WITH {};'
+Q_ALTER_ROLE_WITH = 'ALTER ROLE "{}" WITH {};'
+Q_ALTER_ROLE_SET = 'ALTER ROLE "{}" SET {}="{}"; -- Previous value: {}'
+Q_ALTER_ROLE_RESET = 'ALTER ROLE "{}" RESET {}; -- Previous value: {}'
 Q_ALTER_VALID_UNTIL = "ALTER ROLE \"{}\" WITH VALID UNTIL '{}'; -- Previous value: {}"
 Q_CREATE_ROLE = 'CREATE ROLE "{}";'
 
@@ -67,17 +71,19 @@ def analyze_attributes(spec, cursor, verbose):
                            show_eta=False, item_show_func=common.item_show_func) as all_roles:
         all_sql_to_run = []
         password_all_sql_to_run = []
-        for rolename, spec_config in all_roles:
+        for rolename, spec_settings in all_roles:
             logger.debug('Starting to analyze role {}'.format(rolename))
 
-            spec_config = spec_config or {}
-            spec_attributes = spec_config.get('attributes', [])
+            spec_settings = spec_settings or {}
+            spec_attributes = spec_settings.get('attributes', [])
 
             for keyword, attribute in (('can_login', 'LOGIN'), ('is_superuser', 'SUPERUSER')):
-                is_desired = spec_config.get(keyword, False)
+                is_desired = spec_settings.get(keyword, False)
                 spec_attributes.append(attribute if is_desired else 'NO' + attribute)
 
-            roleconf = AttributeAnalyzer(rolename, spec_attributes, dbcontext)
+            spec_configs = spec_settings.get('configs', {})
+
+            roleconf = AttributeAnalyzer(rolename, spec_attributes, spec_configs, dbcontext)
             roleconf.analyze()
             all_sql_to_run += roleconf.sql_to_run
             password_all_sql_to_run += roleconf.password_sql_to_run
@@ -112,13 +118,15 @@ class AttributeAnalyzer(object):
     make it match the provided spec attributes. Note that spec_attributes is a list whereas
     current_attributes is a dict. """
 
-    def __init__(self, rolename, spec_attributes, dbcontext):
+    def __init__(self, rolename, spec_attributes, spec_configs, dbcontext):
         self.sql_to_run = []
         self.rolename = common.check_name(rolename)
         logger.debug('self.rolename set to {}'.format(self.rolename))
         self.spec_attributes = spec_attributes
+        self.spec_configs = spec_configs
 
         self.current_attributes = dbcontext.get_role_attributes(rolename)
+        self.current_configs = dbcontext.get_role_configs(rolename)
 
         # We keep track of password-related SQL separately as we don't want running this to
         # go into the main SQL stream since that could leak password
@@ -129,7 +137,10 @@ def analyze(self):
             self.create_role()
 
         desired_attributes = self.coalesce_attributes()
+
         self.set_all_attributes(desired_attributes)
+        self.set_all_configs(self.spec_configs)
+
         return self.sql_to_run
 
     def create_role(self):
@@ -203,6 +214,12 @@ def get_attribute_value(self, attribute):
         logger.debug('Returning attribute "{}": "{}"'.format(attribute, value))
         return value
 
+    def get_config_value(self, config):
+        """ Take a config (e.g. `statement_timeout`) and look up that value in our dbcontext """
+        value = self.current_configs.get(config, "")
+        logger.debug('Returning config "{}": "{}"'.format(config, value))
+        return value
+
     def is_same_password(self, value):
         """ Compare stored password hash to the new password.
 
@@ -275,10 +292,22 @@ def set_attribute_value(self, attribute, desired_value, current_value):
             base_keyword = COLUMN_NAME_TO_KEYWORD[attribute]
             # prepend 'NO' if desired_value is False
             keyword = base_keyword if desired_value else 'NO' + base_keyword
-            query = Q_ALTER_ROLE.format(self.rolename, keyword)
+            query = Q_ALTER_ROLE_WITH.format(self.rolename, keyword)
 
         self.sql_to_run.append(query)
 
+    def set_all_configs(self, configs):
+        for config, desired_value in configs.items():
+            current_value = self.get_config_value(config)
+
+            if self.parse_config_value(current_value) != self.parse_config_value(desired_value):
+                self.sql_to_run.append(Q_ALTER_ROLE_SET.format(self.rolename, config, desired_value, current_value))
+
+        if self.current_configs:
+            for current_config, current_value in self.current_configs.items():
+                if current_config not in configs.keys():
+                    self.sql_to_run.append(Q_ALTER_ROLE_RESET.format(self.rolename, current_config, current_value))
+
     def set_password(self, desired_value):
         if desired_value is None:
             actual_query = Q_REMOVE_PASSWORD.format(self.rolename)
@@ -288,3 +317,70 @@ def set_password(self, desired_value):
 
         sanitized_query = Q_ALTER_PASSWORD.format(self.rolename, '******')
         self.sql_to_run.append('--' + sanitized_query)
+
+    @staticmethod
+    def parse_config_value(value):
+        """
+        Parse a config (as in postgresql.conf setting) value and return it’s normalized value.
+
+        If the value matches a well-known unit it is rounded to a multiple of the
+        next smaller unit (if there is one) then it is converted to kB for memory
+        units and to ms for time units
+
+        If the value doesn’t match a well known unit it is returned as an int or a
+        float if conversion is possible or as-is otherwise.
+
+        See: https://www.postgresql.org/docs/current/config-setting.html
+        """
+
+        if not isinstance(value, str):
+            return value
+
+        try:
+            return int(value)
+        except ValueError:
+            pass
+
+        try:
+            return float(value)
+        except ValueError:
+            pass
+
+        # Valid memory units are B (bytes), kB (kilobytes), MB (megabytes), GB (gigabytes), and TB (terabytes).
+        # The multiplier for memory units is 1024, not 1000.
+        m = re.search("(?P<quantity>[\-\+]?[0-9]+(\.[0-9]+)?)\s*(?P<unit>B|kB|MB|GB|TB)", value)
+        if m:
+            quantity = float(m.group("quantity"))
+            unit = m.group("unit")
+
+            if unit == "B":
+                return quantity / 1024
+            if unit == "kB":
+                return math.floor(quantity * 1024) / 1024
+            if unit == "MB":
+                return math.floor(quantity * 1024)
+            if unit == "GB":
+                return math.floor(quantity * 1024) * 1024
+            if unit == "TB":
+                return math.floor(quantity * 1024) * 1024 ** 2
+
+        # Valid time units are us (microseconds), ms (milliseconds), s (seconds), min (minutes), h (hours), and d (days).
+        m = re.search("(?P<quantity>[\-\+]?[0-9]+(\.[0-9]+)?)\s*(?P<unit>us|ms|s|min|h|d)", value)
+        if m:
+            quantity = float(m.group("quantity"))
+            unit = m.group("unit")
+
+            if unit == "us":
+                return quantity / 1000
+            if unit == "ms":
+                return math.floor(quantity * 1000) / 1000
+            if unit == "s":
+                return math.floor(quantity * 1000)
+            if unit == "min":
+                return math.floor(quantity * 60) * 1000
+            if unit == "h":
+                return math.floor(quantity * 60) * 60 * 1000
+            if unit == "d":
+                return math.floor(quantity * 24) * 60 * 60 * 1000
+
+        return value

pgbedrock/context.py

@@ -54,6 +54,7 @@
                ('v', 'tables'),
                ('m', 'tables'),
                ('f', 'tables'),
+               ('p', 'tables'),
                ('S', 'sequences')
     ), tables_and_sequences AS (
         SELECT
@@ -124,6 +125,8 @@
     ;
     """
 
+Q_GET_ALL_ROLE_CONFIGS = "SELECT usename, useconfig FROM pg_shadow;"
+
 Q_GET_ALL_MEMBERSHIPS = """
     SELECT
         auth_member.rolname AS member,
@@ -250,6 +253,7 @@ class DatabaseContext(object):
         'get_all_current_nondefaults',
         'get_all_object_attributes',
         'get_all_role_attributes',
+        'get_all_role_configs',
         'get_all_memberships',
         'get_all_nonschema_objects_and_owners',
         'get_all_personal_schemas',
@@ -428,6 +432,23 @@ def is_superuser(self, rolename):
         role_attributes = self.get_role_attributes(rolename)
         return role_attributes.get('rolsuper', False)
 
+    def get_all_role_configs(self):
+        """ Return a dict of {rolname: {config_name: value}} from pg_shadow"""
+        common.run_query(self.cursor, self.verbose, Q_GET_ALL_ROLE_CONFIGS)
+        role_configs = {}
+        for row in self.cursor.fetchall():
+            role_configs[row['usename']] = {}
+
+            for config in row['useconfig'] or []:
+                k, v = config.split('=')
+                role_configs[row['usename']][k] = v
+
+        return role_configs
+
+    def get_role_configs(self, rolename):
+        all_role_configs = self.get_all_role_configs()
+        return all_role_configs.get(rolename, dict())
+
     def get_all_raw_object_attributes(self):
         """
         Fetch results for all object attributes.

pgbedrock/memberships.py

@@ -24,9 +24,9 @@ def analyze_memberships(spec, cursor, verbose):
     with click.progressbar(spec.items(), label='Analyzing memberships:', bar_template=bar_template,
                            show_eta=False, item_show_func=common.item_show_func) as all_roles:
         all_sql_to_run = []
-        for rolename, spec_config in all_roles:
-            spec_config = spec_config or {}
-            spec_memberships = set(spec_config.get('member_of', []))
+        for rolename, spec_settings in all_roles:
+            spec_settings = spec_settings or {}
+            spec_memberships = set(spec_settings.get('member_of', []))
             sql_to_run = MembershipAnalyzer(rolename, spec_memberships, dbcontext).analyze()
             all_sql_to_run += sql_to_run
 

pgbedrock/spec_inspector.py

@@ -54,6 +54,8 @@
                 - NOLOGIN
                 - SUPERUSER
                 - NOSUPERUSER
+    configs:
+        type: dict
     member_of:
         type: list
         schema:
@@ -490,4 +492,4 @@ def ensure_no_except_on_schema(spec):
         if config and config.get('privileges'):
             if config['privileges'].get('schemas') and config['privileges']['schemas'].get('except') and config['privileges']['schemas']['excepted']:
                 error_messages.append(EXCEPTED_SCHEMAS_MSG.format(rolename))
-    return error_messages
+    return error_messages