Support setting per-role config with ALTER ROLE SET

Author: johanfleury

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/attributes.py

@@ -2,6 +2,8 @@
 import datetime as dt
 import hashlib
 import logging
+import re
+import math
 
 import click
 import psycopg2
@@ -18,7 +20,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 "{}";'
 
@@ -65,17 +69,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('config', {})
+
+            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
@@ -102,13 +108,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
@@ -119,7 +127,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):
@@ -193,6 +204,13 @@ 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):
         """ Convert the input value into a postgres rolname-salted md5 hash and compare
         it with the currently stored hash """
@@ -221,7 +239,7 @@ def set_all_attributes(self, attributes):
 
             elif current_value != desired_value and attribute != 'rolpassword':
                 self.set_attribute_value(attribute, desired_value, current_value)
-
+    
     def set_attribute_value(self, attribute, desired_value, current_value):
         if attribute == 'rolconnlimit':
             query = Q_ALTER_CONN_LIMIT.format(self.rolename, desired_value, current_value)
@@ -231,10 +249,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)
@@ -244,3 +274,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>.+)\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.floort(quantity * 24) * 60 * 60 * 1000
+
+        return value

pgbedrock/context.py

@@ -124,6 +124,8 @@
     ;
     """
 
+Q_GET_ALL_ROLE_CONFIGS = "SELECT usename, useconfig FROM pg_shadow;"
+
 Q_GET_ALL_MEMBERSHIPS = """
     SELECT
         auth_member.rolname AS member,
@@ -427,6 +429,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