Refactor source partitioning

Author: nadove-ucsc

scripts/post_deploy_tdr.py

@@ -101,7 +101,7 @@ def verify_source(self,
                 require(source_spec.prefix.common == '', source_spec)
             self.tdr.check_bigquery_access(source_spec)
         else:
-            ref = plugin.partition_source(catalog, ref)
+            ref = plugin.partition_source_for_indexing(catalog, ref)
             subgraph_count = plugin.count_bundles(ref.spec)
             require(subgraph_count > 0, 'Common prefix is too long', ref.spec)
             require(subgraph_count <= 512, 'Common prefix is too short', ref.spec)

src/azul/azulclient.py

@@ -325,7 +325,7 @@ def remote_reindex(self,
         plugin = self.repository_plugin(catalog)
         for source_spec in sources:
             source_ref = plugin.resolve_source(source_spec)
-            source_ref = plugin.partition_source(catalog, source_ref)
+            source_ref = plugin.partition_source_for_indexing(catalog, source_ref)
 
             def message(partition_prefix: str) -> SQSMessage:
                 log.info('Remotely reindexing prefix %r of source_ref %r into catalog %r',
@@ -607,7 +607,7 @@ def message(source: SourceRef):
     def mirror_source(self, catalog: CatalogName, source_json: JSON):
         plugin = self.repository_plugin(catalog)
         source = plugin.source_ref_cls.from_json(source_json)
-        source = plugin.partition_source(catalog, source)
+        source = plugin.partition_source_for_indexing(catalog, source)
 
         def message(prefix: str) -> SQSMessage:
             log.info('Mirroring files in partition %r of source %r from catalog %r',

src/azul/plugins/__init__.py

@@ -11,6 +11,7 @@
 )
 from typing import (
     AbstractSet,
+    Callable,
     ClassVar,
     Iterable,
     Literal,
@@ -671,27 +672,35 @@ def count_bundles(self, source: SOURCE_SPEC) -> int:
         """
         raise NotImplementedError
 
-    def partition_source(self,
-                         catalog: CatalogName,
-                         source: SOURCE_REF
-                         ) -> SOURCE_REF:
+    def partition_source_for_indexing(self,
+                                      catalog: CatalogName,
+                                      source: SOURCE_REF
+                                      ) -> SOURCE_REF:
         """
         If the source already has a prefix, return the source. Otherwise, return
         an updated copy of the source with a heuristically computed prefix that
         should be appropriate for indexing in the given catalog.
         """
+        return self._partition_source(catalog, source, self.count_bundles)
+
+    def _partition_source(self,
+                          catalog: CatalogName,
+                          source: SOURCE_REF,
+                          counter: Callable[[SOURCE_SPEC], int]
+                          ) -> SOURCE_REF:
         if source.spec.prefix is None:
-            count = self.count_bundles(source.spec)
+            count = counter(source.spec)
             is_main = config.deployment.is_main
             is_it = catalog in config.integration_test_catalogs
-            # We use the "lesser" heuristic during IT to avoid indexing an
-            # excessive number of bundles
+            # We use the "lesser" heuristic during IT to keep the cost and
+            # performance of the tests within reasonable limits
             if is_main and not is_it:
                 prefix = Prefix.for_main_deployment(count)
             else:
                 prefix = Prefix.for_lesser_deployment(count)
-            source = source.with_prefix(prefix)
-        return source
+            return source.with_prefix(prefix)
+        else:
+            return source
 
     @abstractmethod
     def list_bundles(self,

test/integration_test.py

@@ -1277,7 +1277,7 @@ def _prepare_notifications(self,
         bundle_fqids = set()
         notifications = []
         for source in sources:
-            source = plugin.partition_source(catalog, source)
+            source = plugin.partition_source_for_indexing(catalog, source)
             # Some partitions may be empty, but we include them anyway to
             # ensure test coverage for handling multiple partitions per source
             for partition_prefix in source.spec.prefix.partition_prefixes():