fixup! [u r] Fix: Partition sizing ignores supplementary bundles (#5207)

scripts/generate_sources.py → scripts/update_subgraph_counts.py

@@ -1,12 +1,9 @@
 """
-Determine sizes
+Count the number of subgraphs per configured source and produce output to
+expedite updated source configurations.
 """
 
 import argparse
-import json
-from pathlib import (
-    Path,
-)
 import sys
 
 import attr
@@ -23,6 +20,9 @@
 from azul.indexer import (
     Prefix,
 )
+from azul.openapi import (
+    format_description,
+)
 from azul.terra import (
     TDRSourceSpec,
 )
@@ -52,27 +52,17 @@ def __str__(self) -> str:
 
 azul = AzulClient()
 
-active_deployment_dir = Path(config.project_root) / 'deployments' / '.active'
-with open(active_deployment_dir / 'sources.json') as f:
-    raw_sources = json.load(f)
-
 
 def generate_sources(catalog: str) -> list[SourceSpecArgs]:
-    catalog_sources = raw_sources[catalog]
     plugin = azul.repository_plugin(catalog)
     sources = []
-    for source in catalog_sources:
-        project, snapshot = source.split('.')
-        spec = TDRSourceSpec(project=project,
-                             name=snapshot,
-                             is_snapshot=True,
-                             prefix=Prefix.of_everything)
+    for spec in plugin.sources:
+        spec: TDRSourceSpec = attr.evolve(spec, prefix=Prefix.of_everything)
         ref = plugin.resolve_source(str(spec))
         partitions = plugin.list_partitions(ref)
-        sources.append(SourceSpecArgs(project=project,
-                                      snapshot=snapshot,
+        sources.append(SourceSpecArgs(project=spec.project,
+                                      snapshot=spec.name,
                                       subgraph_count=sum(partitions.values())))
-
     return sources
 
 
@@ -86,9 +76,16 @@ def main(args: list[str]):
         spec_args_list.sort(key=lambda spec_args: spec_args.snapshot)
         print(',\n'.join(map(str, spec_args_list)))
 
-    msg = "Don't forget to manually specify flags for `ma` and `pop`."
-    print('-' * len(msg))
-    print(msg)
+    print(format_description('''
+        -----------------
+        !!! IMPORTANT !!!
+        -----------------
+
+        This script does *not* populate the `ma` or `pop` flags for the source
+        specs. Do not copy/paste the above output without checking whether these
+        flags should be applied. If `mksrc` generates a common prefix, manual
+        adjustment of the generated common prefix may be required.
+    '''))
 
 
 if __name__ == '__main__':