add size to snapshot list

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)
 

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_<timestamp>_<snapshot_name>
     """
     # 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]
     ]
 
 

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"})