[u r 2/2] Configure mirroring per catalog and source (#7066)

Author: nadove-ucsc

deployments/anvildev/environment.py

@@ -108,6 +108,7 @@ def env() -> Mapping[str, Optional[str]]:
         'AZUL_CATALOGS': json.dumps({
             f'{catalog}{suffix}': dict(atlas=atlas,
                                        internal=internal,
+                                       mirror_max_file_size=1.5 * 1024 ** 3,
                                        plugins=dict(metadata=dict(name='anvil'),
                                                     repository=dict(name='tdr_anvil')),
                                        sources=list(filter(None, sources.values())))

deployments/dev/environment.py

@@ -182,6 +182,15 @@ def mkdict(previous_catalog: dict[str, str],
 ]))
 
 
+def mirror_max_file_size(atlas: str, internal: bool) -> int | None:
+    if atlas == 'hca':
+        return 1.5 * 1024 ** 3
+    elif atlas == 'lungmap':
+        return -1
+    else:
+        assert False, atlas
+
+
 def env() -> Mapping[str, Optional[str]]:
     """
     Returns a dictionary that maps environment variable names to values. The
@@ -217,6 +226,7 @@ def env() -> Mapping[str, Optional[str]]:
         'AZUL_CATALOGS': json.dumps({
             f'{catalog}{suffix}': dict(atlas=atlas,
                                        internal=internal,
+                                       mirror_max_file_size=mirror_max_file_size(atlas, internal),
                                        plugins=dict(metadata=dict(name='hca'),
                                                     repository=dict(name='tdr_hca')),
                                        sources=list(filter(None, sources.values())))

deployments/hammerbox/environment.py

@@ -1075,6 +1075,7 @@ def env() -> Mapping[str, Optional[str]]:
         'AZUL_CATALOGS': json.dumps({
             f'{catalog}{suffix}': dict(atlas=atlas,
                                        internal=internal,
+                                       mirror_max_file_size=1.5 * 1024 ** 3,
                                        plugins=dict(metadata=dict(name='anvil'),
                                                     repository=dict(name='tdr_anvil')),
                                        sources=list(filter(None, sources.values())))

deployments/prod/environment.py

@@ -1763,6 +1763,15 @@ def mkdict(previous_catalog: dict[str, str],
 ]))
 
 
+def mirror_max_file_size(atlas: str, internal: bool) -> int | None:
+    if atlas == 'hca':
+        return 1.5 * 1024 ** 3 if internal else None
+    elif atlas == 'lungmap':
+        return -1
+    else:
+        assert False, atlas
+
+
 def env() -> Mapping[str, Optional[str]]:
     """
     Returns a dictionary that maps environment variable names to values. The
@@ -1797,6 +1806,7 @@ def env() -> Mapping[str, Optional[str]]:
         'AZUL_CATALOGS': base64.b64encode(bz2.compress(json.dumps({
             f'{catalog}{suffix}': dict(atlas=atlas,
                                        internal=internal,
+                                       mirror_max_file_size=mirror_max_file_size(atlas, internal),
                                        plugins=dict(metadata=dict(name='hca'),
                                                     repository=dict(name='tdr_hca')),
                                        sources=mklist(sources))

deployments/sandbox/environment.py

@@ -184,6 +184,15 @@ def mkdict(previous_catalog: dict[str, str],
 ]))
 
 
+def mirror_max_file_size(atlas: str, internal: bool) -> int | None:
+    if atlas == 'hca':
+        return 1.5 * 1024 ** 3
+    elif atlas == 'lungmap':
+        return -1
+    else:
+        assert False, atlas
+
+
 def env() -> Mapping[str, Optional[str]]:
     """
     Returns a dictionary that maps environment variable names to values. The
@@ -230,6 +239,7 @@ def env() -> Mapping[str, Optional[str]]:
         'AZUL_CATALOGS': json.dumps({
             f'{catalog}{suffix}': dict(atlas=atlas,
                                        internal=internal,
+                                       mirror_max_file_size=mirror_max_file_size(atlas, internal),
                                        plugins=dict(metadata=dict(name='hca'),
                                                     repository=dict(name='tdr_hca')),
                                        sources=list(filter(None, sources.values())))

environment.py

@@ -47,6 +47,7 @@ def env() -> Mapping[str, Optional[str]]:
         #   'name': {
         #       'atlas': 'bar',
         #       'internal': True,
+        #       'mirror_max_file_size': -1,
         #       'plugins': {
         #           plugin_type: {'name'=plugin_package},
         #           plugin_type: {'name'=plugin_package},
@@ -68,6 +69,10 @@ def env() -> Mapping[str, Optional[str]]:
         # `plugin_package` denotes the concrete implementation of how to fulfill
         # that purpose.
         #
+        # Files larger than `mirror_max_file_size` will not be mirrored for
+        # that catalog. This property may be absent, null, or negative, all of
+        # which will result in no files being mirrored for that catalog.
+        #
         # The first catalog listed is the default catalog.
         #
         # A source represents a TDR snapshot or canned staging area to index.

src/azul/__init__.py

@@ -63,6 +63,7 @@
     JSON,
     MutableJSON,
     json_bool,
+    json_int,
     json_mapping,
     json_sequence,
     json_str,
@@ -883,6 +884,7 @@ class Catalog:
         Config.Catalog(name='dcp',
                        atlas='hca',
                        internal=False,
+                       mirror_max_file_size=None,
                        plugins={'metadata': Config.Catalog.Plugin(name='hca'),
                                 'repository': Config.Catalog.Plugin(name='tdr_hca')},
                        sources=set())
@@ -918,6 +920,7 @@ def from_json(cls, spec: JSON) -> Self:
         name: str
         atlas: str
         internal: bool
+        mirror_max_file_size: int | None
         plugins: Mapping[str, Plugin]
         sources: Set[str]
 
@@ -981,9 +984,15 @@ def from_json(cls, name: str, spec: JSON) -> Self:
                 plugin_type: cls.Plugin.from_json(json_mapping(plugin_spec))
                 for plugin_type, plugin_spec in json_mapping(spec['plugins']).items()
             }
+            mirror_max_file_size = spec.get('mirror_max_file_size')
             return cls(name=name,
                        atlas=json_str(spec['atlas']),
                        internal=json_bool(spec['internal']),
+                       mirror_max_file_size=(
+                           None
+                           if mirror_max_file_size is None
+                           else json_int(mirror_max_file_size)
+                       ),
                        plugins=plugins,
                        sources=set(map(json_str, json_sequence(spec['sources']))))
 

src/azul/indexer/mirror_controller.py

@@ -130,17 +130,21 @@ def mirror_source(self, catalog: CatalogName, source_json: JSON):
         source = plugin.source_ref_cls.from_json(source_json)
         assert source.id in plugin.list_source_ids(authentication=None), R(
             'Cannot mirror non-public source', source)
-        source = plugin.partition_source_for_mirroring(catalog, source)
-        prefix = source.spec.prefix
-        log.info('Queueing %d partitions of source %r in catalog %r',
-                 prefix.num_partitions, str(source.spec), catalog)
+        if 'no_mirror' in source.spec.flags:
+            log.info('Not mirroring source % r in catalog %r because `no_mirror` flag is present',
+                     str(source.spec), str(catalog))
+        else:
+            source = plugin.partition_source_for_mirroring(catalog, source)
+            prefix = source.spec.prefix
+            log.info('Queueing %d partitions of source %r in catalog %r',
+                     prefix.num_partitions, str(source.spec), catalog)
 
-        def message(partition: str) -> SQSMessage:
-            log.debug('Queueing partition %r', partition)
-            return self.mirror_partition_message(catalog, source, partition)
+            def message(partition: str) -> SQSMessage:
+                log.debug('Queueing partition %r', partition)
+                return self.mirror_partition_message(catalog, source, partition)
 
-        messages = map(message, prefix.partition_prefixes())
-        self.client.queue_mirror_messages(messages)
+            messages = map(message, prefix.partition_prefixes())
+            self.client.queue_mirror_messages(messages)
 
     def mirror_partition(self,
                          catalog: CatalogName,
@@ -167,13 +171,9 @@ def mirror_file(self,
         file = self.load_file(catalog, file_json)
         assert file.size is not None, R('File size unknown', file)
 
-        file_is_large = file.size > 1.5 * 1024 ** 3
-        deployment_is_stable = (config.deployment.is_stable
-                                and not config.deployment.is_unit_test
-                                and catalog not in config.integration_test_catalogs)
-
+        max_file_size = config.catalogs[catalog].mirror_max_file_size
         service = self.service(catalog)
-        if file_is_large and not deployment_is_stable:
+        if max_file_size is not None and file.size > max_file_size:
             log.info('Not mirroring file to save cost: %r', file)
         elif service.info_exists(file):
             log.info('File is already mirrored, skipping upload: %r', file)