diff --git a/test-docker-compose.yml b/docker-compose.yml similarity index 100% rename from test-docker-compose.yml rename to docker-compose.yml diff --git a/redactdump/app.py b/redactdump/app.py index 1d4534f..5c98a79 100644 --- a/redactdump/app.py +++ b/redactdump/app.py @@ -99,15 +99,25 @@ def dump(self, table: str) -> Tuple[str, int, Union[str, None]]: """ self.console.print(f":construction: [blue]Working on table:[/blue] {table}") - row_count = self.database.count_rows(table) + row_count = ( + self.database.count_rows(table) + if "limits" not in self.config.config + or "max_rows_per_table" not in self.config.config["limits"] + else int(self.config.config["limits"]["max_rows_per_table"]) + ) rows = self.database.get_row_names(table) last_num = 0 - step = 100 + step = ( + 100 + if "performance" not in self.config.config + or "rows_per_request" not in self.config.config["performance"] + else int(self.config.config["performance"]["rows_per_request"]) + ) location = None for x in range(0, row_count, step): - if x == 0: + if x == 0 and step < row_count: continue limit = step if x + step < row_count else step + row_count - x @@ -143,9 +153,18 @@ async def run(self) -> None: sorted_output = sorted(result, key=lambda d: d[1], reverse=True) + row_count_limited = ( + "" + if "limits" not in self.config.config + or "max_rows_per_table" not in self.config.config["limits"] + else " (Limited via config)" + ) + for res in sorted_output: table.add_row( - res[0], str(res[1]), res[2] if res[2] is not None else "No data" + res[0], + f"{str(res[1])}{row_count_limited}", + res[2] if res[2] is not None else "No data", ) self.console.print(table) diff --git a/redactdump/core/config.py b/redactdump/core/config.py index 555301c..b2fddd4 100644 --- a/redactdump/core/config.py +++ b/redactdump/core/config.py @@ -37,6 +37,9 @@ def load_config(self) -> dict: Optional("username"): str, Optional("password"): str, }, + Optional("limits"): {"max_rows_per_table": int, "select_columns": list}, + Optional("performance"): {"rows_per_request": int}, + Optional("debug"): {"enabled": bool}, "redact": { Optional("columns"): { str: [ @@ -83,4 +86,8 @@ def load_config(self) -> dict: except SchemaError as se: raise se + if "debug" not in config or "enabled" not in config["debug"]: + config["debug"] = {} + config["debug"]["enabled"] = False + return config diff --git a/redactdump/core/database.py b/redactdump/core/database.py index 335e130..168d2df 100644 --- a/redactdump/core/database.py +++ b/redactdump/core/database.py @@ -105,17 +105,29 @@ def get_data(self, table: str, rows: list, offset: int, limit: int) -> list: conn = conn.execution_options( postgresql_readonly=True, postgresql_deferrable=True ) + + if not set(self.config.config["limits"]["select_columns"]).issubset(rows): + return [] + with conn.begin(): - self.console.print( - f"[cyan]DEBUG: Running 'SELECT * FROM {table} OFFSET {offset} LIMIT {limit}'[/cyan]" + select = ( + "*" + if "limits" not in self.config.config + or "select_columns" not in self.config.config["limits"] + else ",".join(self.config.config["limits"]["select_columns"]) ) + + if self.config.config["debug"]["enabled"]: + self.console.print( + f"[cyan]DEBUG: Running 'SELECT {select} FROM {table} OFFSET {offset} LIMIT {limit}'[/cyan]" + ) + result = conn.execute( - text(f"SELECT * FROM {table} OFFSET {offset} LIMIT {limit}") + text(f"SELECT {select} FROM {table} OFFSET {offset} LIMIT {limit}") ) records = [dict(zip(row.keys(), row)) for row in result] - for item in records: - + print(item) if self.redactor.data_rules or self.redactor.column_rules: item = self.redactor.redact(item, rows) @@ -144,6 +156,14 @@ def get_row_names(self, table: str) -> list: ) ) + select_columns = ( + [] + if "limits" not in self.config.config + or "select_columns" not in self.config.config["limits"] + else self.config.config["limits"]["select_columns"] + ) + for item in result: - names.append(item[0]) + if not select_columns or item[0] in select_columns: + names.append(item[0]) return names