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_tablespace(cursor, db, tablespace): query = "ALTER DATABASE %s SET TABLESPACE %s" % ( pg_quote_identifier(db, 'database'), pg_quote_identifier(tablespace, 'tablespace')) executed_commands.append(query) cursor.execute(query) return True
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 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) executed_commands.append(cursor.mogrify(query, params)) 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 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 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_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 set_conn_limit(cursor, db, conn_limit): query = "ALTER DATABASE %s CONNECTION LIMIT %s" % ( pg_quote_identifier(db, 'database'), conn_limit) executed_commands.append(query) cursor.execute(query) return True
def set_owner(cursor, db, owner): query = 'ALTER DATABASE %s OWNER TO "%s"' % ( pg_quote_identifier(db, 'database'), owner) executed_commands.append(query) cursor.execute(query) 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), return_bool=True): self.changed = True
def db_delete(cursor, db): if db_exists(cursor, db): query = "DROP DATABASE %s" % pg_quote_identifier(db, 'database') executed_commands.append(query) cursor.execute(query) return True else: return False
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 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 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), return_bool=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 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 drop(self): """Drop tablespace. Return True if success, otherwise, return False. """ return exec_sql(self, "DROP TABLESPACE %s" % pg_quote_identifier(self.name, 'database'), 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), return_bool=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), return_bool=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 __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 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 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 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)