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 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')) 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) 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 test_how_many_dots(self): eq_(pg_quote_identifier('role', 'role'), '"role"') assert_raises_regexp(SQLParseError, "PostgreSQL does not support role with more than 1 dots", pg_quote_identifier, *('role.more', 'role')) eq_(pg_quote_identifier('db', 'database'), '"db"') assert_raises_regexp(SQLParseError, "PostgreSQL does not support database with more than 1 dots", pg_quote_identifier, *('db.more', 'database')) eq_(pg_quote_identifier('db.schema', 'schema'), '"db"."schema"') assert_raises_regexp(SQLParseError, "PostgreSQL does not support schema with more than 2 dots", pg_quote_identifier, *('db.schema.more', 'schema')) eq_(pg_quote_identifier('db.schema.table', 'table'), '"db"."schema"."table"') assert_raises_regexp(SQLParseError, "PostgreSQL does not support table with more than 3 dots", pg_quote_identifier, *('db.schema.table.more', 'table')) eq_(pg_quote_identifier('db.schema.table.column', 'column'), '"db"."schema"."table"."column"') assert_raises_regexp(SQLParseError, "PostgreSQL does not support column with more than 4 dots", pg_quote_identifier, *('db.schema.table.column.more', 'column'))
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 schema_delete(cursor, schema): if schema_exists(cursor, schema): query = "DROP SCHEMA %s" % pg_quote_identifier(schema, 'schema') cursor.execute(query) return True else: return False
def user_delete(cursor, user): """Try to remove a user. Returns True if successful otherwise False""" cursor.execute("SAVEPOINT ansible_pgsql_user_delete") try: cursor.execute("DROP USER %s" % pg_quote_identifier(user, 'role')) except: 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 db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype): params = dict(enc=encoding, collate=lc_collate, ctype=lc_ctype) 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') 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'] ) elif owner and owner != db_info['owner']: return set_owner(cursor, db, owner) else: return False
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: 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) cursor.execute(query, query_password_data) 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 check_valid_quotes(self, identifier, quoted_identifier): eq_(pg_quote_identifier(identifier, 'table'), quoted_identifier)
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 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 main(): argument_spec = pgutils.postgres_common_argument_spec() argument_spec.update( dict( user=dict(required=True, aliases=['name']), password=dict(default=None, no_log=True), state=dict(default="present", choices=["absent", "present"]), priv=dict(default=None), db=dict(default=''), fail_on_user=dict(type='bool', default='yes'), role_attr_flags=dict(default=''), encrypted=dict(type='bool', default='yes'), no_password_changes=dict(type='bool', default='no'), expires=dict(default=None), ssl_mode=dict(default='prefer', choices=[ 'disable', 'allow', 'prefer', 'require', 'verify-ca', 'verify-full' ]), ssl_rootcert=dict(default=None), conn_limit=dict(type='int', default=None), session_role=dict(), )) module = AnsibleModule(argument_spec=argument_spec, supports_check_mode=True) user = module.params["user"] password = module.params["password"] state = module.params["state"] fail_on_user = module.params["fail_on_user"] db = module.params["db"] session_role = module.params["session_role"] if db == '' and module.params["priv"] is not None: module.fail_json(msg="privileges require a database to be specified") privs = parse_privs(module.params["priv"], db) no_password_changes = module.params["no_password_changes"] if module.params["encrypted"]: encrypted = "ENCRYPTED" else: encrypted = "UNENCRYPTED" expires = module.params["expires"] sslrootcert = module.params["ssl_rootcert"] conn_limit = module.params["conn_limit"] if not postgresqldb_found: module.fail_json(msg="the python psycopg2 module is required") # 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", "db": "database", "ssl_mode": "sslmode", "ssl_rootcert": "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 psycopg2.__version__ < '2.4.3' and sslrootcert is not None: module.fail_json( msg= 'psycopg2 must be at least 2.4.3 in order to user the ssl_rootcert parameter' ) try: db_connection = psycopg2.connect(**kw) cursor = db_connection.cursor( cursor_factory=psycopg2.extras.DictCursor) 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' ) 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: role_attr_flags = parse_role_attrs(cursor, module.params["role_attr_flags"]) except InvalidFlagsError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) kw = dict(user=user) changed = False user_removed = False if state == "present": if user_exists(cursor, user): try: changed = user_alter(db_connection, module, user, password, role_attr_flags, encrypted, expires, no_password_changes, conn_limit) except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) else: try: changed = user_add(cursor, user, password, role_attr_flags, encrypted, expires, conn_limit) except psycopg2.ProgrammingError as e: module.fail_json( msg="Unable to add user with given requirement " "due to : %s" % to_native(e), exception=traceback.format_exc()) except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) try: changed = grant_privileges(cursor, user, privs) or changed except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) else: if user_exists(cursor, user): if module.check_mode: changed = True kw['user_removed'] = True else: try: changed = revoke_privileges(cursor, user, privs) user_removed = user_delete(cursor, user) except SQLParseError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) changed = changed or user_removed if fail_on_user and not user_removed: msg = "unable to remove user" module.fail_json(msg=msg) kw['user_removed'] = user_removed if changed: if module.check_mode: db_connection.rollback() else: db_connection.commit() kw['changed'] = changed module.exit_json(**kw)
def main(): module = AnsibleModule(argument_spec=dict( database=dict(required=True, aliases=['db']), state=dict(default='present', choices=['present', 'absent']), privs=dict(required=False, aliases=['priv']), type=dict(default='table', choices=[ 'table', 'sequence', 'function', 'database', 'schema', 'language', 'tablespace', 'group', 'default_privs', 'foreign_data_wrapper', 'foreign_server' ]), objs=dict(required=False, aliases=['obj']), schema=dict(required=False), roles=dict(required=True, aliases=['role']), session_role=dict(required=False), target_roles=dict(required=False), grant_option=dict(required=False, type='bool', aliases=['admin_option']), host=dict(default='', aliases=['login_host']), port=dict(type='int', default=5432), unix_socket=dict(default='', aliases=['login_unix_socket']), login=dict(default='postgres', aliases=['login_user']), password=dict(default='', aliases=['login_password'], no_log=True), ssl_mode=dict(default="prefer", choices=[ 'disable', 'allow', 'prefer', 'require', 'verify-ca', 'verify-full' ]), ca_cert=dict(default=None, aliases=['ssl_rootcert']), fail_on_role=dict(type='bool', default=True), ), supports_check_mode=True) fail_on_role = module.params['fail_on_role'] # Create type object as namespace for module params p = type('Params', (), module.params) # param "schema": default, allowed depends on param "type" if p.type in ['table', 'sequence', 'function', 'default_privs']: p.schema = p.schema or 'public' elif p.schema: module.fail_json(msg='Argument "schema" is not allowed ' 'for type "%s".' % p.type) # param "objs": default, required depends on param "type" if p.type == 'database': p.objs = p.objs or p.database elif not p.objs: module.fail_json(msg='Argument "objs" is required ' 'for type "%s".' % p.type) # param "privs": allowed, required depends on param "type" if p.type == 'group': if p.privs: module.fail_json(msg='Argument "privs" is not allowed ' 'for type "group".') elif not p.privs: module.fail_json(msg='Argument "privs" is required ' 'for type "%s".' % p.type) # Connect to Database if not psycopg2: module.fail_json(msg=missing_required_lib('psycopg2'), exception=PSYCOPG2_IMP_ERR) try: conn = Connection(p, module) except psycopg2.Error as e: module.fail_json(msg='Could not connect to database: %s' % 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' ) module.fail_json(msg="unable to connect to database: %s" % to_native(e), exception=traceback.format_exc()) except ValueError as e: # We raise this when the psycopg library is too old module.fail_json(msg=to_native(e)) if p.session_role: try: conn.cursor.execute('SET ROLE %s' % pg_quote_identifier(p.session_role, 'role')) except Exception as e: module.fail_json(msg="Could not switch to role %s: %s" % (p.session_role, to_native(e)), exception=traceback.format_exc()) try: # privs if p.privs: privs = frozenset(pr.upper() for pr in p.privs.split(',')) if not privs.issubset(VALID_PRIVS): module.fail_json(msg='Invalid privileges specified: %s' % privs.difference(VALID_PRIVS)) else: privs = None # objs: if p.type == 'table' and p.objs == 'ALL_IN_SCHEMA': objs = conn.get_all_tables_in_schema(p.schema) elif p.type == 'sequence' and p.objs == 'ALL_IN_SCHEMA': objs = conn.get_all_sequences_in_schema(p.schema) elif p.type == 'function' and p.objs == 'ALL_IN_SCHEMA': objs = conn.get_all_functions_in_schema(p.schema) elif p.type == 'default_privs': if p.objs == 'ALL_DEFAULT': objs = frozenset(VALID_DEFAULT_OBJS.keys()) else: objs = frozenset(obj.upper() for obj in p.objs.split(',')) if not objs.issubset(VALID_DEFAULT_OBJS): module.fail_json( msg='Invalid Object set specified: %s' % objs.difference(VALID_DEFAULT_OBJS.keys())) # Again, do we have valid privs specified for object type: valid_objects_for_priv = frozenset( obj for obj in objs if privs.issubset(VALID_DEFAULT_OBJS[obj])) if not valid_objects_for_priv == objs: module.fail_json( msg= 'Invalid priv specified. Valid object for priv: {0}. Objects: {1}' .format(valid_objects_for_priv, objs)) else: objs = p.objs.split(',') # function signatures are encoded using ':' to separate args if p.type == 'function': objs = [obj.replace(':', ',') for obj in objs] # roles if p.roles == 'PUBLIC': roles = 'PUBLIC' else: roles = p.roles.split(',') if len(roles) == 1 and not role_exists(module, conn.cursor, roles[0]): module.exit_json(changed=False) if fail_on_role: module.fail_json(msg="Role '%s' does not exist" % roles[0].strip()) else: module.warn("Role '%s' does not exist, nothing to do" % roles[0].strip()) # check if target_roles is set with type: default_privs if p.target_roles and not p.type == 'default_privs': module.warn( '"target_roles" will be ignored ' 'Argument "type: default_privs" is required for usage of "target_roles".' ) # target roles if p.target_roles: target_roles = p.target_roles.split(',') else: target_roles = None changed = conn.manipulate_privs( obj_type=p.type, privs=privs, objs=objs, roles=roles, target_roles=target_roles, state=p.state, grant_option=p.grant_option, schema_qualifier=p.schema, fail_on_role=fail_on_role, ) except Error as e: conn.rollback() module.fail_json(msg=e.message, exception=traceback.format_exc()) except psycopg2.Error as e: conn.rollback() module.fail_json(msg=to_native(e.message)) if module.check_mode: conn.rollback() else: conn.commit() module.exit_json(changed=changed)
def __set_tablespace_owner(self): query = "ALTER TABLESPACE %s OWNER TO %s" % (pg_quote_identifier( self.obj_name, 'database'), pg_quote_identifier(self.role, 'role')) self.changed = self.__exec_sql(query, ddl=True)
def manipulate_privs(self, obj_type, privs, objs, roles, state, grant_option, schema_qualifier=None): """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 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 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: raise Error('Illegal function signature: "%s".' % obj) obj_ids.append('"%s"."%s"(%s' % (schema_qualifier, f, args)) elif obj_type in ['table', 'sequence']: 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(pg_quote_identifier(i, 'role') for i in obj_ids) 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 set_what = '%s ON %s %s' % (','.join(privs), obj_type, ','.join(obj_ids)) # for_whom: SQL-fragment specifying for whom to set the above if roles == 'PUBLIC': for_whom = 'PUBLIC' else: for_whom = ','.join(pg_quote_identifier(r, 'role') for r in roles) status_before = get_status(objs) if state == 'present': if grant_option: if obj_type == 'group': query = 'GRANT %s TO %s WITH ADMIN OPTION' else: query = 'GRANT %s TO %s WITH GRANT OPTION' else: query = 'GRANT %s TO %s' self.cursor.execute(query % (set_what, for_whom)) # Only revoke GRANT/ADMIN OPTION if grant_option actually is False. if grant_option is False: if obj_type == 'group': query = 'REVOKE ADMIN OPTION FOR %s FROM %s' else: query = 'REVOKE GRANT OPTION FOR %s FROM %s' self.cursor.execute(query % (set_what, for_whom)) else: query = 'REVOKE %s FROM %s' self.cursor.execute(query % (set_what, for_whom)) status_after = get_status(objs) return status_before != status_after
def set_stor_params(self, params): query = "ALTER TABLE %s SET (%s)" % (pg_quote_identifier( self.name, 'table'), params) return exec_sql(self, query, ddl=True)
def set_tblspace(self, tblspace): query = "ALTER TABLE %s SET TABLESPACE %s" % (pg_quote_identifier( self.name, 'table'), pg_quote_identifier(tblspace, 'database')) return exec_sql(self, query, ddl=True)
def set_owner(self, username): query = "ALTER TABLE %s OWNER TO %s" % (pg_quote_identifier( self.name, 'table'), pg_quote_identifier(username, 'role')) return exec_sql(self, query, ddl=True)
def rename(self, newname): query = "ALTER TABLE %s RENAME TO %s" % (pg_quote_identifier( self.name, 'table'), pg_quote_identifier(newname, 'table')) return exec_sql(self, query, ddl=True)
def truncate(self): query = "TRUNCATE TABLE %s" % pg_quote_identifier(self.name, 'table') return exec_sql(self, query, ddl=True)
def __set_mat_view_owner(self): query = "ALTER MATERIALIZED VIEW %s OWNER TO %s" % ( pg_quote_identifier(self.obj_name, 'table'), pg_quote_identifier(self.role, 'role')) self.changed = self.__exec_sql(query, ddl=True)
def main(): module = AnsibleModule( argument_spec=dict( login_user=dict(default="postgres"), login_password=dict(default="", no_log=True), login_host=dict(default=""), login_unix_socket=dict(default=""), port=dict(default="5432"), db=dict(required=True), ext=dict(required=True, aliases=['name']), schema=dict(default=""), state=dict(default="present", choices=["absent", "present"]), cascade=dict(type='bool', default=False), ssl_mode=dict(default='prefer', choices=[ 'disable', 'allow', 'prefer', 'require', 'verify-ca', 'verify-full']), ssl_rootcert=dict(default=None), session_role=dict(), ), supports_check_mode=True ) if not postgresqldb_found: module.fail_json(msg=missing_required_lib('psycopg2'), exception=PSYCOPG2_IMP_ERR) db = module.params["db"] ext = module.params["ext"] schema = module.params["schema"] state = module.params["state"] cascade = module.params["cascade"] sslrootcert = module.params["ssl_rootcert"] session_role = module.params["session_role"] changed = False # 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", "db": "database", "ssl_mode": "sslmode", "ssl_rootcert": "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 psycopg2.__version__ < '2.4.3' and sslrootcert is not None: module.fail_json(msg='psycopg2 must be at least 2.4.3 in order to user the ssl_rootcert parameter') try: db_connection = psycopg2.connect(**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 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') 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 == "present": changed = not ext_exists(cursor, ext) elif state == "absent": changed = ext_exists(cursor, ext) else: if state == "absent": changed = ext_delete(cursor, ext, cascade) elif state == "present": changed = ext_create(cursor, ext, schema, cascade) except NotSupportedError as e: module.fail_json(msg=to_native(e), exception=traceback.format_exc()) 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, ext=ext)
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) return True
def manipulate_privs(self, obj_type, privs, objs, roles, state, grant_option, schema_qualifier=None): """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 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 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: raise Error('Illegal function signature: "%s".' % obj) obj_ids.append('"%s"."%s"(%s' % (schema_qualifier, f, args)) elif obj_type in ['table', 'sequence']: 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( pg_quote_identifier(i, 'role') for i in obj_ids) 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 set_what = '%s ON %s %s' % (','.join(privs), obj_type, ','.join(obj_ids)) # for_whom: SQL-fragment specifying for whom to set the above if roles == 'PUBLIC': for_whom = 'PUBLIC' else: for_whom = ','.join(pg_quote_identifier(r, 'role') for r in roles) status_before = get_status(objs) if state == 'present': if grant_option: if obj_type == 'group': query = 'GRANT %s TO %s WITH ADMIN OPTION' else: query = 'GRANT %s TO %s WITH GRANT OPTION' else: query = 'GRANT %s TO %s' self.cursor.execute(query % (set_what, for_whom)) # Only revoke GRANT/ADMIN OPTION if grant_option actually is False. if grant_option is False: if obj_type == 'group': query = 'REVOKE ADMIN OPTION FOR %s FROM %s' else: query = 'REVOKE GRANT OPTION FOR %s FROM %s' self.cursor.execute(query % (set_what, for_whom)) else: query = 'REVOKE %s FROM %s' self.cursor.execute(query % (set_what, for_whom)) status_after = get_status(objs) return status_before != status_after
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 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']: 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( pg_quote_identifier(i, 'role') 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(pg_quote_identifier(r, 'role')) if not for_whom: return False for_whom = ','.join(for_whom) # as_who: as_who = None if target_roles: as_who = ','.join( pg_quote_identifier(r, 'role') 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() self.cursor.execute(query) status_after = get_status(objs) return status_before != status_after
def drop(self): query = "DROP TABLE %s" % pg_quote_identifier(self.name, 'table') self.executed_queries.append(query) return self.__exec_sql(query, ddl=True)
def user_alter(db_connection, module, user, password, role_attr_flags, encrypted, expires, no_password_changes, conn_limit): """Change user password and/or attributes. Return True if changed, False otherwise.""" changed = False cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor) # Note: role_attr_flags escaped by parse_role_attrs and encrypted is a # literal if user == 'PUBLIC': if password is not None: module.fail_json(msg="cannot change the password for PUBLIC user") elif role_attr_flags != '': module.fail_json( msg="cannot change the role_attr_flags for PUBLIC user") else: return False # Handle passwords. if not no_password_changes and (password is not None or role_attr_flags != '' or expires is not None or conn_limit is not None): # Select password and all flag-like columns in order to verify changes. try: select = "SELECT * FROM pg_authid where rolname=%(user)s" cursor.execute(select, {"user": user}) # Grab current role attributes. current_role_attrs = cursor.fetchone() except psycopg2.ProgrammingError: current_role_attrs = None db_connection.rollback() pwchanging = user_should_we_change_password(current_role_attrs, user, password, encrypted) if current_role_attrs is None: try: # AWS RDS instances does not allow user to access pg_authid # so try to get current_role_attrs from pg_roles tables select = "SELECT * FROM pg_roles where rolname=%(user)s" cursor.execute(select, {"user": user}) # Grab current role attributes from pg_roles current_role_attrs = cursor.fetchone() except psycopg2.ProgrammingError as e: db_connection.rollback() module.fail_json( msg="Failed to get role details for current user %s: %s" % (user, e)) role_attr_flags_changing = False if role_attr_flags: role_attr_flags_dict = {} for r in role_attr_flags.split(' '): if r.startswith('NO'): role_attr_flags_dict[r.replace('NO', '', 1)] = False else: role_attr_flags_dict[r] = True for role_attr_name, role_attr_value in role_attr_flags_dict.items( ): if current_role_attrs[PRIV_TO_AUTHID_COLUMN[ role_attr_name]] != role_attr_value: role_attr_flags_changing = True if expires is not None: cursor.execute("SELECT %s::timestamptz;", (expires, )) expires_with_tz = cursor.fetchone()[0] expires_changing = expires_with_tz != current_role_attrs.get( 'rolvaliduntil') else: expires_changing = False conn_limit_changing = ( conn_limit is not None and conn_limit != current_role_attrs['rolconnlimit']) if not pwchanging and not role_attr_flags_changing and not expires_changing and not conn_limit_changing: return False alter = [ 'ALTER USER %(user)s' % { "user": pg_quote_identifier(user, 'role') } ] if pwchanging: if password != '': alter.append("WITH %(crypt)s" % {"crypt": encrypted}) alter.append("PASSWORD %(password)s") else: alter.append("WITH PASSWORD NULL") alter.append(role_attr_flags) elif role_attr_flags: alter.append('WITH %s' % role_attr_flags) if expires is not None: alter.append("VALID UNTIL %(expires)s") if conn_limit is not None: alter.append("CONNECTION LIMIT %(conn_limit)s" % {"conn_limit": conn_limit}) query_password_data = dict(password=password, expires=expires) try: cursor.execute(' '.join(alter), query_password_data) changed = True except psycopg2.InternalError as e: if e.pgcode == '25006': # Handle errors due to read-only transactions indicated by pgcode 25006 # ERROR: cannot execute ALTER ROLE in a read-only transaction changed = False module.fail_json(msg=e.pgerror, exception=traceback.format_exc()) return changed else: raise psycopg2.InternalError(e) except psycopg2.NotSupportedError as e: module.fail_json(msg=e.pgerror, exception=traceback.format_exc()) elif no_password_changes and role_attr_flags != '': # Grab role information from pg_roles instead of pg_authid select = "SELECT * FROM pg_roles where rolname=%(user)s" cursor.execute(select, {"user": user}) # Grab current role attributes. current_role_attrs = cursor.fetchone() role_attr_flags_changing = False if role_attr_flags: role_attr_flags_dict = {} for r in role_attr_flags.split(' '): if r.startswith('NO'): role_attr_flags_dict[r.replace('NO', '', 1)] = False else: role_attr_flags_dict[r] = True for role_attr_name, role_attr_value in role_attr_flags_dict.items( ): if current_role_attrs[PRIV_TO_AUTHID_COLUMN[ role_attr_name]] != role_attr_value: role_attr_flags_changing = True if not role_attr_flags_changing: return False alter = [ 'ALTER USER %(user)s' % { "user": pg_quote_identifier(user, 'role') } ] if role_attr_flags: alter.append('WITH %s' % role_attr_flags) try: cursor.execute(' '.join(alter)) except psycopg2.InternalError as e: if e.pgcode == '25006': # Handle errors due to read-only transactions indicated by pgcode 25006 # ERROR: cannot execute ALTER ROLE in a read-only transaction changed = False module.fail_json(msg=e.pgerror, exception=traceback.format_exc()) return changed else: raise psycopg2.InternalError(e) # Grab new role attributes. cursor.execute(select, {"user": user}) new_role_attrs = cursor.fetchone() # Detect any differences between current_ and new_role_attrs. changed = current_role_attrs != new_role_attrs return changed
def create(self, columns='', params='', tblspace='', unlogged=False, owner=''): """ Create table. If table exists, check passed args (params, tblspace, owner) and, if they're different from current, change them. Arguments: params - storage params (passed by "WITH (...)" in SQL), comma separated. tblspace - tablespace. owner - table owner. unlogged - create unlogged table. columns - column string (comma separated). """ name = pg_quote_identifier(self.name, 'table') changed = False if self.exists: if tblspace == 'pg_default' and self.info['tblspace'] is None: pass # Because they have the same meaning elif tblspace and self.info['tblspace'] != tblspace: self.set_tblspace(tblspace) changed = True if owner and self.info['owner'] != owner: self.set_owner(owner) changed = True if params: param_list = [p.strip(' ') for p in params.split(',')] new_param = False for p in param_list: if p not in self.info['storage_params']: new_param = True if new_param: self.set_stor_params(params) changed = True if changed: return True return False query = "CREATE" if unlogged: query += " UNLOGGED TABLE %s" % name else: query += " TABLE %s" % name if columns: query += " (%s)" % columns else: query += " ()" if params: query += " WITH (%s)" % params if tblspace: query += " TABLESPACE %s" % pg_quote_identifier( tblspace, 'database') if self.__exec_sql(query, ddl=True): self.executed_queries.append(query) changed = True if owner: changed = self.set_owner(owner) return changed
def main(): argument_spec = postgres_common_argument_spec() argument_spec.update( db=dict(type="str", required=True, aliases=["login_db"]), lang=dict(type="str", required=True, aliases=["name"]), state=dict(type="str", default="present", choices=["absent", "present"]), trust=dict(type="bool", default="no"), force_trust=dict(type="bool", default="no"), cascade=dict(type="bool", default="no"), fail_on_drop=dict(type="bool", default="yes"), session_role=dict(type="str"), ) module = AnsibleModule( argument_spec=argument_spec, supports_check_mode=True, ) db = module.params["db"] lang = module.params["lang"] state = module.params["state"] trust = module.params["trust"] force_trust = module.params["force_trust"] cascade = module.params["cascade"] fail_on_drop = module.params["fail_on_drop"] sslrootcert = module.params["ca_cert"] session_role = module.params["session_role"] if not HAS_PSYCOPG2: module.fail_json(msg=missing_required_lib('psycopg2'), exception=PSYCOPG2_IMP_ERR) # 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", "db": "database", "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 psycopg2.__version__ < '2.4.3' and sslrootcert is not None: module.fail_json( msg= 'psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter' ) db_connection = connect_to_db(module, kw, autocommit=False) cursor = db_connection.cursor() 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()) changed = False kw = {'db': db, 'lang': lang, 'trust': trust} if state == "present": if lang_exists(cursor, lang): lang_trusted = lang_istrusted(cursor, lang) if (lang_trusted and not trust) or (not lang_trusted and trust): if module.check_mode: changed = True else: changed = lang_altertrust(cursor, lang, trust) else: if module.check_mode: changed = True else: changed = lang_add(cursor, lang, trust) if force_trust: changed = lang_altertrust(cursor, lang, trust) else: if lang_exists(cursor, lang): if module.check_mode: changed = True kw['lang_dropped'] = True else: changed = lang_drop(cursor, lang, cascade) if fail_on_drop and not changed: msg = "unable to drop language, use cascade to delete dependencies or fail_on_drop=no to ignore" module.fail_json(msg=msg) kw['lang_dropped'] = changed if changed: if module.check_mode: db_connection.rollback() else: db_connection.commit() kw['changed'] = changed kw['queries'] = executed_queries module.exit_json(**kw)
def truncate(self): query = "TRUNCATE TABLE %s" % pg_quote_identifier(self.name, 'table') self.executed_queries.append(query) return self.__exec_sql(query, ddl=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"' % 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 rename(self, newname): query = "ALTER TABLE %s RENAME TO %s" % (pg_quote_identifier( self.name, 'table'), pg_quote_identifier(newname, 'table')) self.executed_queries.append(query) return self.__exec_sql(query, ddl=True)
def grant_table_privileges(cursor, user, table, privs): # Note: priv escaped by parse_privs privs = ', '.join(privs) query = 'GRANT %s ON TABLE %s TO %s' % ( privs, pg_quote_identifier(table, 'table'), pg_quote_identifier(user, 'role')) cursor.execute(query)
def user_alter(db_connection, module, user, password, role_attr_flags, encrypted, expires, no_password_changes, conn_limit): """Change user password and/or attributes. Return True if changed, False otherwise.""" changed = False cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor) # Note: role_attr_flags escaped by parse_role_attrs and encrypted is a # literal if user == 'PUBLIC': if password is not None: module.fail_json(msg="cannot change the password for PUBLIC user") elif role_attr_flags != '': module.fail_json(msg="cannot change the role_attr_flags for PUBLIC user") else: return False # Handle passwords. if not no_password_changes and (password is not None or role_attr_flags != '' or expires is not None or conn_limit is not None): # Select password and all flag-like columns in order to verify changes. try: select = "SELECT * FROM pg_authid where rolname=%(user)s" cursor.execute(select, {"user": user}) # Grab current role attributes. current_role_attrs = cursor.fetchone() except psycopg2.ProgrammingError: current_role_attrs = None db_connection.rollback() pwchanging = user_should_we_change_password(current_role_attrs, user, password, encrypted) role_attr_flags_changing = False if role_attr_flags: role_attr_flags_dict = {} for r in role_attr_flags.split(' '): if r.startswith('NO'): role_attr_flags_dict[r.replace('NO', '', 1)] = False else: role_attr_flags_dict[r] = True for role_attr_name, role_attr_value in role_attr_flags_dict.items(): if current_role_attrs[PRIV_TO_AUTHID_COLUMN[role_attr_name]] != role_attr_value: role_attr_flags_changing = True if expires is not None: cursor.execute("SELECT %s::timestamptz;", (expires,)) expires_with_tz = cursor.fetchone()[0] expires_changing = expires_with_tz != current_role_attrs.get('rolvaliduntil') else: expires_changing = False conn_limit_changing = (conn_limit is not None and conn_limit != current_role_attrs['rolconnlimit']) if not pwchanging and not role_attr_flags_changing and not expires_changing and not conn_limit_changing: return False alter = ['ALTER USER %(user)s' % {"user": pg_quote_identifier(user, 'role')}] if pwchanging: alter.append("WITH %(crypt)s" % {"crypt": encrypted}) alter.append("PASSWORD %(password)s") alter.append(role_attr_flags) elif role_attr_flags: alter.append('WITH %s' % role_attr_flags) if expires is not None: alter.append("VALID UNTIL %(expires)s") if conn_limit is not None: alter.append("CONNECTION LIMIT %(conn_limit)s" % {"conn_limit": conn_limit}) query_password_data = dict(password=password, expires=expires) try: cursor.execute(' '.join(alter), query_password_data) changed = True except psycopg2.InternalError as e: if e.pgcode == '25006': # Handle errors due to read-only transactions indicated by pgcode 25006 # ERROR: cannot execute ALTER ROLE in a read-only transaction changed = False module.fail_json(msg=e.pgerror, exception=traceback.format_exc()) return changed else: raise psycopg2.InternalError(e) elif no_password_changes and role_attr_flags != '': # Grab role information from pg_roles instead of pg_authid select = "SELECT * FROM pg_roles where rolname=%(user)s" cursor.execute(select, {"user": user}) # Grab current role attributes. current_role_attrs = cursor.fetchone() role_attr_flags_changing = False if role_attr_flags: role_attr_flags_dict = {} for r in role_attr_flags.split(' '): if r.startswith('NO'): role_attr_flags_dict[r.replace('NO', '', 1)] = False else: role_attr_flags_dict[r] = True for role_attr_name, role_attr_value in role_attr_flags_dict.items(): if current_role_attrs[PRIV_TO_AUTHID_COLUMN[role_attr_name]] != role_attr_value: role_attr_flags_changing = True if not role_attr_flags_changing: return False alter = ['ALTER USER %(user)s' % {"user": pg_quote_identifier(user, 'role')}] if role_attr_flags: alter.append('WITH %s' % role_attr_flags) try: cursor.execute(' '.join(alter)) except psycopg2.InternalError as e: if e.pgcode == '25006': # Handle errors due to read-only transactions indicated by pgcode 25006 # ERROR: cannot execute ALTER ROLE in a read-only transaction changed = False module.fail_json(msg=e.pgerror, exception=traceback.format_exc()) return changed else: raise psycopg2.InternalError(e) # Grab new role attributes. cursor.execute(select, {"user": user}) new_role_attrs = cursor.fetchone() # Detect any differences between current_ and new_role_attrs. changed = current_role_attrs != new_role_attrs return changed
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 test_invalid_quotes(identifier, id_type, msg): with pytest.raises(SQLParseError) as ex: pg_quote_identifier(identifier, id_type) ex.match(msg)
def main(): argument_spec = pgutils.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="postgres"), session_role=dict(type='str'), conn_limit=dict(type='str', 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'] raw_connection = state in ("dump", "restore") if not HAS_PSYCOPG2 and not raw_connection: module.fail_json(msg=missing_required_lib('psycopg2'), exception=PSYCOPG2_IMP_ERR) # 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: pgutils.ensure_libs(sslrootcert=module.params.get('ca_cert')) 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) 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) 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='Dump of database %s failed' % db, stdout=stdout, stderr=stderr, rc=rc, cmd=cmd) elif stderr and 'warning' not in str(stderr): module.fail_json(msg='Dump of database %s failed' % db, stdout=stdout, stderr=stderr, rc=1, cmd=cmd) else: module.exit_json(changed=True, msg='Dump of database %s has been done' % db, stdout=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_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 __set_db_owner(self): """Set the database owner.""" query = "ALTER DATABASE %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 __add_schema(self): return '.'.join([ pg_quote_identifier(self.schema, 'schema'), pg_quote_identifier(self.name, 'sequence') ])
def set_owner(self, username): query = "ALTER TABLE %s OWNER TO %s" % (pg_quote_identifier( self.name, 'table'), pg_quote_identifier(username, 'role')) self.executed_queries.append(query) return self.__exec_sql(query, ddl=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')) cursor.execute(query)
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_tblspace(self, tblspace): query = "ALTER TABLE %s SET TABLESPACE %s" % (pg_quote_identifier( self.name, 'table'), pg_quote_identifier(tblspace, 'database')) self.executed_queries.append(query) return self.__exec_sql(query, ddl=True)
def set_stor_params(self, params): query = "ALTER TABLE %s SET (%s)" % (pg_quote_identifier( self.name, 'table'), params) self.executed_queries.append(query) return self.__exec_sql(query, ddl=True)
def test_valid_quotes(identifier, quoted_identifier): assert pg_quote_identifier(identifier, 'table') == quoted_identifier