From e7398ad58872c1dbc727aa7b6b6ff59e6586d7ab Mon Sep 17 00:00:00 2001 From: Ahmed Et-tanany Date: Wed, 3 Jan 2024 18:37:49 +0100 Subject: [PATCH] Support excluding database roles --- aiven_db_migrate/migrate/pgmigrate.py | 50 ++++++++++++++++++++------- test/test_pg_roles.py | 30 +++++++++++++++- 2 files changed, 66 insertions(+), 14 deletions(-) diff --git a/aiven_db_migrate/migrate/pgmigrate.py b/aiven_db_migrate/migrate/pgmigrate.py index 0eda888..22110b7 100644 --- a/aiven_db_migrate/migrate/pgmigrate.py +++ b/aiven_db_migrate/migrate/pgmigrate.py @@ -109,7 +109,13 @@ class PGCluster: _mangle: bool _has_aiven_gatekeper: Optional[bool] - def __init__(self, conn_info: Union[str, Dict[str, Any]], filtered_db: Optional[str] = None, mangle: bool = False): + def __init__( + self, + conn_info: Union[str, Dict[str, Any]], + filtered_db: Optional[str] = None, + excluded_roles: Optional[str] = None, + mangle: bool = False, + ): self.log = logging.getLogger(self.__class__.__name__) self.conn_info = get_connection_info(conn_info) self.conn_lock = threading.RLock() @@ -122,10 +128,8 @@ def __init__(self, conn_info: Union[str, Dict[str, Any]], filtered_db: Optional[ self._pg_ext_whitelist = None self._pg_lang = None self._pg_roles = dict() - if filtered_db: - self.filtered_db = filtered_db.split(",") - else: - self.filtered_db = [] + self.filtered_db = filtered_db.split(",") if filtered_db else [] + self.excluded_roles = excluded_roles.split(",") if excluded_roles else [] if "application_name" not in self.conn_info: self.conn_info["application_name"] = f"aiven-db-migrate/{__version__}" self._mangle = mangle @@ -320,6 +324,9 @@ def pg_roles(self) -> Dict[str, PGRole]: if not self._pg_roles: # exclude system roles roles = self.c("SELECT quote_ident(rolname) as safe_rolname, * FROM pg_catalog.pg_roles WHERE oid > 16384") + # Filter out excluded roles. + roles = [r for r in roles if r["rolname"] not in self.excluded_roles] + for r in roles: rolname = r["rolname"] # create semi-random placeholder password for role with login @@ -789,6 +796,7 @@ def __init__( verbose: bool = False, mangle: bool = False, filtered_db: Optional[str] = None, + excluded_roles: Optional[str] = None, skip_tables: Optional[List[str]] = None, with_tables: Optional[List[str]] = None, replicate_extensions: bool = True, @@ -796,8 +804,18 @@ def __init__( if skip_tables and with_tables: raise Exception("Can only specify a skip table list or a with table list") self.log = logging.getLogger(self.__class__.__name__) - self.source = PGSource(conn_info=source_conn_info, filtered_db=filtered_db, mangle=mangle) - self.target = PGTarget(conn_info=target_conn_info, filtered_db=filtered_db, mangle=mangle) + self.source = PGSource( + conn_info=source_conn_info, + filtered_db=filtered_db, + excluded_roles=excluded_roles, + mangle=mangle, + ) + self.target = PGTarget( + conn_info=target_conn_info, + filtered_db=filtered_db, + excluded_roles=excluded_roles, + mangle=mangle, + ) self.skip_tables = self._convert_table_names(skip_tables) self.with_tables = self._convert_table_names(with_tables) if pgbin is None: @@ -1047,14 +1065,13 @@ def _check_aiven_pg_security_agent(self) -> None: if unauthorized_roles: self.log.error("Some superuser roles from source database are not allowed in target database.") - exc_message_with_hints = ( - "Some superuser roles from source database are not allowed in target database. " - "Hints: To see all superusers in the source DB, run 'SELECT rolname FROM pg_roles WHERE rolsuper;'. " - "To see which roles are allowed in the target DB, search in 'postgresql.conf' file for " - "'aiven.pg_security_agent_reserved_roles'." + exc_message = ( + f"Some superuser roles from source database {tuple(unauthorized_roles)} are not " + "allowed in target database. You can delete them from the source or exclude them from " + "the migration if you are sure your database will function properly without them." ) - raise PGMigrateValidationFailedError(exc_message_with_hints) + raise PGMigrateValidationFailedError(exc_message) self.log.info("All superuser roles from source database are allowed in target database.") @@ -1387,6 +1404,12 @@ def main(args=None, *, prog="pg_migrate"): parser.add_argument( "-f", "--filtered-db", help="Comma separated list of databases to filter out during migrations", required=False ) + parser.add_argument( + "-xr", + "--excluded-roles", + help="Comma separated list of database roles to exclude during migrations", + required=False, + ) parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output.") parser.add_argument( "--no-createdb", action="store_false", dest="createdb", help="Don't automatically create database(s) in target." @@ -1479,6 +1502,7 @@ def main(args=None, *, prog="pg_migrate"): stop_replication=args.stop_replication, verbose=args.verbose, filtered_db=args.filtered_db, + excluded_roles=args.excluded_roles, mangle=args.mangle, skip_tables=args.skip_table, with_tables=args.with_table, diff --git a/test/test_pg_roles.py b/test/test_pg_roles.py index ca6cc9f..b716112 100644 --- a/test/test_pg_roles.py +++ b/test/test_pg_roles.py @@ -247,11 +247,39 @@ def test_migration_fails_with_additional_superuser_roles(pg_source_and_target: T with pytest.raises( PGMigrateValidationFailedError, - match=r"Some superuser roles from source database are not allowed in target database.*", + match=r"Some superuser roles from source database .* are not allowed in target database.*", ): pg_mig.migrate() +def test_migration_succeeds_when_additional_superuser_roles_are_excluded( + pg_source_and_target: Tuple[PGRunner, PGRunner], +) -> None: + """Test that migration succeeds when non allowed superuser roles are excluded.""" + source, target = pg_source_and_target + regularuser = random_string() + superuser1 = random_string() + superuser2 = random_string() + source.add_cleanup(lambda: source.drop_user(username=regularuser)) + source.add_cleanup(lambda: source.drop_user(username=superuser1)) + source.add_cleanup(lambda: source.drop_user(username=superuser2)) + source.create_role(username=regularuser, password=random_string()) + source.create_role(username=superuser1, password=random_string(), superuser=True) + source.create_role(username=superuser2, password=random_string(), superuser=True) + + pg_mig = PGMigrate( + source_conn_info=source.conn_info(), + target_conn_info=target.super_conn_info(), + createdb=True, + verbose=True, + excluded_roles=f"{superuser1},{superuser2}", + ) + assert pg_mig.target.is_pg_security_agent_enabled + result = pg_mig.migrate() + + assert result.pg_roles.keys() == {regularuser, source.testuser} + + def test_migration_succeeds_with_authorized_superuser_role(pg_source_and_target: Tuple[PGRunner, PGRunner]) -> None: """Test that it succeeds when we try to migrate a superuser role that is in the reserved roles list of the target.""" source, target = pg_source_and_target