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 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_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'), pg_quote_identifier(user, 'role')) executed_queries.append(query) cursor.execute(query)
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'), pg_quote_identifier(user, 'role')) executed_queries.append(query) cursor.execute(query)
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 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 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 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 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 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 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 grant(self): for group in self.groups: self.granted[group] = [] for role in self.target_roles: # If role is in a group now, pass: if self.__check_membership(group, role): continue query = "GRANT %s TO %s" % ((pg_quote_identifier(group, 'role'), (pg_quote_identifier(role, 'role')))) self.changed = exec_sql(self, query, ddl=True) if self.changed: self.granted[group].append(role) return self.changed
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 revoke(self): for group in self.groups: self.revoked[group] = [] for role in self.target_roles: # If role is not in a group now, pass: if not self.__check_membership(group, role): continue query = "REVOKE %s FROM %s" % ((pg_quote_identifier(group, 'role'), (pg_quote_identifier(role, 'role')))) self.changed = exec_sql(self, query, ddl=True) if self.changed: self.revoked[group].append(role) return self.changed
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 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 user_delete(cursor, user): """Try to remove a user. Returns True if successful otherwise False""" cursor.execute("SAVEPOINT ansible_pgsql_user_delete") try: query = "DROP USER %s" % pg_quote_identifier(user, 'role') executed_queries.append(query) cursor.execute(query) except Exception: cursor.execute("ROLLBACK TO SAVEPOINT ansible_pgsql_user_delete") cursor.execute("RELEASE SAVEPOINT ansible_pgsql_user_delete") return False cursor.execute("RELEASE SAVEPOINT ansible_pgsql_user_delete") return True
def user_add(cursor, user, password, role_attr_flags, encrypted, expires, conn_limit): """Create a new database user (role).""" # Note: role_attr_flags escaped by parse_role_attrs and encrypted is a # literal query_password_data = dict(password=password, expires=expires) query = ['CREATE USER %(user)s' % {"user": pg_quote_identifier(user, 'role')}] if password is not None and password != '': query.append("WITH %(crypt)s" % {"crypt": encrypted}) query.append("PASSWORD %(password)s") if expires is not None: query.append("VALID UNTIL %(expires)s") if conn_limit is not None: query.append("CONNECTION LIMIT %(conn_limit)s" % {"conn_limit": conn_limit}) query.append(role_attr_flags) query = ' '.join(query) executed_queries.append(query) cursor.execute(query, query_password_data) 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 main(): argument_spec = pgutils.jctanner.database_postgresql.postgres_common_argument_spec( ) argument_spec.update( db=dict(type='str', required=True, aliases=['name']), owner=dict(type='str', default=''), template=dict(type='str', default=''), encoding=dict(type='str', default=''), lc_collate=dict(type='str', default=''), lc_ctype=dict(type='str', default=''), state=dict(type='str', default='present', choices=['absent', 'dump', 'present', 'restore']), target=dict(type='path', default=''), target_opts=dict(type='str', default=''), maintenance_db=dict(type='str', default="jctanner.database_postgresql.postgres"), session_role=dict(type='str'), conn_limit=dict(type='str', default=''), tablespace=dict(type='path', default=''), ) module = AnsibleModule(argument_spec=argument_spec, supports_check_mode=True) db = module.params["db"] owner = module.params["owner"] template = module.params["template"] encoding = module.params["encoding"] lc_collate = module.params["lc_collate"] lc_ctype = module.params["lc_ctype"] target = module.params["target"] target_opts = module.params["target_opts"] state = module.params["state"] changed = False maintenance_db = module.params['maintenance_db'] session_role = module.params["session_role"] conn_limit = module.params['conn_limit'] tablespace = module.params['tablespace'] raw_connection = state in ("dump", "restore") if not raw_connection: pgutils.ensure_required_libs(module) # To use defaults values, keyword arguments must be absent, so # check which values are empty and don't include in the **kw # dictionary params_map = { "login_host": "host", "login_user": "******", "login_password": "******", "port": "port", "ssl_mode": "sslmode", "ca_cert": "sslrootcert" } kw = dict((params_map[k], v) for (k, v) in iteritems(module.params) if k in params_map and v != '' and v is not None) # If a login_unix_socket is specified, incorporate it here. is_localhost = "host" not in kw or kw["host"] == "" or kw[ "host"] == "localhost" if is_localhost and module.params["login_unix_socket"] != "": kw["host"] = module.params["login_unix_socket"] if target == "": target = "{0}/{1}.sql".format(os.getcwd(), db) target = os.path.expanduser(target) if not raw_connection: try: db_connection = psycopg2.connect(database=maintenance_db, **kw) # Enable autocommit so we can create databases if psycopg2.__version__ >= '2.4.2': db_connection.autocommit = True else: db_connection.set_isolation_level( psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT) cursor = db_connection.cursor( cursor_factory=psycopg2.extras.DictCursor) except pgutils.LibraryError as e: module.fail_json(msg="unable to connect to database: {0}".format( to_native(e)), exception=traceback.format_exc()) except TypeError as e: if 'sslrootcert' in e.args[0]: module.fail_json( msg= 'Postgresql server must be at least version 8.4 to support sslrootcert. Exception: {0}' .format(to_native(e)), exception=traceback.format_exc()) module.fail_json(msg="unable to connect to database: %s" % to_native(e), exception=traceback.format_exc()) except Exception as e: module.fail_json(msg="unable to connect to database: %s" % to_native(e), exception=traceback.format_exc()) if session_role: try: cursor.execute('SET ROLE %s' % pg_quote_identifier(session_role, 'role')) except Exception as e: module.fail_json(msg="Could not switch role: %s" % to_native(e), exception=traceback.format_exc()) try: if module.check_mode: if state == "absent": changed = db_exists(cursor, db) elif state == "present": changed = not db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace) module.exit_json(changed=changed, db=db) if state == "absent": try: changed = db_delete(cursor, db) except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) elif state == "present": try: changed = db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace) except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) elif state in ("dump", "restore"): method = state == "dump" and db_dump or db_restore try: rc, stdout, stderr, cmd = method(module, target, target_opts, db, **kw) if rc != 0: module.fail_json(msg=stderr, stdout=stdout, rc=rc, cmd=cmd) else: module.exit_json(changed=True, msg=stdout, stderr=stderr, rc=rc, cmd=cmd) except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) except NotSupportedError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) except SystemExit: # Avoid catching this on Python 2.4 raise except Exception as e: module.fail_json(msg="Database query failed: %s" % to_native(e), exception=traceback.format_exc()) module.exit_json(changed=changed, db=db)
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 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' % pg_quote_identifier(owner, 'role')) 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 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'), pg_quote_identifier(owner, 'role')) 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)