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)