redis-sizer 0.2.0__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: redis-sizer
3
- Version: 0.2.0
3
+ Version: 0.3.0
4
4
  Summary: A simple command-line tool for analyzing memory usage across keys in a Redis database.
5
5
  Author-Email: Mienxiu <mienxiu@gmail.com>
6
6
  License-Expression: MIT
@@ -17,8 +17,7 @@ Description-Content-Type: text/markdown
17
17
 
18
18
  A simple command-line tool for analyzing memory usage across keys in a Redis database.
19
19
 
20
-
21
- ![Example table](docs/example_table.png)
20
+ ![Sample output](https://raw.githubusercontent.com/mienxiu/redis-sizer/refs/heads/main/docs/sample_output.png)
22
21
 
23
22
  redis-sizer can help with the following problems:
24
23
  - Optimize caching strategy.
@@ -51,7 +50,7 @@ Options:
51
50
  - `--pattern`: Pattern to filter keys [default: *]
52
51
  - `--sample-size`: Number of keys to sample [default: None]
53
52
  - `--namespace-separator`: Separator for key namespaces [default: :]
54
- - `--namespace-level`: Number of namespace levels to aggregate keys by. 0 means no aggregation. [default: 0]
53
+ - `--namespace-level`: Maximum number of namespace levels to aggregate keys by. 0 means no aggregation. [default: 0]
55
54
  - `--memory-unit`: Memory unit for display in result table [default: B]
56
55
  - `--top`: Maximum number of rows to display in result table [default: 10]
57
56
  - `--scan-count`: `COUNT` option for scanning keys [default: 1000]
@@ -62,49 +61,23 @@ Options:
62
61
  Many applications using Redis adopt hierarchical key naming conventions.
63
62
  redis-sizer leverages this structure to aggregate and analyze memory usage, helping you quickly identify which key groups consume the most memory.
64
63
 
65
- For example, if your keys follow a colon-delimited format (e.g., `users:profiles:001`, `users:logs:login:002`), the hierarchy can be visualized like a tree structure:
64
+ For example, if your keys follow a colon-delimited format (e.g., `users:profiles:...`, `users:logs:...`, `sessions:active:...`), the hierarchy can be visualized like a tree structure:
66
65
 
67
66
  ```
68
67
  DB
69
68
  ├── users:
70
69
  │ ├── profiles:
71
- │ │ ├── 001
72
- │ │ └── 002
70
+ │ │ ├── ...
73
71
  │ └── logs:
74
- │ ├── 001
75
- │ └── 002
72
+ │ ├── ...
76
73
  ├── sessions:
77
74
  │ └── active:
78
75
  ...
79
76
  ```
80
77
 
81
- Use `--namespace-level 1` to aggregate memory usage at the top-level namespace:
82
-
83
- ```console
84
- Key │ Count │ Size (MB) │ ... │ Memory Usage (%)
85
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
86
- users │ 350000 │ 2460.86 │ ... │ 93.82
87
- sessions │ 60000 │ 140.23 │ ... │ 5.35
88
- cache │ 10000 │ 20.07 │ ... │ 0.77
89
- tokens │ 700 │ 1.75 │ ... │ 0.07
90
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
91
- Total Keys Scanned │ 420700 │ 2622.90 │ ... │ 100.00
92
- ```
78
+ Using `--namespace-level 2` prints out memory usage tabels up to the second-level namespace:
93
79
 
94
- Use `--namespace-level 2` to aggregate memory usage at the second-level namespace:
95
-
96
- ```console
97
- Key │ Count │ Size (MB) │ ... │ Memory Usage (%)
98
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
99
- users:profiles │ 100000 │ 1177.22 │ ... │ 44.88
100
- users:logs │ 200000 │ 988.01 │ ... │ 37.67
101
- sessions:active │ 50000 │ 295.64 │ ... │ 11.27
102
- cache:images │ 40000 │ 100.02 │ ... │ 3.81
103
- cache:queries │ 20000 │ 40.21 │ ... │ 1.53
104
- tokens:auth │ 700 │ 1.75 │ ... │ 0.07
105
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
106
- Total Keys Scanned │ 420700 │ 2622.90 │ ... | 100.00
107
- ```
80
+ ![Example output](https://raw.githubusercontent.com/mienxiu/redis-sizer/refs/heads/main/docs/example_output.png)
108
81
 
109
82
  ## Considerations
110
83
 
@@ -2,8 +2,7 @@
2
2
 
3
3
  A simple command-line tool for analyzing memory usage across keys in a Redis database.
4
4
 
5
-
6
- ![Example table](docs/example_table.png)
5
+ ![Sample output](https://raw.githubusercontent.com/mienxiu/redis-sizer/refs/heads/main/docs/sample_output.png)
7
6
 
8
7
  redis-sizer can help with the following problems:
9
8
  - Optimize caching strategy.
@@ -36,7 +35,7 @@ Options:
36
35
  - `--pattern`: Pattern to filter keys [default: *]
37
36
  - `--sample-size`: Number of keys to sample [default: None]
38
37
  - `--namespace-separator`: Separator for key namespaces [default: :]
39
- - `--namespace-level`: Number of namespace levels to aggregate keys by. 0 means no aggregation. [default: 0]
38
+ - `--namespace-level`: Maximum number of namespace levels to aggregate keys by. 0 means no aggregation. [default: 0]
40
39
  - `--memory-unit`: Memory unit for display in result table [default: B]
41
40
  - `--top`: Maximum number of rows to display in result table [default: 10]
42
41
  - `--scan-count`: `COUNT` option for scanning keys [default: 1000]
@@ -47,49 +46,23 @@ Options:
47
46
  Many applications using Redis adopt hierarchical key naming conventions.
48
47
  redis-sizer leverages this structure to aggregate and analyze memory usage, helping you quickly identify which key groups consume the most memory.
49
48
 
50
- For example, if your keys follow a colon-delimited format (e.g., `users:profiles:001`, `users:logs:login:002`), the hierarchy can be visualized like a tree structure:
49
+ For example, if your keys follow a colon-delimited format (e.g., `users:profiles:...`, `users:logs:...`, `sessions:active:...`), the hierarchy can be visualized like a tree structure:
51
50
 
52
51
  ```
53
52
  DB
54
53
  ├── users:
55
54
  │ ├── profiles:
56
- │ │ ├── 001
57
- │ │ └── 002
55
+ │ │ ├── ...
58
56
  │ └── logs:
59
- │ ├── 001
60
- │ └── 002
57
+ │ ├── ...
61
58
  ├── sessions:
62
59
  │ └── active:
63
60
  ...
64
61
  ```
65
62
 
66
- Use `--namespace-level 1` to aggregate memory usage at the top-level namespace:
67
-
68
- ```console
69
- Key │ Count │ Size (MB) │ ... │ Memory Usage (%)
70
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
71
- users │ 350000 │ 2460.86 │ ... │ 93.82
72
- sessions │ 60000 │ 140.23 │ ... │ 5.35
73
- cache │ 10000 │ 20.07 │ ... │ 0.77
74
- tokens │ 700 │ 1.75 │ ... │ 0.07
75
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
76
- Total Keys Scanned │ 420700 │ 2622.90 │ ... │ 100.00
77
- ```
63
+ Using `--namespace-level 2` prints out memory usage tabels up to the second-level namespace:
78
64
 
79
- Use `--namespace-level 2` to aggregate memory usage at the second-level namespace:
80
-
81
- ```console
82
- Key │ Count │ Size (MB) │ ... │ Memory Usage (%)
83
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
84
- users:profiles │ 100000 │ 1177.22 │ ... │ 44.88
85
- users:logs │ 200000 │ 988.01 │ ... │ 37.67
86
- sessions:active │ 50000 │ 295.64 │ ... │ 11.27
87
- cache:images │ 40000 │ 100.02 │ ... │ 3.81
88
- cache:queries │ 20000 │ 40.21 │ ... │ 1.53
89
- tokens:auth │ 700 │ 1.75 │ ... │ 0.07
90
- ╶────────────────────┼─────────┼───────────┼─────┼──────────────────╴
91
- Total Keys Scanned │ 420700 │ 2622.90 │ ... | 100.00
92
- ```
65
+ ![Example output](https://raw.githubusercontent.com/mienxiu/redis-sizer/refs/heads/main/docs/example_output.png)
93
66
 
94
67
  ## Considerations
95
68
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "redis-sizer"
3
- version = "0.2.0"
3
+ version = "0.3.0"
4
4
  description = "A simple command-line tool for analyzing memory usage across keys in a Redis database."
5
5
  authors = [
6
6
  { name = "Mienxiu", email = "mienxiu@gmail.com" },
@@ -0,0 +1,401 @@
1
+ import time
2
+ from dataclasses import dataclass
3
+ from enum import Enum
4
+ from typing import Annotated
5
+
6
+ import typer
7
+ from redis import Redis
8
+ from redis.exceptions import RedisError
9
+ from rich import box
10
+ from rich.console import Console
11
+ from rich.progress import track
12
+ from rich.table import Table
13
+
14
+ app = typer.Typer()
15
+
16
+
17
+ class MemoryUnit(str, Enum):
18
+ """
19
+ For displaying memory usage in different units.
20
+ """
21
+
22
+ B = "B"
23
+ KB = "KB"
24
+ MB = "MB"
25
+ GB = "GB"
26
+
27
+
28
+ @dataclass
29
+ class TableRow:
30
+ """
31
+ Represent a row in the result table.
32
+ """
33
+
34
+ key: str
35
+ count: str
36
+ size: str
37
+ avg_size: str
38
+ min_size: str
39
+ max_size: str
40
+ percentage: str
41
+
42
+
43
+ @app.command()
44
+ def analyze(
45
+ host: str,
46
+ port: Annotated[int, typer.Option(help="Port number")] = 6379,
47
+ db: Annotated[int, typer.Option(help="DB number")] = 0,
48
+ password: Annotated[str | None, typer.Option(help="Password")] = None,
49
+ socket_timeout: Annotated[int, typer.Option(help="Socket timeout in seconds")] = 10,
50
+ socket_connect_timeout: Annotated[
51
+ int, typer.Option(help="Socket connect timeout in seconds")
52
+ ] = 10,
53
+ pattern: Annotated[str, typer.Option(help="Pattern to filter keys")] = "*",
54
+ sample_size: Annotated[int | None, typer.Option(help="Number of keys to sample")] = None,
55
+ namespace_separator: Annotated[str, typer.Option(help="Separator for key namespaces")] = ":",
56
+ namespace_level: Annotated[
57
+ int,
58
+ typer.Option(
59
+ min=0,
60
+ help="Maximum number of namespace levels to aggregate keys by. 0 means no aggregation.",
61
+ ),
62
+ ] = 0,
63
+ memory_unit: Annotated[
64
+ MemoryUnit, typer.Option(help="Memory unit for display in result table")
65
+ ] = MemoryUnit.B,
66
+ top: Annotated[
67
+ int | None, typer.Option(help="Maximum number of rows to display in result table")
68
+ ] = 10,
69
+ scan_count: Annotated[int, typer.Option(help="COUNT option for scanning keys")] = 1000,
70
+ batch_size: Annotated[int, typer.Option(help="Batch size for calculating memory usage")] = 1000,
71
+ ):
72
+ """
73
+ Analyze memory usage across keys in a Redis database and display the results in a table.
74
+ """
75
+ # Start the timer to measure execution time
76
+ start_time = time.time()
77
+
78
+ # Create Console instance for rich output
79
+ console = Console()
80
+
81
+ # Create Redis client
82
+ redis = Redis(
83
+ host=host,
84
+ port=port,
85
+ db=db,
86
+ password=password,
87
+ socket_timeout=socket_timeout,
88
+ socket_connect_timeout=socket_connect_timeout,
89
+ )
90
+
91
+ # Get the total number of keys in the database
92
+ try:
93
+ total_size: int = redis.dbsize() # type: ignore
94
+ except RedisError as error:
95
+ console.print(f"[red]Error occured: {error}[/red]")
96
+ redis.close()
97
+ exit(1)
98
+
99
+ # If the total size is 0, stop the process
100
+ if total_size == 0:
101
+ console.print("[yellow]No keys found in the database.[/yellow]")
102
+ redis.close()
103
+ return
104
+
105
+ console.print(f"The total number of keys: {total_size}")
106
+
107
+ # Scan the keys in the database
108
+ keys = _scan_keys(
109
+ redis=redis,
110
+ pattern=pattern,
111
+ count=scan_count,
112
+ sample_size=sample_size,
113
+ console=console,
114
+ total=total_size,
115
+ )
116
+ if not keys:
117
+ console.print(f"[yellow]No keys found matching the pattern: {pattern}[/yellow]")
118
+ redis.close()
119
+ return
120
+
121
+ # Get memory usage for each key (in batches)
122
+ memory_usage_by_key: dict[str, int | None] = {} # key -> memory usage
123
+ for i in track(
124
+ range(0, len(keys), batch_size),
125
+ description="Calculating memory usage...",
126
+ console=console,
127
+ ):
128
+ batch = keys[i : i + batch_size]
129
+ try:
130
+ memory_usage_by_key.update(_get_memory_usage(redis=redis, keys=batch))
131
+ except RedisError as error:
132
+ console.print(f"[red]Error occured: {error}[/red]")
133
+ redis.close()
134
+ exit(2)
135
+
136
+ redis.close()
137
+
138
+ for level in range(0, namespace_level + 1):
139
+ # Aggregate memory usage for this namespace level
140
+ group_data: dict[str, dict] = {} # group -> {keys: [], size: 0, sizes: []}
141
+ overall_min: int | None = None
142
+ overall_max: int = 0
143
+ valid_keys_count: int = 0
144
+ for key in keys:
145
+ memory_usage: int | None = memory_usage_by_key[key]
146
+ if memory_usage is None:
147
+ continue
148
+ group: str = _parse_key_group(key, namespace_separator, level)
149
+ if group not in group_data:
150
+ group_data[group] = {"keys": [], "size": 0, "sizes": []}
151
+ group_data[group]["keys"].append(key)
152
+ group_data[group]["size"] += memory_usage
153
+ group_data[group]["sizes"].append(memory_usage)
154
+ overall_min = memory_usage if overall_min is None else min(overall_min, memory_usage)
155
+ overall_max = max(overall_max, memory_usage)
156
+ valid_keys_count += 1
157
+
158
+ if valid_keys_count == 0:
159
+ console.print(
160
+ f"[yellow]No valid keys found for namespace level {level}. The scanned keys might have expired.[/yellow]"
161
+ )
162
+ continue
163
+
164
+ # Sort the groups by size in descending order and limit if 'top' is specified
165
+ sorted_groups = sorted(
166
+ (
167
+ (group, data["size"], data["sizes"], len(data["keys"]))
168
+ for group, data in group_data.items()
169
+ ),
170
+ key=lambda x: x[1],
171
+ reverse=True,
172
+ )
173
+ total_key_groups: int = len(sorted_groups)
174
+ if top is not None and top > 0 and top < total_key_groups:
175
+ group_data_for_rows = sorted_groups[:top]
176
+ hidden_count = total_key_groups - top
177
+ else:
178
+ group_data_for_rows = sorted_groups
179
+ hidden_count = 0
180
+
181
+ # Prepare table rows for current namespace level
182
+ rows: list[TableRow] = []
183
+ factor: int = _get_memory_unit_factor(memory_unit)
184
+ total_size: int = sum(group["size"] for group in group_data.values())
185
+ for group, size, sizes, count in group_data_for_rows:
186
+ display_key: str = group if level == 0 else f"{group}[dim]*[/dim]"
187
+ size_converted: float = size / factor
188
+ display_size: str = (
189
+ f"{size_converted:.2f}" if memory_unit != MemoryUnit.B else f"{int(size_converted)}"
190
+ )
191
+ avg_usage = (sum(sizes) / len(sizes)) / factor if sizes else 0
192
+ min_usage = min(sizes) / factor if sizes else 0
193
+ max_usage = max(sizes) / factor if sizes else 0
194
+ percentage = (size / total_size * 100) if total_size else 0
195
+ rows.append(
196
+ TableRow(
197
+ key=display_key,
198
+ count=str(count),
199
+ size=display_size,
200
+ avg_size=f"{avg_usage:.4f}"
201
+ if memory_unit != MemoryUnit.B
202
+ else f"{int(avg_usage)}",
203
+ min_size=f"{min_usage:.4f}"
204
+ if memory_unit != MemoryUnit.B
205
+ else f"{int(min_usage)}",
206
+ max_size=f"{max_usage:.4f}"
207
+ if memory_unit != MemoryUnit.B
208
+ else f"{int(max_usage)}",
209
+ percentage=f"{percentage:.2f}",
210
+ )
211
+ )
212
+
213
+ # Prepare the total row
214
+ total_usage_display: str = (
215
+ f"{total_size / factor:.2f}"
216
+ if memory_unit != MemoryUnit.B
217
+ else f"{int(total_size / factor)}"
218
+ )
219
+ overall_avg: float = (total_size / valid_keys_count) / factor
220
+ overall_min_conv: float = (overall_min or 0) / factor
221
+ overall_max_conv: float = overall_max / factor
222
+ overall_avg_display: str = (
223
+ f"{overall_avg:.4f}" if memory_unit != MemoryUnit.B else f"{int(overall_avg)}"
224
+ )
225
+ overall_min_display: str = (
226
+ f"{overall_min_conv:.4f}" if memory_unit != MemoryUnit.B else f"{int(overall_min_conv)}"
227
+ )
228
+ overall_max_display: str = (
229
+ f"{overall_max_conv:.4f}" if memory_unit != MemoryUnit.B else f"{int(overall_max_conv)}"
230
+ )
231
+ total_count: int = sum(len(group["keys"]) for group in group_data.values())
232
+ total_row: TableRow = TableRow(
233
+ key="Total Keys Scanned",
234
+ count=str(total_count),
235
+ size=total_usage_display,
236
+ avg_size=overall_avg_display,
237
+ min_size=overall_min_display,
238
+ max_size=overall_max_display,
239
+ percentage="100.00",
240
+ )
241
+
242
+ _print_memory_usage_table(
243
+ title=f"Memory Usage Table for Namespace Level {level}",
244
+ rows=rows,
245
+ total_row=total_row,
246
+ hidden_count=hidden_count,
247
+ memory_unit=memory_unit,
248
+ show_extra_stats=(level != 0),
249
+ console=console,
250
+ )
251
+ console.print("") # Add a blank line between tables
252
+
253
+ console.print(f"Took {(time.time() - start_time):.2f} seconds")
254
+
255
+
256
+ def _scan_keys(
257
+ redis: Redis,
258
+ pattern: str,
259
+ count: int,
260
+ sample_size: int | None,
261
+ console: Console,
262
+ total: int | None = None,
263
+ ) -> list[str]:
264
+ """
265
+ Scan keys in the Redis database using the given pattern.
266
+ If sample_size is None, perform a full iteration and estimate progress using total.
267
+ """
268
+ keys = []
269
+ for key in track(
270
+ redis.scan_iter(match=pattern, count=count),
271
+ description="Scanning keys...",
272
+ total=sample_size or total,
273
+ console=console,
274
+ show_speed=False,
275
+ ):
276
+ keys.append(key.decode())
277
+ if sample_size and sample_size <= len(keys):
278
+ break
279
+ return keys
280
+
281
+
282
+ def _parse_key_group(key: str, separator: str, level: int) -> str:
283
+ """
284
+ Parse the key group of a key up to the specified level using the given separator.
285
+ For example, _parse_key_group("a:b:c:d", ":", 2) -> "a:b"
286
+
287
+ If level is 0, return the entire key.
288
+ """
289
+ if level == 0:
290
+ return key
291
+ return separator.join(key.split(separator)[:level])
292
+
293
+
294
+ def _get_memory_usage(redis: Redis, keys: list[str]) -> dict[str, int | None]:
295
+ """
296
+ Get memory usage for a list of keys using Lua script.
297
+ If the key doesn't exist (expired), memory usage is None.
298
+ """
299
+ script = """
300
+ local result = {}
301
+ for i, key in ipairs(KEYS) do
302
+ result[i] = redis.call('MEMORY', 'USAGE', key, 'SAMPLES', 0)
303
+ end
304
+ return result
305
+ """
306
+ func = redis.register_script(script)
307
+ return dict(zip(keys, func(keys=keys))) # type: ignore
308
+
309
+
310
+ def _get_memory_unit_factor(memory_unit: MemoryUnit) -> int:
311
+ """
312
+ Get the conversion factor for the specified memory unit.
313
+ """
314
+ if memory_unit == MemoryUnit.B:
315
+ return 1
316
+ elif memory_unit == MemoryUnit.KB:
317
+ return 1024
318
+ elif memory_unit == MemoryUnit.MB:
319
+ return 1024 * 1024
320
+ elif memory_unit == MemoryUnit.GB:
321
+ return 1024 * 1024 * 1024
322
+ else:
323
+ raise ValueError("Invalid value for memory_unit. Use B, KB, MB, or GB.")
324
+
325
+
326
+ def _print_memory_usage_table(
327
+ title: str,
328
+ rows: list[TableRow],
329
+ total_row: TableRow,
330
+ hidden_count: int,
331
+ memory_unit: MemoryUnit,
332
+ show_extra_stats: bool, # new parameter to control detail columns
333
+ console: Console,
334
+ ):
335
+ """
336
+ - show_extra_stats: If True, show (Avg, Min, Max) size in the table.
337
+ """
338
+ table: Table = Table(title=title, box=box.MINIMAL)
339
+ table.add_column("Key", justify="left")
340
+ table.add_column("Count", justify="right", style="green")
341
+ table.add_column(f"Size ({memory_unit.upper()})", justify="right", style="magenta")
342
+ if show_extra_stats:
343
+ table.add_column(f"Avg Size ({memory_unit.upper()})", justify="right", style="orange1")
344
+ table.add_column(f"Min Size ({memory_unit.upper()})", justify="right", style="yellow")
345
+ table.add_column(f"Max Size ({memory_unit.upper()})", justify="right", style="red")
346
+ table.add_column("Memory Usage (%)", justify="right", style="cyan")
347
+
348
+ for row in rows:
349
+ if show_extra_stats:
350
+ table.add_row(
351
+ row.key,
352
+ row.count,
353
+ row.size,
354
+ row.avg_size,
355
+ row.min_size,
356
+ row.max_size,
357
+ row.percentage,
358
+ )
359
+ else:
360
+ table.add_row(row.key, row.count, row.size, row.percentage)
361
+
362
+ if hidden_count > 0:
363
+ if show_extra_stats:
364
+ table.add_row(
365
+ f"... {hidden_count} more entries not shown ...",
366
+ "",
367
+ "",
368
+ "",
369
+ "",
370
+ "",
371
+ "",
372
+ style="dim",
373
+ )
374
+ else:
375
+ table.add_row(f"... {hidden_count} more entries not shown ...", "", "", "", style="dim")
376
+ table.add_section()
377
+
378
+ if show_extra_stats:
379
+ table.add_row(
380
+ total_row.key,
381
+ total_row.count,
382
+ total_row.size,
383
+ total_row.avg_size,
384
+ total_row.min_size,
385
+ total_row.max_size,
386
+ total_row.percentage,
387
+ style="bold",
388
+ )
389
+ else:
390
+ table.add_row(
391
+ total_row.key,
392
+ total_row.count,
393
+ total_row.size,
394
+ total_row.percentage,
395
+ style="bold",
396
+ )
397
+ console.print(table)
398
+
399
+
400
+ if __name__ == "__main__":
401
+ app()
@@ -198,12 +198,13 @@ class TestGetMemoryUnitFactor(unittest.TestCase):
198
198
  class TestPrintMemoryUsage(unittest.TestCase):
199
199
  """Test the _print_memory_usage_table function."""
200
200
 
201
- def _test_print_memory_usage(self) -> None:
201
+ def test_print_memory_usage(self) -> None:
202
202
  """
203
203
  Sanity check for _print_memory_usage_table.
204
204
  NOTE Verifying console output is flaky as results may vary based on the terminal width.
205
205
  """
206
206
  _print_memory_usage_table(
207
+ title="Test Memory Usage",
207
208
  rows=[
208
209
  TableRow(
209
210
  key="key",
@@ -226,6 +227,7 @@ class TestPrintMemoryUsage(unittest.TestCase):
226
227
  ),
227
228
  hidden_count=0,
228
229
  memory_unit=MemoryUnit.B,
230
+ show_extra_stats=False,
229
231
  console=Console(),
230
232
  )
231
233
 
@@ -1,355 +0,0 @@
1
- import time
2
- from dataclasses import dataclass
3
- from enum import Enum
4
- from typing import Annotated
5
-
6
- import typer
7
- from redis import Redis
8
- from redis.exceptions import RedisError
9
- from rich import box
10
- from rich.console import Console
11
- from rich.progress import track
12
- from rich.table import Table
13
-
14
- app = typer.Typer()
15
-
16
-
17
- class MemoryUnit(str, Enum):
18
- """
19
- For displaying memory usage in different units.
20
- """
21
-
22
- B = "B"
23
- KB = "KB"
24
- MB = "MB"
25
- GB = "GB"
26
-
27
-
28
- @dataclass
29
- class TableRow:
30
- """
31
- Represent a row in the result table.
32
- """
33
-
34
- key: str
35
- count: str
36
- size: str
37
- avg_size: str
38
- min_size: str
39
- max_size: str
40
- percentage: str
41
-
42
-
43
- @app.command()
44
- def analyze(
45
- host: str,
46
- port: Annotated[int, typer.Option(help="Port number")] = 6379,
47
- db: Annotated[int, typer.Option(help="DB number")] = 0,
48
- password: Annotated[str | None, typer.Option(help="Password")] = None,
49
- socket_timeout: Annotated[int, typer.Option(help="Socket timeout in seconds")] = 10,
50
- socket_connect_timeout: Annotated[
51
- int, typer.Option(help="Socket connect timeout in seconds")
52
- ] = 10,
53
- pattern: Annotated[str, typer.Option(help="Pattern to filter keys")] = "*",
54
- sample_size: Annotated[int | None, typer.Option(help="Number of keys to sample")] = None,
55
- namespace_separator: Annotated[str, typer.Option(help="Separator for key namespaces")] = ":",
56
- namespace_level: Annotated[
57
- int,
58
- typer.Option(
59
- min=0, help="Number of namespace levels to aggregate keys by. 0 means no aggregation."
60
- ),
61
- ] = 0,
62
- memory_unit: Annotated[
63
- MemoryUnit, typer.Option(help="Memory unit for display in result table")
64
- ] = MemoryUnit.B,
65
- top: Annotated[
66
- int | None, typer.Option(help="Maximum number of rows to display in result table")
67
- ] = 10,
68
- scan_count: Annotated[int, typer.Option(help="COUNT option for scanning keys")] = 1000,
69
- batch_size: Annotated[int, typer.Option(help="Batch size for calculating memory usage")] = 1000,
70
- ):
71
- """
72
- Analyze memory usage across keys in a Redis database and display the results in a table.
73
- """
74
- # Start the timer to measure execution time
75
- start_time = time.time()
76
-
77
- # Create Console instance for rich output
78
- console = Console()
79
-
80
- # Create Redis client
81
- redis = Redis(
82
- host=host,
83
- port=port,
84
- db=db,
85
- password=password,
86
- socket_timeout=socket_timeout,
87
- socket_connect_timeout=socket_connect_timeout,
88
- )
89
-
90
- # Get the total number of keys in the database
91
- try:
92
- total_size: int = redis.dbsize() # type: ignore
93
- except RedisError as error:
94
- console.print(f"[red]Error occured: {error}[/red]")
95
- redis.close()
96
- exit(1)
97
-
98
- # If the total size is 0, stop the process
99
- if total_size == 0:
100
- console.print("[yellow]No keys found in the database.[/yellow]")
101
- redis.close()
102
- return
103
-
104
- console.print(f"The total number of keys: {total_size}")
105
-
106
- # Scan the keys in the database
107
- keys = _scan_keys(
108
- redis=redis,
109
- pattern=pattern,
110
- count=scan_count,
111
- sample_size=sample_size,
112
- console=console,
113
- total=total_size,
114
- )
115
- if not keys:
116
- console.print(f"[yellow]No keys found matching the pattern: {pattern}[/yellow]")
117
- redis.close()
118
- return
119
-
120
- # Get memory usage for each key (in batches)
121
- memory_usage_by_key: dict[str, int | None] = {} # key -> memory usage
122
- for i in track(
123
- range(0, len(keys), batch_size),
124
- description="Calculating memory usage...",
125
- console=console,
126
- ):
127
- batch = keys[i : i + batch_size]
128
- try:
129
- memory_usage_by_key.update(_get_memory_usage(redis=redis, keys=batch))
130
- except RedisError as error:
131
- console.print(f"[red]Error occured: {error}[/red]")
132
- redis.close()
133
- exit(2)
134
-
135
- redis.close()
136
-
137
- # Aggregate memory usage by key groups
138
- group_data: dict[str, dict] = {} # group -> {keys: [], size: 0, sizes: []}
139
- overall_min: int | None = None
140
- overall_max: int = 0
141
- valid_keys_count: int = 0
142
- for key in keys:
143
- memory_usage: int | None = memory_usage_by_key[key]
144
- # Memory usage is None for expired keys
145
- if memory_usage is None:
146
- continue
147
- group: str = _parse_key_group(key, namespace_separator, namespace_level)
148
- if group not in group_data:
149
- group_data[group] = {"keys": [], "size": 0, "sizes": []}
150
- group_data[group]["keys"].append(key)
151
- group_data[group]["size"] += memory_usage
152
- group_data[group]["sizes"].append(memory_usage)
153
- overall_min = memory_usage if overall_min is None else min(overall_min, memory_usage)
154
- overall_max = max(overall_max, memory_usage)
155
- valid_keys_count += 1
156
-
157
- if valid_keys_count == 0:
158
- console.print("[yellow]No valid keys found. The scanned keys might have expired. [/yellow]")
159
- return
160
-
161
- # Sort the groups by size in descending order
162
- sorted_groups = sorted(
163
- (
164
- (group, data["size"], data["sizes"], len(data["keys"]))
165
- for group, data in group_data.items()
166
- ),
167
- key=lambda x: x[1],
168
- reverse=True,
169
- )
170
-
171
- # Get hidden_count to limit the number of displayed groups if top is specified
172
- total_key_groups: int = len(sorted_groups)
173
- if top is not None and top > 0 and top < total_key_groups:
174
- group_data_for_rows = sorted_groups[:top]
175
- hidden_count = total_key_groups - top
176
- else:
177
- group_data_for_rows = sorted_groups
178
- hidden_count = 0
179
-
180
- # Prepare the rows
181
- rows: list[TableRow] = []
182
- factor: int = _get_memory_unit_factor(memory_unit)
183
- total_size: int = sum(group["size"] for group in group_data.values())
184
- for group, size, sizes, count in group_data_for_rows:
185
- display_key: str = group if namespace_level == 0 else f"{group}[dim]*[/dim]"
186
- size_converted: float = size / factor
187
- display_size: str = (
188
- f"{size_converted:.2f}" if memory_unit != MemoryUnit.B else f"{int(size_converted)}"
189
- )
190
- avg_usage = (sum(sizes) / len(sizes)) / factor if sizes else 0
191
- min_usage = min(sizes) / factor if sizes else 0
192
- max_usage = max(sizes) / factor if sizes else 0
193
- percentage = (size / total_size * 100) if total_size else 0
194
- rows.append(
195
- TableRow(
196
- key=display_key,
197
- count=str(count),
198
- size=display_size,
199
- avg_size=f"{avg_usage:.4f}" if memory_unit != MemoryUnit.B else f"{int(avg_usage)}",
200
- min_size=f"{min_usage:.4f}" if memory_unit != MemoryUnit.B else f"{int(min_usage)}",
201
- max_size=f"{max_usage:.4f}" if memory_unit != MemoryUnit.B else f"{int(max_usage)}",
202
- percentage=f"{percentage:.2f}",
203
- )
204
- )
205
-
206
- # Prepare the total row
207
- total_usage_display: str = (
208
- f"{total_size / factor:.2f}"
209
- if memory_unit != MemoryUnit.B
210
- else f"{int(total_size / factor)}"
211
- )
212
- overall_avg: float = (total_size / valid_keys_count) / factor
213
- overall_min_conv: float = (overall_min or 0) / factor
214
- overall_max_conv: float = overall_max / factor
215
- overall_avg_display: str = (
216
- f"{overall_avg:.4f}" if memory_unit != MemoryUnit.B else f"{int(overall_avg)}"
217
- )
218
- overall_min_display: str = (
219
- f"{overall_min_conv:.4f}" if memory_unit != MemoryUnit.B else f"{int(overall_min_conv)}"
220
- )
221
- overall_max_display: str = (
222
- f"{overall_max_conv:.4f}" if memory_unit != MemoryUnit.B else f"{int(overall_max_conv)}"
223
- )
224
- total_count: int = sum(len(group["keys"]) for group in group_data.values())
225
- total_row: TableRow = TableRow(
226
- key="Total Keys Scanned",
227
- count=str(total_count),
228
- size=total_usage_display,
229
- avg_size=overall_avg_display,
230
- min_size=overall_min_display,
231
- max_size=overall_max_display,
232
- percentage="100.00",
233
- )
234
-
235
- _print_memory_usage_table(
236
- rows=rows,
237
- total_row=total_row,
238
- hidden_count=hidden_count,
239
- memory_unit=memory_unit,
240
- console=console,
241
- )
242
-
243
- # Print the total execution time taken
244
- console.print(f"Took {(time.time() - start_time):.2f} seconds")
245
-
246
-
247
- def _scan_keys(
248
- redis: Redis,
249
- pattern: str,
250
- count: int,
251
- sample_size: int | None,
252
- console: Console,
253
- total: int | None = None,
254
- ) -> list[str]:
255
- """
256
- Scan keys in the Redis database using the given pattern.
257
- If sample_size is None, perform a full iteration and estimate progress using total.
258
- """
259
- keys = []
260
- for key in track(
261
- redis.scan_iter(match=pattern, count=count),
262
- description="Scanning keys...",
263
- total=sample_size or total,
264
- console=console,
265
- show_speed=False,
266
- ):
267
- keys.append(key.decode())
268
- if sample_size and sample_size <= len(keys):
269
- break
270
- return keys
271
-
272
-
273
- def _parse_key_group(key: str, separator: str, level: int) -> str:
274
- """
275
- Parse the key group of a key up to the specified level using the given separator.
276
- For example, _parse_key_group("a:b:c:d", ":", 2) -> "a:b"
277
-
278
- If level is 0, return the entire key.
279
- """
280
- if level == 0:
281
- return key
282
- return separator.join(key.split(separator)[:level])
283
-
284
-
285
- def _get_memory_usage(redis: Redis, keys: list[str]) -> dict[str, int | None]:
286
- """
287
- Get memory usage for a list of keys using Lua script.
288
- If the key doesn't exist (expired), memory usage is None.
289
- """
290
- script = """
291
- local result = {}
292
- for i, key in ipairs(KEYS) do
293
- result[i] = redis.call('MEMORY', 'USAGE', key, 'SAMPLES', 0)
294
- end
295
- return result
296
- """
297
- func = redis.register_script(script)
298
- return dict(zip(keys, func(keys=keys))) # type: ignore
299
-
300
-
301
- def _get_memory_unit_factor(memory_unit: MemoryUnit) -> int:
302
- """
303
- Get the conversion factor for the specified memory unit.
304
- """
305
- if memory_unit == MemoryUnit.B:
306
- return 1
307
- elif memory_unit == MemoryUnit.KB:
308
- return 1024
309
- elif memory_unit == MemoryUnit.MB:
310
- return 1024 * 1024
311
- elif memory_unit == MemoryUnit.GB:
312
- return 1024 * 1024 * 1024
313
- else:
314
- raise ValueError("Invalid value for memory_unit. Use B, KB, MB, or GB.")
315
-
316
-
317
- def _print_memory_usage_table(
318
- rows: list[TableRow],
319
- total_row: TableRow,
320
- hidden_count: int,
321
- memory_unit: MemoryUnit,
322
- console: Console,
323
- ):
324
- table: Table = Table(box=box.MINIMAL)
325
- table.add_column("Key", justify="left")
326
- table.add_column("Count", justify="right", style="green")
327
- table.add_column(f"Size ({memory_unit.upper()})", justify="right", style="magenta")
328
- table.add_column(f"Avg Size ({memory_unit.upper()})", justify="right", style="orange1")
329
- table.add_column(f"Min Size ({memory_unit.upper()})", justify="right", style="yellow")
330
- table.add_column(f"Max Size ({memory_unit.upper()})", justify="right", style="red")
331
- table.add_column("Memory Usage (%)", justify="right", style="cyan")
332
- for row in rows:
333
- table.add_row(
334
- row.key, row.count, row.size, row.avg_size, row.min_size, row.max_size, row.percentage
335
- )
336
- if 0 < hidden_count:
337
- table.add_row(
338
- f"... {hidden_count} more entries not shown ...", "", "", "", "", "", "", style="dim"
339
- )
340
- table.add_section()
341
- table.add_row(
342
- total_row.key,
343
- total_row.count,
344
- total_row.size,
345
- total_row.avg_size,
346
- total_row.min_size,
347
- total_row.max_size,
348
- total_row.percentage,
349
- style="bold",
350
- )
351
- console.print(table)
352
-
353
-
354
- if __name__ == "__main__":
355
- app()
File without changes