def reassign(self, old_owners, fail_on_role): """Implements REASSIGN OWNED BY command. If success, set self.changed as True. Arguments: old_owners (list): The ownership of all the objects within the current database, and of all shared objects (databases, tablespaces), owned by these roles will be reassigned to self.role. fail_on_role (bool): If True, fail when a role from old_owners does not exist. Otherwise just warn and continue. """ roles = [] for r in old_owners: if self.check_role_exists(r, fail_on_role): roles.append(pg_quote_identifier(r, 'role')) # Roles do not exist, nothing to do, exit: if not roles: return False old_owners = ','.join(roles) query = ['REASSIGN OWNED BY'] query.append(old_owners) query.append('TO %s' % pg_quote_identifier(self.role, 'role')) query = ' '.join(query) self.changed = exec_sql(self, query, ddl=True)
def test_how_many_dots(identifier, id_type, quoted_identifier, msg): assert pg_quote_identifier(identifier, id_type) == quoted_identifier with pytest.raises(SQLParseError) as ex: pg_quote_identifier('%s.more' % identifier, id_type) ex.match(msg)
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 revoke_database_privileges(cursor, user, db, privs): # Note: priv escaped by parse_privs privs = ', '.join(privs) if user == "PUBLIC": query = 'REVOKE %s ON DATABASE %s FROM PUBLIC' % ( privs, pg_quote_identifier(db, 'database')) else: query = 'REVOKE %s ON DATABASE %s FROM "%s"' % ( privs, pg_quote_identifier(db, 'database'), user) executed_queries.append(query) cursor.execute(query)
def create_like(self, src_table, including='', tblspace='', unlogged=False, params='', owner=''): """ Create table like another table (with similar DDL). Arguments: src_table - source table. including - corresponds to optional INCLUDING expression in CREATE TABLE ... LIKE statement. params - storage params (passed by "WITH (...)" in SQL), comma separated. tblspace - tablespace. owner - table owner. unlogged - create unlogged table. """ changed = False name = pg_quote_identifier(self.name, 'table') query = "CREATE" if unlogged: query += " UNLOGGED TABLE %s" % name else: query += " TABLE %s" % name query += " (LIKE %s" % pg_quote_identifier(src_table, 'table') if including: including = including.split(',') for i in including: query += " INCLUDING %s" % i query += ')' if params: query += " WITH (%s)" % params if tblspace: query += " TABLESPACE %s" % pg_quote_identifier( tblspace, 'database') if exec_sql(self, query, ddl=True): changed = True if owner: changed = self.set_owner(owner) return changed
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 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 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 __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 __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 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 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 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 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 create(self, location): """Create tablespace. Return True if success, otherwise, return False. args: location (str) -- tablespace directory path in the FS """ query = ("CREATE TABLESPACE %s LOCATION '%s'" % (pg_quote_identifier(self.name, 'database'), location)) return exec_sql(self, query, ddl=True)
def __set_setting(self, setting): """Set tablespace setting. Return True if success, otherwise, return False. args: setting (str) -- string in format "setting_name = 'setting_value'" """ query = "ALTER TABLESPACE %s SET (%s)" % (pg_quote_identifier(self.name, 'database'), setting) return exec_sql(self, query, ddl=True)
def rename(self, newname): """Rename tablespace. Return True if success, otherwise, return False. args: newname (str) -- new name for the tablespace" """ query = "ALTER TABLESPACE %s RENAME TO %s" % (pg_quote_identifier(self.name, 'database'), newname) self.new_name = newname return exec_sql(self, query, ddl=True)
def set_owner(self, new_owner): """Set tablespace owner. Return True if success, otherwise, return False. args: new_owner (str) -- name of a new owner for the tablespace" """ if new_owner == self.owner: return False query = "ALTER TABLESPACE %s OWNER TO %s" % (pg_quote_identifier(self.name, 'database'), new_owner) return exec_sql(self, query, ddl=True)
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 __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 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 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 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