Allow to limit selection of data. (#15)

* Allow to limit selection of data. * Lint + Formatting
math280h · Dec 10, 2022 · ce07bce · ce07bce
1 parent add3d34
commit ce07bce
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 10 deletions.
diff --git a/test-docker-compose.yml → docker-compose.yml b/test-docker-compose.yml → docker-compose.yml
diff --git 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
@@ -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
@@ -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