fix: remove executor secret from flow mermaid chart (#4801)

jina/hubble/helper.py

@@ -189,6 +189,23 @@ def parse_hub_uri(uri_path: str) -> Tuple[str, str, str, str]:
     return scheme, name, tag, secret
 
 
+def replace_secret_of_hub_uri(uri_path: str, txt: str = '<secret>') -> str:
+    """Replace the secret of the Jina Hub URI.
+
+    :param uri_path: the uri of Jina Hub URI
+    :param txt: text to replace
+    :return: the new URI
+    """
+
+    try:
+        secret = parse_hub_uri(uri_path)[-1]
+        if secret:
+            return uri_path.replace(secret, txt)
+    except ValueError:
+        pass  # ignore if the URI is not a valid Jina Hub URI
+    return uri_path
+
+
 def is_valid_huburi(uri: str) -> bool:
     """Return True if it is a valid Hubble URI
 

jina/orchestrate/deployments/__init__.py

@@ -11,6 +11,7 @@
 from jina import __default_executor__, __default_host__, __docker_host__, helper
 from jina.enums import DeploymentRoleType, PodRoleType, PollingType
 from jina.helper import CatchAllCleanupContextManager
+from jina.hubble.helper import replace_secret_of_hub_uri
 from jina.hubble.hubio import HubIO
 from jina.jaml.helper import complete_path
 from jina.orchestrate.pods.container import ContainerPod
@@ -781,6 +782,7 @@ def _mermaid_str(self) -> List[str]:
         .. # noqa: DAR201
         """
         mermaid_graph = []
+        secret = '&ltsecret&gt'
         if self.role != DeploymentRoleType.GATEWAY and not self.external:
             mermaid_graph = [f'subgraph {self.name};', f'\ndirection LR;\n']
 
@@ -790,15 +792,17 @@ def _mermaid_str(self) -> List[str]:
                 else None
             )
             uses_before_uses = (
-                self.uses_before_args.uses
+                replace_secret_of_hub_uri(self.uses_before_args.uses, secret)
                 if self.uses_before_args is not None
                 else None
             )
             uses_after_name = (
                 self.uses_after_args.name if self.uses_after_args is not None else None
             )
             uses_after_uses = (
-                self.uses_after_args.uses if self.uses_after_args is not None else None
+                replace_secret_of_hub_uri(self.uses_after_args.uses, secret)
+                if self.uses_after_args is not None
+                else None
             )
             shard_names = []
             if len(self.pod_args['pods']) > 1:
@@ -818,7 +822,7 @@ def _mermaid_str(self) -> List[str]:
                     ]  # all the uses should be the same but let's keep it this
                     # way
                     for rep_i, (name, use) in enumerate(zip(names, uses)):
-                        escaped_uses = f'"{use}"'
+                        escaped_uses = f'"{replace_secret_of_hub_uri(use, secret)}"'
                         shard_mermaid_graph.append(f'{name}[{escaped_uses}]:::pod;')
                     shard_mermaid_graph.append('end;')
                     shard_mermaid_graph = [
@@ -828,7 +832,9 @@ def _mermaid_str(self) -> List[str]:
                     mermaid_graph.append('\n')
                 if uses_before_name is not None:
                     for shard_name in shard_names:
-                        escaped_uses_before_uses = f'"{uses_before_uses}"'
+                        escaped_uses_before_uses = (
+                            f'"{replace_secret_of_hub_uri(uses_before_uses, secret)}"'
+                        )
                         mermaid_graph.append(
                             f'{self.args.name}-head[{escaped_uses_before_uses}]:::HEADTAIL --> {shard_name};'
                         )
@@ -840,16 +846,17 @@ def _mermaid_str(self) -> List[str]:
                         )
             else:
                 # single shard case, no uses_before or uses_after_considered
-                name = list(self.pod_args['pods'].values())[0][0].name
-                uses = f'"{list(self.pod_args["pods"].values())[0][0].uses}"'
-                num_replicas = list(self.pod_args['pods'].values())[0][0].replicas
+                pod_args = list(self.pod_args['pods'].values())[0][0]
+                uses = f'"{replace_secret_of_hub_uri(pod_args.uses, secret)}"'
 
                 # just put the replicas in parallel
-                if num_replicas > 1:
-                    for rep_i in range(num_replicas):
-                        mermaid_graph.append(f'{name}/rep-{rep_i}[{uses}]:::pod;')
+                if pod_args.replicas > 1:
+                    for rep_i in range(pod_args.replicas):
+                        mermaid_graph.append(
+                            f'{pod_args.name}/rep-{rep_i}["{uses}"]:::pod;'
+                        )
                 else:
-                    mermaid_graph.append(f'{name}[{uses}]:::pod;')
+                    mermaid_graph.append(f'{pod_args.name}["{uses}"]:::pod;')
 
             mermaid_graph.append('end;')
         return mermaid_graph

tests/unit/hubble/test_helper.py

@@ -5,7 +5,6 @@
 
 from jina.hubble import helper
 from jina.hubble.helper import disk_cache_offline
-from jina.parsers.hubble import set_hub_pull_parser
 
 
 @pytest.fixture
@@ -44,6 +43,19 @@ def test_parse_wrong_hub_uri(uri_path):
     assert f'{uri_path} is not a valid Hub URI.' == str(info.value)
 
 
+def test_replace_secret_of_hub_uri():
+    result = helper.replace_secret_of_hub_uri('jinahub://hello', '_secret_')
+    assert result == 'jinahub://hello'
+
+    result = helper.replace_secret_of_hub_uri(
+        'jinahub://hello:dummy@secret/path', '*secret*'
+    )
+    assert result == 'jinahub://hello:*secret*/path'
+
+    result = helper.replace_secret_of_hub_uri('hello:magic/world')
+    assert result == 'hello:magic/world'
+
+
 def test_md5file(dummy_zip_file):
     md5sum = helper.md5file(dummy_zip_file)
     assert md5sum == '7ffd1501f24fe5a66dc45883550c2005'

tests/unit/orchestrate/deployments/test_deployments.py

@@ -1,6 +1,5 @@
 import json
 import os
-from multiprocessing import Process
 
 import pytest
 
@@ -116,13 +115,14 @@ def test_uses_before_after(pod_args, shards):
         assert pod.num_pods == 5 if shards == 2 else 1
 
 
-def test_mermaid_str_no_error(pod_args):
+def test_mermaid_str_no_secret(pod_args):
     pod_args.replicas = 3
-    pod_args.uses_before = 'MyDummyExecutor'
+    pod_args.shards = 3
+    pod_args.uses_before = 'jinahub+docker://MyDummyExecutor:Dummy@Secret'
     pod_args.uses_after = 'ChildDummyExecutor2'
-    pod_args.uses = 'ChildDummyExecutor'
+    pod_args.uses = 'jinahub://ChildDummyExecutor:Dummy@Secret'
     pod = Deployment(pod_args)
-    print(pod._mermaid_str)
+    assert 'Dummy@Secret' not in ''.join(pod._mermaid_str)
 
 
 @pytest.mark.slow