def get_all_object_owners(self):
""" Return a dict of the form:
{objkindA: {
'schemaA': {
'objnameA': ownerA,
'objnameB': ownerB,
}
'schemaB': ...
...
},
objkindB:
....
}
i.e. we can access an object's owner via output[objkind][schema][objname]
This structure is chosen to allow for looking up:
1) object ownership
2) all owners in a given schema
3) all objects in a given schema
"""
OwnerRow = namedtuple('OwnerRow',
['objkind', 'schema', 'objname', 'owner'])
common.run_query(self.cursor, self.verbose, Q_GET_ALL_OBJECT_OWNERS)
all_object_owners = defaultdict(dict)
for i in self.cursor.fetchall():
row = OwnerRow(*i)
objkind_owners = all_object_owners[row.objkind]
if row.schema not in objkind_owners:
objkind_owners[row.schema] = dict()
objkind_owners[row.schema][row.objname] = row.owner
return all_object_owners
def test_run_query_fails_not_verbose_mode(capsys, cursor):
cursor.close()
with pytest.raises(SystemExit):
common.run_query(cursor, verbose=False, query='SELECT 1+1')
expected_msg = common.FAILED_QUERY_MSG.format('SELECT 1+1',
'cursor already closed\n')
assert expected_msg == capsys.readouterr()[0]
def test_run_query_fails_in_verbose_mode(capsys, cursor):
cursor.close()
with pytest.raises(SystemExit):
common.run_query(cursor, verbose=True, query='SELECT 1+1')
expected_msg = common.FAILED_QUERY_MSG.format('SELECT 1+1', '')
error = capsys.readouterr()[0]
assert error.startswith(expected_msg)
assert 'pgbedrock/common.py", line' in error
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)
role_attributes = {
row['rolname']: dict(row)
for row in self.cursor.fetchall()
}
return role_attributes
def get_all_personal_schemas(self):
""" Return all personal schemas
Returns:
set: A set of common.ObjectName instances
"""
common.run_query(self.cursor, self.verbose, Q_GET_ALL_PERSONAL_SCHEMAS)
personal_schemas = set([common.ObjectName(schema=row[0]) for row in self.cursor.fetchall()])
return personal_schemas
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
def get_all_raw_object_attributes(self):
"""
Fetch results for all object attributes.
The results are used in several subsequent methods, so having consistent results is
important. Thus, this helper method is here to ensure that we only run this query once.
"""
common.run_query(self.cursor, self.verbose,
Q_GET_ALL_RAW_OBJECT_ATTRIBUTES)
results = [ObjectAttributes(*row) for row in self.cursor.fetchall()]
return results
def get_all_nonschema_objects_and_owners(self):
""" For all objkinds other than schemas return a dict of the form
{schema_name: [(objkind, objname, objowner, is_dependent), ...]}
"""
common.run_query(self.cursor, self.verbose,
Q_GET_ALL_NONSCHEMA_OBJECTS_AND_OWNERS)
schema_objects = collections.defaultdict(list)
for i in self.cursor.fetchall():
schema = i[0]
objinfo = ObjectInfo(*i[1:])
schema_objects[schema].append(objinfo)
return schema_objects
def get_all_raw_object_attributes(self):
"""
Fetch results for all object attributes.
The results are used in several subsequent methods, so having consistent results is
important. Thus, this helper method is here to ensure that we only run this query once.
"""
common.run_query(self.cursor, self.verbose, Q_GET_ALL_RAW_OBJECT_ATTRIBUTES)
results = []
NamedRow = namedtuple('NamedRow', ['kind', 'schema', 'unqualified_name', 'owner', 'is_dependent'])
for i in self.cursor.fetchall():
row = NamedRow(*i)
objname = common.ObjectName(schema=row.schema, unqualified_name=row.unqualified_name)
entry = ObjectAttributes(row.kind, row.schema, objname, row.owner, row.is_dependent)
results.append(entry)
return results
def get_all_current_nondefaults(self):
""" Return a dict of the form:
{roleA: {
objkindA: {
'read': set([
(objname, privilege),
...
]),
'write': set([
(objname, privilege),
...
]),
},
}
roleB:
....
}
This will not include privileges granted by this role to itself
"""
NamedRow = namedtuple(
'NamedRow',
['grantee', 'objkind', 'schema', 'unqualified_name', 'privilege'])
common.run_query(self.cursor, self.verbose,
Q_GET_ALL_CURRENT_NONDEFAULTS)
current_nondefaults = defaultdict(dict)
for i in self.cursor.fetchall():
row = NamedRow(*i)
is_read_priv = row.privilege in PRIVILEGE_MAP[row.objkind]['read']
access_key = 'read' if is_read_priv else 'write'
role_nondefaults = current_nondefaults[row.grantee]
# Create this role's dict substructure for the first entry we come across
if row.objkind not in role_nondefaults:
role_nondefaults[row.objkind] = {
'read': set(),
'write': set(),
}
objname = common.ObjectName(schema=row.schema,
unqualified_name=row.unqualified_name)
entry = (objname, row.privilege)
role_nondefaults[row.objkind][access_key].add(entry)
return current_nondefaults
def test_run_query(cursor):
common.run_query(cursor, verbose=True, query='SELECT 1+1')
assert cursor.fetchone() == [2]
def get_version_info(self):
""" Return information for this Postgres instance """
common.run_query(self.cursor, self.verbose, Q_GET_VERSIONS)
results = self.cursor.fetchone()
info = VersionInfo(*results)
return info
def get_all_personal_schemas(self):
""" Return all personal schemas as a set """
common.run_query(self.cursor, self.verbose, Q_GET_ALL_PERSONAL_SCHEMAS)
personal_schemas = set([i[0] for i in self.cursor.fetchall()])
return personal_schemas
def get_all_memberships(self):
""" Return a list of tuple, where each tuple is (member, group) """
common.run_query(self.cursor, self.verbose, Q_GET_ALL_MEMBERSHIPS)
return self.cursor.fetchall()
def run_module_sql(module_sql, cursor, verbose):
if module_sql and has_changes(module_sql):
# Put all SQL into 1 string to reduce network IO of sending many small calls to Postgres
combined_sql = '\n'.join(module_sql)
common.run_query(cursor, verbose, combined_sql)
def get_all_schemas_and_owners(self):
""" Return a dict of {schema_name: schema_owner} """
common.run_query(self.cursor, self.verbose,
Q_GET_ALL_SCHEMAS_AND_OWNERS)
return {row['schema']: row['owner'] for row in self.cursor.fetchall()}