Skip to content

Commit

Permalink
filter roles feature improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
runwuf committed Apr 28, 2022
1 parent 2759321 commit 1a6cc73
Showing 1 changed file with 15 additions and 13 deletions.
28 changes: 15 additions & 13 deletions aiven_db_migrate/migrate/pgmigrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ def __init__(
self,
conn_info: Union[str, Dict[str, Any]],
filtered_db: Optional[str] = None,
filtered_roles: Optional[str] = None,
filtered_roles: Tuple[str] = ((),),
mangle: bool = False
):
self.log = logging.getLogger(self.__class__.__name__)
Expand All @@ -131,10 +131,8 @@ def __init__(
self.filtered_db = filtered_db.split(",")
else:
self.filtered_db = []
if filtered_roles:
self.filtered_roles = filtered_roles.split(",")
else:
self.filtered_roles = []
self.filtered_roles = filtered_roles

if "application_name" not in self.conn_info:
self.conn_info["application_name"] = f"aiven-db-migrate/{__version__}"
self._mangle = mangle
Expand Down Expand Up @@ -328,18 +326,19 @@ def pg_lang(self) -> List[Dict[str, Any]]:

@property
def pg_roles(self) -> Dict[str, PGRole]:
filtered_roles = ["rdstopmgr"]
if self.filtered_roles:
filtered_roles.extend(self.filtered_roles)
roles_list = ",".join(f"'{role}'" for role in filtered_roles)

if not self._pg_roles:
# exclude system roles
roles = self.c(
f"SELECT quote_ident(rolname) as safe_rolname, * FROM pg_catalog.pg_roles WHERE oid > 16384 AND rolname NOT IN ({roles_list})"
)
roles = self.c(f"SELECT quote_ident(rolname) as safe_rolname, * FROM pg_catalog.pg_roles WHERE oid > 16384")

for r in roles:
rolname = r["rolname"]

if rolname in self.filtered_roles:
self.log.debug(f'Skipping filtered role: [{r["rolname"]}]')
continue

self.log.debug(f'Migrating role: [{r["rolname"]}]')
# create semi-random placeholder password for role with login
rolpassword = (
"placeholder_{}".format("".join(random.choices(string.ascii_lowercase, k=16)))
Expand Down Expand Up @@ -755,7 +754,7 @@ def __init__(
verbose: bool = False,
mangle: bool = False,
filtered_db: Optional[str] = None,
filtered_roles: Optional[str] = None,
filtered_roles: Tuple[str] = ((),),
skip_tables: Optional[List[str]] = None,
with_tables: Optional[List[str]] = None,
replicate_extensions: bool = True,
Expand Down Expand Up @@ -1373,6 +1372,9 @@ def main(args=None, *, prog="pg_migrate"):
else:
logging.basicConfig(level=logging.INFO, format=log_format)

if not args.filtered_roles:
args.filtered_roles = ()

pg_mig = PGMigrate(
source_conn_info=args.source,
target_conn_info=args.target,
Expand Down

0 comments on commit 1a6cc73

Please sign in to comment.