def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace):
params = dict(enc=encoding, collate=lc_collate, ctype=lc_ctype, conn_limit=conn_limit, tablespace=tablespace)
if not db_exists(cursor, db):
query_fragments = ['CREATE DATABASE %s' % pg_quote_identifier(db, 'database')]
if owner:
query_fragments.append('OWNER "%s"' % owner)
if template:
query_fragments.append('TEMPLATE %s' % pg_quote_identifier(template, 'database'))
if encoding:
query_fragments.append('ENCODING %(enc)s')
if lc_collate:
query_fragments.append('LC_COLLATE %(collate)s')
if lc_ctype:
query_fragments.append('LC_CTYPE %(ctype)s')
if tablespace:
query_fragments.append('TABLESPACE %s' % pg_quote_identifier(tablespace, 'tablespace'))
if conn_limit:
query_fragments.append("CONNECTION LIMIT %(conn_limit)s" % {"conn_limit": conn_limit})
query = ' '.join(query_fragments)
cursor.execute(query, params)
return True
else:
db_info = get_db_info(cursor, db)
if (encoding and get_encoding_id(cursor, encoding) != db_info['encoding_id']):
raise NotSupportedError(
'Changing database encoding is not supported. '
'Current encoding: %s' % db_info['encoding']
)
elif lc_collate and lc_collate != db_info['lc_collate']:
raise NotSupportedError(
'Changing LC_COLLATE is not supported. '
'Current LC_COLLATE: %s' % db_info['lc_collate']
)
elif lc_ctype and lc_ctype != db_info['lc_ctype']:
raise NotSupportedError(
'Changing LC_CTYPE is not supported.'
'Current LC_CTYPE: %s' % db_info['lc_ctype']
)
else:
changed = False
if owner and owner != db_info['owner']:
changed = set_owner(cursor, db, owner)
if conn_limit and conn_limit != str(db_info['conn_limit']):
changed = set_conn_limit(cursor, db, conn_limit)
if tablespace and tablespace != db_info['tablespace']:
changed = set_tablespace(cursor, db, tablespace)
return changed
def __pub_set_owner(self, role, check_mode=False):
"""Set a publication owner.
Args:
role (str): Role (user) name that needs to be set as a publication owner.
Kwargs:
check_mode (bool): If True, don't actually change anything,
just make SQL, add it to ``self.executed_queries`` and return True.
Returns:
True if successful, False otherwise.
"""
query = ("ALTER PUBLICATION %s OWNER TO %s" % (pg_quote_identifier(
self.name, 'publication'), pg_quote_identifier(role, 'role')))
return self.__exec_sql(query, check_mode=check_mode)
def __pub_drop_table(self, table, check_mode=False):
"""Drop a table from the publication.
Args:
table (str): Table name.
Kwargs:
check_mode (bool): If True, don't actually change anything,
just make SQL, add it to ``self.executed_queries`` and return True.
Returns:
True if successful, False otherwise.
"""
query = ("ALTER PUBLICATION %s DROP TABLE %s" % (pg_quote_identifier(
self.name, 'publication'), pg_quote_identifier(table, 'table')))
return self.__exec_sql(query, check_mode=check_mode)
def db_delete(cursor, db):
if db_exists(cursor, db):
query = "DROP DATABASE %s" % pg_quote_identifier(db, 'database')
cursor.execute(query)
return True
else:
return False
def copy_from(self):
"""Implements COPY FROM command behavior."""
self.src = self.module.params['copy_from']
self.dst = self.module.params['dst']
query_fragments = ['COPY %s' % pg_quote_identifier(self.dst, 'table')]
if self.module.params.get('columns'):
query_fragments.append('(%s)' %
','.join(self.module.params['columns']))
query_fragments.append('FROM')
if self.module.params.get('program'):
query_fragments.append('PROGRAM')
query_fragments.append("'%s'" % self.src)
if self.module.params.get('options'):
query_fragments.append(self.__transform_options())
# Note: check mode is implemented here:
if self.module.check_mode:
self.changed = self.__check_table(self.dst)
if self.changed:
self.executed_queries.append(' '.join(query_fragments))
else:
if exec_sql(self, ' '.join(query_fragments), ddl=True):
self.changed = True
def revoke_table_privileges(cursor, user, table, privs):
# Note: priv escaped by parse_privs
privs = ', '.join(privs)
query = 'REVOKE %s ON TABLE %s FROM "%s"' % (
privs, pg_quote_identifier(table, 'table'), user)
executed_queries.append(query)
cursor.execute(query)
def set_owner(self):
"""Implements ALTER SEQUENCE OWNER TO command behavior."""
query = ['ALTER SEQUENCE']
query.append(self.__add_schema())
query.append('OWNER TO %s' %
pg_quote_identifier(self.module.params['owner'], 'role'))
return exec_sql(self, ' '.join(query), ddl=True)
def __pub_set_tables(self, tables, check_mode=False):
"""Set a table suit that need to be published by the publication.
Args:
tables (list): List of tables.
Kwargs:
check_mode (bool): If True, don't actually change anything,
just make SQL, add it to ``self.executed_queries`` and return True.
Returns:
True if successful, False otherwise.
"""
quoted_tables = [pg_quote_identifier(t, 'table') for t in tables]
query = ("ALTER PUBLICATION %s SET TABLE %s" % (pg_quote_identifier(
self.name, 'publication'), ', '.join(quoted_tables)))
return self.__exec_sql(query, check_mode=check_mode)
def drop(self, cascade=False):
if not self.exists:
return False
query = "DROP TABLE %s" % pg_quote_identifier(self.name, 'table')
if cascade:
query += " CASCADE"
return exec_sql(self, query, ddl=True)
def transform_tables_representation(tbl_list):
"""Add 'public.' to names of tables where a schema identifier is absent
and add quotes to each element.
Args:
tbl_list (list): List of table names.
Returns:
tbl_list (list): Changed list.
"""
for i, table in enumerate(tbl_list):
if '.' not in table:
tbl_list[i] = pg_quote_identifier('public.%s' % table.strip(),
'table')
else:
tbl_list[i] = pg_quote_identifier(table.strip(), 'table')
return tbl_list
def set_schema(self):
"""Implements ALTER SEQUENCE SET SCHEMA command behavior."""
query = ['ALTER SEQUENCE']
query.append(self.__add_schema())
query.append(
'SET SCHEMA %s' %
pg_quote_identifier(self.module.params['newschema'], 'schema'))
return exec_sql(self, ' '.join(query), ddl=True)
def rename(self):
"""Implements ALTER SEQUENCE RENAME TO command behavior."""
query = ['ALTER SEQUENCE']
query.append(self.__add_schema())
query.append(
'RENAME TO %s' %
pg_quote_identifier(self.module.params['rename_to'], 'sequence'))
return exec_sql(self, ' '.join(query), ddl=True)
def schema_create(cursor, schema, owner):
if not schema_exists(cursor, schema):
query_fragments = [
'CREATE SCHEMA %s' % pg_quote_identifier(schema, 'schema')
]
if owner:
query_fragments.append('AUTHORIZATION %s' %
pg_quote_identifier(owner, 'role'))
query = ' '.join(query_fragments)
cursor.execute(query)
executed_queries.append(query)
return True
else:
schema_info = get_schema_info(cursor, schema)
if owner and owner != schema_info['owner']:
return set_owner(cursor, schema, owner)
else:
return False
def schema_delete(cursor, schema, cascade):
if schema_exists(cursor, schema):
query = "DROP SCHEMA %s" % pg_quote_identifier(schema, 'schema')
if cascade:
query += " CASCADE"
cursor.execute(query)
executed_queries.append(query)
return True
else:
return False
def __pub_set_param(self, param, value, check_mode=False):
"""Set an optional publication parameter.
Args:
param (str): Name of the parameter.
value (str): Parameter value.
Kwargs:
check_mode (bool): If True, don't actually change anything,
just make SQL, add it to ``self.executed_queries`` and return True.
Returns:
True if successful, False otherwise.
"""
query = ("ALTER PUBLICATION %s SET (%s = '%s')" %
(pg_quote_identifier(self.name, 'publication'), param, value))
return self.__exec_sql(query, check_mode=check_mode)
def create(self, tables, params, owner, check_mode=True):
"""Create the publication.
Args:
tables (list): List with names of the tables that need to be added to the publication.
params (dict): Dict contains optional publication parameters and their values.
owner (str): Name of the publication owner.
Kwargs:
check_mode (bool): If True, don't actually change anything,
just make SQL, add it to ``self.executed_queries`` and return True.
Returns:
changed (bool): True if publication has been created, otherwise False.
"""
changed = True
query_fragments = [
"CREATE PUBLICATION %s" %
pg_quote_identifier(self.name, 'publication')
]
if tables:
query_fragments.append("FOR TABLE %s" % ', '.join(tables))
else:
query_fragments.append("FOR ALL TABLES")
if params:
params_list = []
# Make list ["param = 'value'", ...] from params dict:
for (key, val) in iteritems(params):
params_list.append("%s = '%s'" % (key, val))
# Add the list to query_fragments:
query_fragments.append("WITH (%s)" % ', '.join(params_list))
changed = self.__exec_sql(' '.join(query_fragments),
check_mode=check_mode)
if owner:
# If check_mode, just add possible SQL to
# executed_queries and return:
self.__pub_set_owner(owner, check_mode=check_mode)
return changed
def check_pub(self):
"""Check the publication and refresh ``self.attrs`` publication attribute.
Returns:
True if the publication with ``self.name`` exists, False otherwise.
"""
pub_info = self.__get_general_pub_info()
if not pub_info:
# Publication does not exist:
return False
self.attrs['owner'] = pub_info.get('pubowner')
# Publication DML operations:
self.attrs['parameters']['publish'] = {}
self.attrs['parameters']['publish']['insert'] = pub_info.get(
'pubinsert', False)
self.attrs['parameters']['publish']['update'] = pub_info.get(
'pubupdate', False)
self.attrs['parameters']['publish']['delete'] = pub_info.get(
'pubdelete', False)
if pub_info.get('pubtruncate'):
self.attrs['parameters']['publish']['truncate'] = pub_info.get(
'pubtruncate')
# If alltables flag is False, get the list of targeted tables:
if not pub_info.get('puballtables'):
table_info = self.__get_tables_pub_info()
# Join sublists [['schema', 'table'], ...] to ['schema.table', ...]
# for better representation:
for i, schema_and_table in enumerate(table_info):
table_info[i] = pg_quote_identifier('.'.join(schema_and_table),
'table')
self.attrs['tables'] = table_info
else:
self.attrs['alltables'] = True
# Publication exists:
return True
def drop(self, cascade=False, check_mode=True):
"""Drop the publication.
Kwargs:
cascade (bool): Flag indicates that publication needs to be deleted
with its dependencies.
check_mode (bool): If True, don't actually change anything,
just make SQL, add it to ``self.executed_queries`` and return True.
Returns:
changed (bool): True if publication has been updated, otherwise False.
"""
if self.exists:
query_fragments = []
query_fragments.append(
"DROP PUBLICATION %s" %
pg_quote_identifier(self.name, 'publication'))
if cascade:
query_fragments.append("CASCADE")
return self.__exec_sql(' '.join(query_fragments),
check_mode=check_mode)
def __check_table(self, table):
"""Check table or SQL in transaction mode for check_mode.
Return True if it is OK.
Arguments:
table (str) - Table name that needs to be checked.
It can be SQL SELECT statement that was passed
instead of the table name.
"""
if 'SELECT ' in table.upper():
# In this case table is actually SQL SELECT statement.
# If SQL fails, it's handled by exec_sql():
exec_sql(self, table, add_to_executed=False)
# If exec_sql was passed, it means all is OK:
return True
exec_sql(self,
'SELECT 1 FROM %s' % pg_quote_identifier(table, 'table'),
add_to_executed=False)
# If SQL was executed successfully:
return True
def copy_to(self):
"""Implements COPY TO command behavior."""
self.src = self.module.params['src']
self.dst = self.module.params['copy_to']
if 'SELECT ' in self.src.upper():
# If src is SQL SELECT statement:
query_fragments = ['COPY (%s)' % self.src]
else:
# If src is a table:
query_fragments = [
'COPY %s' % pg_quote_identifier(self.src, 'table')
]
if self.module.params.get('columns'):
query_fragments.append('(%s)' %
','.join(self.module.params['columns']))
query_fragments.append('TO')
if self.module.params.get('program'):
query_fragments.append('PROGRAM')
query_fragments.append("'%s'" % self.dst)
if self.module.params.get('options'):
query_fragments.append(self.__transform_options())
# Note: check mode is implemented here:
if self.module.check_mode:
self.changed = self.__check_table(self.src)
if self.changed:
self.executed_queries.append(' '.join(query_fragments))
else:
if exec_sql(self, ' '.join(query_fragments), ddl=True):
self.changed = True
def set_owner(cursor, schema, owner):
query = "ALTER SCHEMA %s OWNER TO %s" % (pg_quote_identifier(
schema, 'schema'), pg_quote_identifier(owner, 'role'))
cursor.execute(query)
executed_queries.append(query)
return True
def manipulate_privs(self,
obj_type,
privs,
objs,
roles,
target_roles,
state,
grant_option,
schema_qualifier=None,
fail_on_role=True):
"""Manipulate database object privileges.
:param obj_type: Type of database object to grant/revoke
privileges for.
:param privs: Either a list of privileges to grant/revoke
or None if type is "group".
:param objs: List of database objects to grant/revoke
privileges for.
:param roles: Either a list of role names or "PUBLIC"
for the implicitly defined "PUBLIC" group
:param target_roles: List of role names to grant/revoke
default privileges as.
:param state: "present" to grant privileges, "absent" to revoke.
:param grant_option: Only for state "present": If True, set
grant/admin option. If False, revoke it.
If None, don't change grant option.
:param schema_qualifier: Some object types ("TABLE", "SEQUENCE",
"FUNCTION") must be qualified by schema.
Ignored for other Types.
"""
# get_status: function to get current status
if obj_type == 'table':
get_status = partial(self.get_table_acls, schema_qualifier)
elif obj_type == 'sequence':
get_status = partial(self.get_sequence_acls, schema_qualifier)
elif obj_type == 'function':
get_status = partial(self.get_function_acls, schema_qualifier)
elif obj_type == 'schema':
get_status = self.get_schema_acls
elif obj_type == 'language':
get_status = self.get_language_acls
elif obj_type == 'tablespace':
get_status = self.get_tablespace_acls
elif obj_type == 'database':
get_status = self.get_database_acls
elif obj_type == 'group':
get_status = self.get_group_memberships
elif obj_type == 'default_privs':
get_status = partial(self.get_default_privs, schema_qualifier)
elif obj_type == 'foreign_data_wrapper':
get_status = self.get_foreign_data_wrapper_acls
elif obj_type == 'foreign_server':
get_status = self.get_foreign_server_acls
elif obj_type == 'type':
get_status = partial(self.get_type_acls, schema_qualifier)
else:
raise Error('Unsupported database object type "%s".' % obj_type)
# Return False (nothing has changed) if there are no objs to work on.
if not objs:
return False
# obj_ids: quoted db object identifiers (sometimes schema-qualified)
if obj_type == 'function':
obj_ids = []
for obj in objs:
try:
f, args = obj.split('(', 1)
except Exception:
raise Error('Illegal function signature: "%s".' % obj)
obj_ids.append('"%s"."%s"(%s' % (schema_qualifier, f, args))
elif obj_type in ['table', 'sequence', 'type']:
obj_ids = ['"%s"."%s"' % (schema_qualifier, o) for o in objs]
else:
obj_ids = ['"%s"' % o for o in objs]
# set_what: SQL-fragment specifying what to set for the target roles:
# Either group membership or privileges on objects of a certain type
if obj_type == 'group':
set_what = ','.join('"%s"' % i for i in obj_ids)
elif obj_type == 'default_privs':
# We don't want privs to be quoted here
set_what = ','.join(privs)
else:
# function types are already quoted above
if obj_type != 'function':
obj_ids = [pg_quote_identifier(i, 'table') for i in obj_ids]
# Note: obj_type has been checked against a set of string literals
# and privs was escaped when it was parsed
# Note: Underscores are replaced with spaces to support multi-word obj_type
set_what = '%s ON %s %s' % (
','.join(privs), obj_type.replace('_', ' '), ','.join(obj_ids))
# for_whom: SQL-fragment specifying for whom to set the above
if roles == 'PUBLIC':
for_whom = 'PUBLIC'
else:
for_whom = []
for r in roles:
if not role_exists(self.module, self.cursor, r):
if fail_on_role:
self.module.fail_json(msg="Role '%s' does not exist" %
r.strip())
else:
self.module.warn("Role '%s' does not exist, pass it" %
r.strip())
else:
for_whom.append('"%s"' % r)
if not for_whom:
return False
for_whom = ','.join(for_whom)
# as_who:
as_who = None
if target_roles:
as_who = ','.join('"%s"' % r for r in target_roles)
status_before = get_status(objs)
query = QueryBuilder(state) \
.for_objtype(obj_type) \
.with_grant_option(grant_option) \
.for_whom(for_whom) \
.as_who(as_who) \
.for_schema(schema_qualifier) \
.set_what(set_what) \
.for_objs(objs) \
.build()
executed_queries.append(query)
self.cursor.execute(query)
status_after = get_status(objs)
return status_before != status_after
def test_invalid_quotes(identifier, id_type, msg):
with pytest.raises(SQLParseError) as ex:
pg_quote_identifier(identifier, id_type)
ex.match(msg)
def __set_tablespace_owner(self):
"""Set the tablespace owner."""
query = "ALTER TABLESPACE %s OWNER TO %s" % (pg_quote_identifier(
self.obj_name, 'database'), pg_quote_identifier(self.role, 'role'))
self.changed = exec_sql(self, query, ddl=True)
def set_tablespace(cursor, db, tablespace):
query = "ALTER DATABASE %s SET TABLESPACE %s" % (
pg_quote_identifier(db, 'database'),
pg_quote_identifier(tablespace, 'tablespace'))
cursor.execute(query)
return True
def set_conn_limit(cursor, db, conn_limit):
query = "ALTER DATABASE %s CONNECTION LIMIT %s" % (
pg_quote_identifier(db, 'database'),
conn_limit)
cursor.execute(query)
return True
def set_owner(cursor, db, owner):
query = 'ALTER DATABASE %s OWNER TO "%s"' % (
pg_quote_identifier(db, 'database'),
owner)
cursor.execute(query)
return True
def __set_mat_view_owner(self):
"""Set the materialized view owner."""
query = "ALTER MATERIALIZED VIEW %s OWNER TO %s" % (
pg_quote_identifier(self.obj_name, 'table'),
pg_quote_identifier(self.role, 'role'))
self.changed = exec_sql(self, query, ddl=True)
def test_valid_quotes(identifier, quoted_identifier):
assert pg_quote_identifier(identifier, 'table') == quoted_identifier
def __set_schema_owner(self):
"""Set the schema owner."""
query = "ALTER SCHEMA %s OWNER TO %s" % (pg_quote_identifier(
self.obj_name, 'schema'), pg_quote_identifier(self.role, 'role'))
self.changed = exec_sql(self, query, ddl=True)
