Optionally use pg_roles instead of pg_authid.

pgbedrock/attributes.py

@@ -54,9 +54,9 @@
 COLUMN_NAME_TO_KEYWORD = {v: k for k, v in PG_COLUMN_NAME.items()}
 
 
-def analyze_attributes(spec, cursor, verbose):
+def analyze_attributes(spec, cursor, verbose, attributes_source_table):
     logger.debug('Starting analyze_attributes()')
-    dbcontext = DatabaseContext(cursor, verbose)
+    dbcontext = DatabaseContext(cursor, verbose, attributes_source_table)
 
     # We disable the progress bar when showing verbose output (using '' as our bar_template)
     # or # the bar will get lost in the # output

pgbedrock/cli.py

@@ -27,8 +27,9 @@ def entrypoint():
 @click.option('--privileges/--no-privileges', default=True, help='whether to configure privileges (default: --privileges)')
 @click.option('--live/--check', default=False, help='whether to actually make changes ("live") or only show what would be changed ("check") (default: --check)')
 @click.option('--verbose/--no-verbose', default=False, help='whether to show debug-level logging messages while running (default: --no-verbose)')
+@click.option('--alternate-attributes-table/--no-alternate-attributes-table', default=False, help='whether to use pg_roles instead of pg_authid (default: --no-alternate-attributes-table)')
 def configure(spec, host, port, user, password, dbname, prompt, attributes, memberships, ownerships,
-              privileges, live, verbose):
+              privileges, live, verbose, alternate_attributes_table):
     """
     Configure the role attributes, memberships, object ownerships, and/or privileges of a
     database cluster to match a desired spec.
@@ -41,8 +42,11 @@ def configure(spec, host, port, user, password, dbname, prompt, attributes, memb
     In addition, using --verbose will print to STDOUT all debug statements and all SQL queries
     issued by pgbedrock.
     """
+
+    attributes_source_table = 'pg_roles' if alternate_attributes_table else 'pg_authid'
+
     core_configure.configure(spec, host, port, user, password, dbname, prompt, attributes,
-                             memberships, ownerships, privileges, live, verbose)
+                             memberships, ownerships, privileges, live, verbose, attributes_source_table)
 
 
 @entrypoint.command(short_help='Generate a YAML spec for a database')
@@ -53,12 +57,15 @@ def configure(spec, host, port, user, password, dbname, prompt, attributes, memb
 @click.option('-d', '--dbname', default=USER, help='database to connect to (default: "{}")'.format(USER))
 @click.option('--prompt/--no-prompt', default=False, help='prompt the user to input a password (default: --no-prompt)')
 @click.option('--verbose/--no-verbose', default=False, help='whether to show debug-level logging messages while running (default: --no-verbose)')
-def generate(host, port, user, password, dbname, prompt, verbose):
+@click.option('--alternate-attributes-table/--no-alternate-attributes-table', default=False, help='whether to use pg_roles instead of pg_authid (default: --no-alternate-attributes-table)')
+def generate(host, port, user, password, dbname, prompt, verbose, alternate_attributes_table):
     """
     Generate a YAML spec that represents the roles, memberships, ownerships, and/or privileges of a
     database.
     """
-    core_generate.generate(host, port, user, password, dbname, prompt, verbose)
+
+    attributes_source_table = 'pg_roles' if alternate_attributes_table else 'pg_authid'
+    core_generate.generate(host, port, user, password, dbname, prompt, verbose, attributes_source_table)
 
 
 if __name__ == '__main__':

pgbedrock/context.py

@@ -24,7 +24,7 @@
             (aclexplode(def.defaclacl)).privilege_type
         FROM
             pg_default_acl def
-            JOIN pg_authid auth
+            JOIN pg_roles auth
                     ON def.defaclrole = auth.oid
             JOIN pg_namespace nsp
                     ON def.defaclnamespace = nsp.oid
@@ -41,7 +41,7 @@
         subq.privilege_type
     FROM
         subq
-        JOIN pg_authid t_grantee
+        JOIN pg_roles t_grantee
             ON subq.grantee_oid = t_grantee.oid
     WHERE
         subq.grantor_oid != subq.grantee_oid
@@ -65,7 +65,7 @@
             (aclexplode(c.relacl)).privilege_type
         FROM
             pg_class c
-            JOIN pg_authid t_owner
+            JOIN pg_roles t_owner
                 ON c.relowner = t_owner.OID
             JOIN pg_namespace nsp
                 ON c.relnamespace = nsp.oid
@@ -83,7 +83,7 @@
              t_owner.rolname AS owner,
              (aclexplode(nsp.nspacl)).privilege_type
         FROM pg_namespace nsp
-        JOIN pg_authid t_owner
+        JOIN pg_roles t_owner
             ON nsp.nspowner = t_owner.OID
     ), combined AS (
         SELECT *
@@ -100,7 +100,7 @@
         combined.privilege_type
     FROM
         combined
-        JOIN pg_authid t_grantee
+        JOIN pg_roles t_grantee
             ON combined.grantee_oid = t_grantee.oid
         WHERE combined.owner != t_grantee.rolname
     ;
@@ -119,7 +119,7 @@
         rolreplication,
         rolsuper,
         rolvaliduntil
-    FROM pg_authid
+    FROM {source_table}
     WHERE rolname != 'pg_signal_backend'
     ;
     """
@@ -130,9 +130,9 @@
         auth_group.rolname AS group
     FROM
         pg_auth_members link_table
-        JOIN pg_authid auth_member
+        JOIN pg_roles auth_member
             ON link_table.member = auth_member.oid
-        JOIN pg_authid auth_group
+        JOIN pg_roles auth_group
             ON link_table.roleid = auth_group.oid
     ;
     """
@@ -191,7 +191,7 @@
         t_owner.rolname AS owner,
         co.is_dependent
     FROM combined AS co
-    JOIN pg_authid t_owner
+    JOIN pg_roles t_owner
         ON co.owner_id = t_owner.OID
     WHERE
         co.schema NOT LIKE 'pg\_t%'
@@ -201,7 +201,7 @@
 Q_GET_ALL_PERSONAL_SCHEMAS = """
     SELECT nsp.nspname
     FROM pg_namespace nsp
-        JOIN pg_authid auth
+        JOIN pg_roles auth
             ON  nsp.nspname = auth.rolname
     WHERE auth.rolcanlogin IS TRUE
     ;
@@ -256,10 +256,11 @@ class DatabaseContext(object):
         'get_version_info',
     }
 
-    def __init__(self, cursor, verbose):
+    def __init__(self, cursor, verbose, attributes_source_table='pg_authid'):
         self.cursor = cursor
         self.verbose = verbose
         self._cache = dict()
+        self._attributes_source_table = attributes_source_table
 
     def __getattribute__(self, attr):
         """ If the requested attribute should be cached and hasn't, fetch it and cache it. """
@@ -413,9 +414,13 @@ def get_role_current_nondefaults(self, rolename, object_kind, access):
         except KeyError:
             return set()
 
-    def get_all_role_attributes(self):
-        """ Return a dict with key = rolname and values = all fields in pg_authid """
-        common.run_query(self.cursor, self.verbose, Q_GET_ALL_ROLE_ATTRIBUTES)
+    def get_all_role_attributes(self, source_table=None):
+        """ Return a dict with key = rolname and values = all fields in pg_authid/pg_roles """
+
+        if source_table is None:
+            source_table = self._attributes_source_table
+
+        common.run_query(self.cursor, self.verbose, Q_GET_ALL_ROLE_ATTRIBUTES.format(source_table=source_table))
         role_attributes = {row['rolname']: dict(row) for row in self.cursor.fetchall()}
         return role_attributes
 

pgbedrock/core_configure.py

@@ -66,7 +66,7 @@ def run_password_sql(cursor, all_password_sql_to_run):
 
 
 def configure(spec_path, host, port, user, password, dbname, prompt, attributes, memberships,
-              ownerships, privileges, live, verbose):
+              ownerships, privileges, live, verbose, attributes_source_table):
     """
     Configure the role attributes, memberships, object ownerships, and/or privileges of a
     database cluster to match a desired spec.
@@ -106,6 +106,8 @@ def configure(spec_path, host, port, user, password, dbname, prompt, attributes,
 
         verbose - bool; whether to show all queries that are executed and all debug log
             messages during execution
+
+        attributes_source_table - str; the table to read use attributes from (pg_authid or pg_roles)
     """
     if verbose:
         root_logger = logging.getLogger('')
@@ -126,7 +128,7 @@ def configure(spec_path, host, port, user, password, dbname, prompt, attributes,
         sql_to_run.append(create_divider('attributes'))
         # Password changes happen within the attributes.py module itself so we don't leak
         # passwords; as a result we need to see if password changes occurred
-        module_sql, all_password_sql_to_run = analyze_attributes(spec, cursor, verbose)
+        module_sql, all_password_sql_to_run = analyze_attributes(spec, cursor, verbose, attributes_source_table)
         run_module_sql(module_sql, cursor, verbose)
         if all_password_sql_to_run:
             password_changed = True

pgbedrock/core_generate.py

@@ -382,13 +382,13 @@ def determine_nonschema_privileges_for_schema(role, objkind, objname, dbcontext)
         return all_writes, only_reads
 
 
-def create_spec(host, port, user, password, dbname, verbose):
+def create_spec(host, port, user, password, dbname, verbose, attributes_source_table):
     db_connection = common.get_db_connection(host, port, dbname, user, password)
     # We will only be reading, so it is worth being safe here and ensuring that we can't write
     db_connection.set_session(readonly=True)
     cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
 
-    dbcontext = DatabaseContext(cursor, verbose)
+    dbcontext = DatabaseContext(cursor, verbose, attributes_source_table)
     spec = initialize_spec(dbcontext)
     spec = add_attributes(spec, dbcontext)
     spec = add_memberships(spec, dbcontext)
@@ -483,7 +483,7 @@ def sort_sublists(data):
     return data
 
 
-def generate(host, port, user, password, dbname, prompt, verbose):
+def generate(host, port, user, password, dbname, prompt, verbose, attributes_source_table):
     """
     Generate a YAML spec that represents the role attributes, memberships, object ownerships,
     and privileges for all roles in a database.
@@ -508,6 +508,8 @@ def generate(host, port, user, password, dbname, prompt, verbose):
 
         verbose - bool; whether to show all queries that are executed and all debug log
             messages during execution
+
+        attributes_source_table - str; the table to read use attributes from (pg_authid or pg_roles)
     """
     if verbose:
         root_logger = logging.getLogger('')
@@ -516,6 +518,6 @@ def generate(host, port, user, password, dbname, prompt, verbose):
     if prompt:
         password = getpass.getpass()
 
-    spec = create_spec(host, port, user, password, dbname, verbose)
+    spec = create_spec(host, port, user, password, dbname, verbose, attributes_source_table)
     sorted_spec = sort_sublists(spec)
     output_spec(sorted_spec)

tests/test_attributes.py

@@ -41,7 +41,8 @@ def roleconf(request, mockdbcontext):
 @run_setup_sql([
     attr.Q_CREATE_ROLE.format(ROLE1),
     ])
-def test_analyze_attributes_modifying_objects(capsys, cursor):
+@pytest.mark.parametrize('attributes_source_table', ['pg_authid', 'pg_roles'])
+def test_analyze_attributes_modifying_objects(capsys, cursor, attributes_source_table):
     """
     End-to-end test.
     ROLE1 exists and has some non-defaults
@@ -82,7 +83,12 @@ def test_analyze_attributes_modifying_objects(capsys, cursor):
     expected.add(attr.Q_ALTER_ROLE.format(ROLE2, 'SUPERUSER'))
     expected.add(attr.Q_CREATE_ROLE.format(ROLE3))
 
-    actual, password_changed = attr.analyze_attributes(spec, cursor, verbose=False)
+    if attributes_source_table == 'pg_roles':
+        # pg_roles masks the user's password, so the password change
+        # logic is always activated (for existing users)
+        expected.add('--' + attr.Q_ALTER_PASSWORD.format(ROLE1, '******'))
+
+    actual, password_changed = attr.analyze_attributes(spec, cursor, verbose=False, attributes_source_table=attributes_source_table)
     # Filter out changes for roles that existed before this test
     actual = set([s for s in actual if ('postgres' not in s and 'test_user' not in s)])
 
@@ -421,6 +427,6 @@ def test_set_password_log_message_is_masked(capsys, cursor):
         ROLE1: {'attributes': ['PASSWORD {}'.format(new_password)]},
     }
 
-    _, password_all_sql_to_run = attr.analyze_attributes(spec, cursor, verbose=False)
+    _, password_all_sql_to_run = attr.analyze_attributes(spec, cursor, verbose=False, attributes_source_table='pg_authid')
 
     assert password_all_sql_to_run == [attr.Q_ALTER_PASSWORD.format(ROLE1, new_password)]

tests/test_core_configure.py

@@ -35,7 +35,8 @@ def test_configure_no_changes_needed(tmpdir, capsys, db_config, base_spec):
              ownerships=True,
              privileges=True,
              live=False,
-             verbose=False
+             verbose=False,
+             attributes_source_table='pg_authid'
              )
     )
     core_configure.configure(**params)
@@ -63,7 +64,8 @@ def test_configure_live_mode_works(capsys, cursor, spec_with_new_user, db_config
              ownerships=True,
              privileges=True,
              live=live_mode,
-             verbose=False
+             verbose=False,
+             attributes_source_table='pg_authid'
              )
     )
     core_configure.configure(**params)
@@ -115,6 +117,7 @@ def test_configure_live_does_not_leak_passwords(tmpdir, capsys, cursor, db_confi
              privileges=True,
              live=True,
              verbose=True,
+             attributes_source_table='pg_authid',
              )
     )
     core_configure.configure(**params)
@@ -162,7 +165,8 @@ def test_no_password_attribute_makes_password_none(cursor, spec_with_new_user, d
              ownerships=True,
              privileges=True,
              live=True,
-             verbose=False
+             verbose=False,
+             attributes_source_table='pg_authid',
              )
     )
     core_configure.configure(**params)
@@ -197,7 +201,8 @@ def test_configure_schema_role_has_dash(tmpdir, capsys, db_config, cursor, base_
              ownerships=True,
              privileges=True,
              live=False,
-             verbose=False
+             verbose=False,
+             attributes_source_table='pg_authid',
              )
     )
     core_configure.configure(**params)