Add get_additional_upstream_nodes to FlytekitPlugin

flytekit/configuration/plugin.py

@@ -18,15 +18,20 @@
 ```
 """
 
-from typing import Optional, Protocol, runtime_checkable
+from contextvars import Context
+from typing import TYPE_CHECKING, List, Optional, Protocol, runtime_checkable
 
 from click import Group
 from importlib_metadata import entry_points
 
 from flytekit.configuration import Config, get_config_file
+from flytekit.core.node import Node
 from flytekit.loggers import logger
 from flytekit.remote import FlyteRemote
 
+if TYPE_CHECKING:
+    from flytekit.core.promise import SupportsNodeCreation
+
 
 @runtime_checkable
 class FlytekitPluginProtocol(Protocol):
@@ -90,6 +95,10 @@ def get_auth_success_html(endpoint: str) -> Optional[str]:
         """Get default success html. Return None to use flytekit's default success html."""
         return None
 
+    @staticmethod
+    def get_additional_upstream_nodes(ctx: Context, entity: "SupportsNodeCreation") -> List[Node]:
+        return []
+
 
 def _get_plugin_from_entrypoint():
     """Get plugin from entrypoint."""

flytekit/core/promise.py

@@ -1160,6 +1160,11 @@ def create_and_link_node(
     # These will be our core Nodes until we can amend the Promise to use NodeOutputs that reference our Nodes
     upstream_nodes = list(set([n for n in nodes if n.id != _common_constants.GLOBAL_INPUT_NODE_ID]))
 
+    from flytekit.configuration.plugin import get_plugin
+
+    additional_nodes = get_plugin().get_additional_upstream_nodes(ctx, entity)
+    upstream_nodes.extend(additional_nodes)
+
     flytekit_node = Node(
         # TODO: Better naming, probably a derivative of the function name.
         id=f"{ctx.compilation_state.prefix}n{len(ctx.compilation_state.nodes)}",