diff --git a/dslr/cli.py b/dslr/cli.py index 411f359..89a060d 100644 --- a/dslr/cli.py +++ b/dslr/cli.py @@ -160,9 +160,10 @@ def list(): table = Table(box=box.SIMPLE) table.add_column("Name", style="cyan") table.add_column("Created") + table.add_column("Size", justify="right") for snapshot in sorted(snapshots, key=lambda s: s.created_at, reverse=True): - table.add_row(snapshot.name, timeago.format(snapshot.created_at)) + table.add_row(snapshot.name, timeago.format(snapshot.created_at), snapshot.size) cprint(table) diff --git a/dslr/operations.py b/dslr/operations.py index 30a8454..b35c43a 100644 --- a/dslr/operations.py +++ b/dslr/operations.py @@ -53,7 +53,7 @@ def drop_database(dbname: str): # Snapshot operations ################################################################################ -Snapshot = namedtuple("Snapshot", ["dbname", "name", "created_at"]) +Snapshot = namedtuple("Snapshot", ["dbname", "name", "created_at", "size"]) def generate_snapshot_db_name( @@ -79,22 +79,27 @@ def get_snapshots() -> List[Snapshot]: dslr__ """ # Find the snapshot databases - result = exec_sql("SELECT datname FROM pg_database WHERE datname LIKE 'dslr_%'") + result = exec_sql( + """ + SELECT + datname, + pg_size_pretty(pg_database_size(datname)) + FROM pg_database + WHERE datname LIKE 'dslr_%' + """ + ) if result is None: raise RuntimeError("Did not get results from database.") - snapshot_dbnames = sorted([row[0] for row in result]) - - # Parse the name into a Snapshot - parts = [dbname.split("_") for dbname in snapshot_dbnames] return [ Snapshot( dbname=line, name="_".join(part[2:]), created_at=datetime.fromtimestamp(int(part[1])), + size=size, ) - for part, line in zip(parts, snapshot_dbnames) + for line, part, size in [(row[0], row[0].split("_"), row[1]) for row in result] ] diff --git a/tests/test_cli.py b/tests/test_cli.py index 1a0b533..8e88ca2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -21,7 +21,7 @@ def stub_exec_sql(*args, **kwargs) -> List[Tuple[Any, ...]]: "existing-snapshot-2", created_at=datetime(2020, 1, 2, 0, 0, 0, 0), ) - return [(fake_snapshot_1,), (fake_snapshot_2,)] + return [(fake_snapshot_1, "100 kB"), (fake_snapshot_2, "100 kB")] @mock.patch.dict(os.environ, {"DATABASE_URL": "postgres://user:pw@test:5432/my_db"})