Run ignore query

flytekit/clis/sdk_in_container/run.py

@@ -26,6 +26,7 @@
 from flytekit.configuration import DefaultImages, FastSerializationSettings, ImageConfig, SerializationSettings
 from flytekit.configuration.plugin import get_plugin
 from flytekit.core import context_manager
+from flytekit.core.artifact import ArtifactQuery
 from flytekit.core.base_task import PythonTask
 from flytekit.core.data_persistence import FileAccessProvider
 from flytekit.core.type_engine import TypeEngine
@@ -376,14 +377,17 @@ def to_click_option(
 
     description_extra = ""
     if literal_var.type.simple == SimpleType.STRUCT:
-        if default_val:
+        if default_val and not isinstance(default_val, ArtifactQuery):
             if type(default_val) == dict or type(default_val) == list:
                 default_val = json.dumps(default_val)
             else:
                 default_val = cast(DataClassJsonMixin, default_val).to_json()
         if literal_var.type.metadata:
             description_extra = f": {json.dumps(literal_var.type.metadata)}"
 
+    # If a query has been specified, the input is never strictly required at this layer
+    required = False if default_val and isinstance(default_val, ArtifactQuery) else required
+
     return click.Option(
         param_decls=[f"--{input_name}"],
         type=literal_converter.click_type,
@@ -508,7 +512,22 @@ def _run(*args, **kwargs):
         try:
             inputs = {}
             for input_name, _ in entity.python_interface.inputs.items():
-                inputs[input_name] = kwargs.get(input_name)
+                processed_click_value = kwargs.get(input_name)
+                if isinstance(processed_click_value, ArtifactQuery):
+                    if run_level_params.is_remote:
+                        click.secho(
+                            click.style(
+                                f"Input '{input_name}' not passed, supported backends will query"
+                                f" for {processed_click_value.get_str(**kwargs)}",
+                                bold=True,
+                            )
+                        )
+                        continue
+                    else:
+                        raise click.UsageError(
+                            f"Default for '{input_name}' is a query, which must be specified when running locally."
+                        )
+                inputs[input_name] = processed_click_value
 
             if not run_level_params.is_remote:
                 with FlyteContextManager.with_context(_update_flyte_context(run_level_params)):

flytekit/core/artifact.py

@@ -186,6 +186,47 @@
     ) -> art_id.ArtifactQuery:
         return Serializer.artifact_query_to_idl(self, **kwargs)
 
+    def get_time_partition_str(self, **kwargs) -> str:
+        tp_str = ""
+        if self.time_partition:
+            tp = self.time_partition.value
+            if tp.HasField("time_value"):
+                tp = tp.time_value.ToDatetime()
+                tp_str += f" Time partition: {tp}"
+            elif tp.HasField("input_binding"):
+                var = tp.input_binding.var
+                if var not in kwargs:
+                    raise ValueError(f"Time partition input binding {var} not found in kwargs")
+                else:
+                    tp_str += f" Time partition from input<{var}>,"
+        return tp_str
+
+    def get_partition_str(self, **kwargs) -> str:
+        p_str = ""
+        if self.partitions and self.partitions.partitions and len(self.partitions.partitions) > 0:
+            p_str = " Partitions: "
+            for k, v in self.partitions.partitions.items():
+                if v.value and v.value.HasField("static_value"):
+                    p_str += f"{k}={v.value.static_value}, "
+                elif v.value and v.value.HasField("input_binding"):
+                    var = v.value.input_binding.var
+                    if var not in kwargs:
+                        raise ValueError(f"Partition input binding {var} not found in kwargs")
+                    else:
+                        p_str += f"{k} from input<{var}>, "
+        return p_str.rstrip("\n\r, ")
+
+    def get_str(self, **kwargs):
+        # Detailed string that explains query a bit more, used in running
+        tp_str = self.get_time_partition_str(**kwargs)
+        p_str = self.get_partition_str(**kwargs)
+
+        return f"'{self.artifact.name}'...{tp_str}{p_str}"
+
+    def __str__(self):
+        # Default string used for printing --help
+        return f"Artifact Query: on {self.artifact.name}"
+
 
 class TimePartition(object):
     def __init__(

flytekit/interaction/click_types.py

@@ -14,6 +14,7 @@
 from pytimeparse import parse
 
 from flytekit import BlobType, FlyteContext, FlyteContextManager, Literal, LiteralType, StructuredDataset
+from flytekit.core.artifact import ArtifactQuery
 from flytekit.core.data_persistence import FileAccessProvider
 from flytekit.core.type_engine import TypeEngine
 from flytekit.models.types import SimpleType
@@ -61,6 +62,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         p = pathlib.Path(value)
         # set remote_directory to false if running pyflyte run locally. This makes sure that the original
         # directory is used and not a random one.
@@ -80,6 +83,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         if isinstance(value, str):
             return StructuredDataset(uri=value)
         elif isinstance(value, StructuredDataset):
@@ -93,6 +98,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         # set remote_directory to false if running pyflyte run locally. This makes sure that the original
         # file is used and not a random one.
         remote_path = None if getattr(ctx.obj, "is_remote", False) else False
@@ -109,6 +116,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         # set remote_directory to false if running pyflyte run locally. This makes sure that the original
         # file is used and not a random one.
         remote_path = None if getattr(ctx.obj, "is_remote", None) else False
@@ -131,6 +140,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         if value in self._ADDITONAL_FORMATS:
             if value == self._NOW_FMT:
                 return datetime.datetime.now()
@@ -143,6 +154,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Duration type.")
         return datetime.timedelta(seconds=parse(value))
@@ -156,6 +169,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> enum.Enum:
+        if isinstance(value, ArtifactQuery):
+            return value
         if isinstance(value, self._enum_type):
             return value
         return self._enum_type(super().convert(value, param, ctx))
@@ -191,6 +206,8 @@
         Important to implement NoneType / Optional.
         Also could we just determine the click types from the python types
         """
+        if isinstance(value, ArtifactQuery):
+            return value
         for t in self._types:
             try:
                 return t.convert(value, param, ctx)
@@ -228,6 +245,8 @@
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
+        if isinstance(value, ArtifactQuery):
+            return value
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Json type.")
 
@@ -274,7 +293,7 @@
     SimpleType.STRING: click.STRING,
     SimpleType.BOOLEAN: click.BOOL,
     SimpleType.DURATION: DurationParamType(),
-    SimpleType.DATETIME: click.DateTime(),
+    SimpleType.DATETIME: DateTimeType(),
 }
 
 
@@ -353,6 +372,8 @@
         """
         Convert the value to a Flyte Literal or a python native type. This is used by click to convert the input.
         """
+        if isinstance(value, ArtifactQuery):
+            return value
         try:
             # If the expected Python type is datetime.date, adjust the value to date
             if self._python_type is datetime.date:

tests/flytekit/unit/interaction/test_click_types.py

@@ -10,13 +10,19 @@
 import yaml
 
 from flytekit import FlyteContextManager
+from flytekit.core.artifact import Artifact
 from flytekit.core.type_engine import TypeEngine
 from flytekit.interaction.click_types import (
     DateTimeType,
+    DirParamType,
     DurationParamType,
+    EnumParamType,
     FileParamType,
     FlyteLiteralConverter,
     JsonParamType,
+    PickleParamType,
+    StructuredDatasetParamType,
+    UnionParamType,
     key_value_callback,
 )
 
@@ -62,6 +68,23 @@ def test_literal_converter(python_type, python_value):
     assert lc.convert(click_ctx, dummy_param, python_value) == TypeEngine.to_literal(ctx, python_value, python_type, lt)
 
 
+def test_literal_converter_query():
+    ctx = FlyteContextManager.current_context()
+    lt = TypeEngine.to_literal_type(int)
+
+    lc = FlyteLiteralConverter(
+        ctx,
+        literal_type=lt,
+        python_type=int,
+        is_remote=True,
+    )
+
+    a = Artifact(name="test-artifact")
+    query = a.query()
+    click_ctx = click.Context(click.Command("test_command"), obj={"remote": True})
+    assert lc.convert(click_ctx, dummy_param, query) == query
+
+
 @pytest.mark.parametrize(
     "python_type, python_str_value, python_value",
     [
@@ -163,3 +186,23 @@ def test_key_value_callback():
         key_value_callback(ctx, "a", ["a=b", "c=d", "e"])
     with pytest.raises(click.BadParameter):
         key_value_callback(ctx, "a", ["a=b", "c=d", "e=f", "g"])
+
+
+@pytest.mark.parametrize(
+    "param_type",
+    [
+        (DateTimeType()),
+        (DurationParamType()),
+        (JsonParamType(typing.Dict[str, str])),
+        (UnionParamType([click.FLOAT, click.INT])),
+        (EnumParamType(Color)),
+        (PickleParamType()),
+        (StructuredDatasetParamType()),
+        (DirParamType()),
+    ],
+)
+def test_query_passing(param_type: click.ParamType):
+    a = Artifact(name="test-artifact")
+    query = a.query()
+
+    assert param_type.convert(value=query, param=None, ctx=None) is query