Merge branch 'master' into feat/databricks-token-support

Resolve merge conflicts in task.py after #3392 (serverless) was merged:
- Combined DatabricksV2 attributes: kept all serverless fields, added
  databricks_token_secret
- Combined get_custom() serialization for both feature sets
- Added auth_token to serverless test metadata assertion
- Removed emoji from error message

Signed-off-by: Rohit Sharma <rohitrsh@gmail.com>
Made-with: Cursor

Author: rohitrsh

Dockerfile.connector

@@ -7,8 +7,15 @@ ARG VERSION
 
 RUN apt-get update && apt-get install build-essential -y \
     && pip install uv
-
-RUN uv pip install --system --no-cache-dir -U flytekit[connector]==$VERSION \
+# Pin pendulum<3.0: Apache Airflow (via flytekitplugins-airflow) imports
+# pendulum.tz.timezone() at module load time (airflow/settings.py).
+# Pendulum 3.x changed the tz API, causing the connector to crash on startup:
+#   airflow/settings.py → TIMEZONE = pendulum.tz.timezone("UTC") → AttributeError
+# Without this pin, uv resolves to pendulum 3.x which breaks the import chain:
+#   pyflyte serve connector → load_implicit_plugins → airflow → pendulum → crash
+RUN uv pip install --system --no-cache-dir -U \
+  "pendulum>=2.0.0,<3.0" \
+  flytekit[connector]==$VERSION \
   flytekitplugins-airflow==$VERSION \
   flytekitplugins-bigquery==$VERSION \
   flytekitplugins-k8sdataservice==$VERSION \

flytekit/configuration/__init__.py

@@ -545,6 +545,7 @@ class S3Config(object):
     backoff: datetime.timedelta = datetime.timedelta(seconds=5)
     access_key_id: typing.Optional[str] = None
     secret_access_key: typing.Optional[str] = None
+    adressing_style: typing.Optional[str] = None
 
     @classmethod
     def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> S3Config:
@@ -561,6 +562,7 @@ def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> S3Config:
         kwargs = set_if_exists(kwargs, "backoff", _internal.AWS.BACKOFF_SECONDS.read(config_file))
         kwargs = set_if_exists(kwargs, "access_key_id", _internal.AWS.S3_ACCESS_KEY_ID.read(config_file))
         kwargs = set_if_exists(kwargs, "secret_access_key", _internal.AWS.S3_SECRET_ACCESS_KEY.read(config_file))
+        kwargs = set_if_exists(kwargs, "adressing_style", _internal.AWS.S3_ADDRESSING_STYLE.read(config_file))
         return S3Config(**kwargs)
 
 

flytekit/configuration/internal.py

@@ -55,6 +55,10 @@ class AWS(object):
         LegacyConfigEntry(SECTION, "backoff_seconds", datetime.timedelta),
         transform=lambda x: datetime.timedelta(seconds=int(x)),
     )
+    S3_ADDRESSING_STYLE = ConfigEntry(
+        LegacyConfigEntry(SECTION, "addressing_style"),
+        YamlConfigEntry("storage.connection.addressing-style"),
+    )
 
 
 class GCP(object):

flytekit/core/data_persistence.py

@@ -57,6 +57,9 @@ def s3_setup_args(s3_cfg: configuration.S3Config, anonymous: bool = False) -> Di
     if s3_cfg.endpoint is not None:
         kwargs["client_kwargs"] = {"endpoint_url": s3_cfg.endpoint}
 
+    if s3_cfg.adressing_style:
+        kwargs["config_kwargs"] = {"s3": {"addressing_style": "virtual"}}
+
     if anonymous:
         kwargs[_ANON] = True
 

flytekit/core/type_engine.py

@@ -225,6 +225,29 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
             f"Conversion to python value expected type {expected_python_type} from literal not implemented"
         )
 
+    def schema_match(self, schema: dict) -> bool:
+        """Check if a JSON schema fragment matches this transformer's python_type.
+
+        For BaseModel subclasses, automatically compares the schema's title, type, and
+        required fields against the type's own JSON schema. For other types, returns
+        False by default — override if needed.
+        """
+        if not isinstance(schema, dict):
+            return False
+        try:
+            from pydantic import BaseModel
+
+            if hasattr(self.python_type, "model_json_schema") and self.python_type is not BaseModel:
+                this_schema = self.python_type.model_json_schema()  # type: ignore[attr-defined]
+                return (
+                    schema.get("title") == this_schema.get("title")
+                    and schema.get("type") == this_schema.get("type")
+                    and set(schema.get("required", [])) == set(this_schema.get("required", []))
+                )
+        except Exception:
+            pass
+        return False
+
     def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> Optional[T]:
         """
         This function primarily handles deserialization for untyped dicts, dataclasses, Pydantic BaseModels, and attribute access.｀
@@ -1180,6 +1203,12 @@ def _handle_json_schema_property(
         attr_union_type = reduce(lambda x, y: typing.Union[x, y], attr_types)
         return (property_key, attr_union_type)  # type: ignore
 
+    # Handle $ref (e.g. when anyOf recurses with a $ref variant like Optional[Inner])
+    # v1 iterates all properties (not just required), so $ref inside anyOf must be resolved here
+    if property_val.get("$ref"):
+        resolved_type = _get_element_type(property_val, schema)
+        return (property_key, resolved_type)
+
     # Handle enum
     if property_val.get("enum"):
         property_type = "enum"
@@ -1203,6 +1232,9 @@ def _handle_json_schema_property(
         elif property_val.get("title"):
             # For nested dataclass
             sub_schema_name = property_val["title"]
+            matched_type = _match_registered_type_from_schema(property_val)
+            if matched_type is not None:
+                return (property_key, typing.cast(GenericAlias, matched_type))
             return (
                 property_key,
                 typing.cast(GenericAlias, convert_mashumaro_json_schema_to_python_class(property_val, sub_schema_name)),
@@ -1217,6 +1249,14 @@ def _handle_json_schema_property(
         return (property_key, _get_element_type(property_val, schema))  # type: ignore
 
 
+def _match_registered_type_from_schema(schema: dict) -> typing.Optional[type]:
+    """Check if a JSON schema fragment matches any registered TypeTransformer."""
+    for transformer in TypeEngine._REGISTRY.values():  # type: ignore[misc]
+        if transformer.schema_match(schema):
+            return transformer.python_type
+    return None
+
+
 def generate_attribute_list_from_dataclass_json_mixin(
     schema: typing.Dict[str, typing.Any],
     schema_name: typing.Any,
@@ -1233,6 +1273,15 @@ def generate_attribute_list_from_dataclass_json_mixin(
             defs = schema.get("$defs", schema.get("definitions", {}))
             if ref_name in defs:
                 ref_schema = defs[ref_name].copy()
+                # Check if the $ref points to an enum definition (no properties)
+                if ref_schema.get("enum"):
+                    attribute_list.append((property_key, str))
+                    continue
+                # Check if the $ref matches a registered custom type
+                matched_type = _match_registered_type_from_schema(ref_schema)
+                if matched_type is not None:
+                    attribute_list.append((property_key, typing.cast(GenericAlias, matched_type)))
+                    continue
                 # Include $defs so nested models can resolve their own $refs
                 if "$defs" not in ref_schema and defs:
                     ref_schema["$defs"] = defs
@@ -2531,32 +2580,53 @@ def convert_mashumaro_json_schema_to_python_class(schema: dict, schema_name: typ
 def _get_element_type(
     element_property: typing.Dict[str, typing.Any], schema: typing.Optional[typing.Dict[str, typing.Any]] = None
 ) -> Type:
-    # Handle $ref for nested models in arrays (e.g., List[NestedModel])
-    # Pydantic generates JSON schema like: {'items': {'$ref': '#/$defs/NestedModel'}, 'type': 'array'}
-    if element_property.get("$ref"):
-        if schema is None:
-            raise ValueError(f"Cannot resolve $ref '{element_property['$ref']}' without schema context")
-        ref_path = element_property["$ref"]
-        ref_name = ref_path.split("/")[-1]
+    # Handle $ref for nested models and enums
+    # Ensure that the element is actually a $ref and we have the entire schema to look up
+    if element_property.get("$ref") and schema is not None:
+        ref_name = element_property["$ref"].split("/")[-1]
         defs = schema.get("$defs", schema.get("definitions", {}))
+        # Look up for ref_name in the defs defined in the schema
         if ref_name in defs:
+            # Don't mutate the original schema
             ref_schema = defs[ref_name].copy()
-            # Include $defs so nested models can resolve their own $refs
+            # Guard the nested enum elements inside containers
+            if ref_schema.get("enum"):
+                return str
+            # Check if the $ref matches a registered custom type
+            if (matched_type := _match_registered_type_from_schema(ref_schema)) is not None:
+                return matched_type
+            # if defs not in the schema, they need to be propagated into the resolved schema
             if "$defs" not in ref_schema and defs:
                 ref_schema["$defs"] = defs
+            # build a dataclass from the resolved schema
             return convert_mashumaro_json_schema_to_python_class(ref_schema, ref_name)
-        else:
-            raise ValueError(f"Cannot find definition for $ref '{ref_path}' in schema")
-
-    element_type = (
-        [e_property["type"] for e_property in element_property["anyOf"]]  # type: ignore
-        if element_property.get("anyOf")
-        else element_property["type"]
-    )
-    element_format = element_property["format"] if "format" in element_property else None
+        # default to str on failure. Shouldn't happen with valid pydantic schemas
+        return str
 
+    # Handle anyOf (e.g. Optional[int], Optional[Inner])
+    # Early return block replacing the previous list comprehension which would fail when an anyOf reference was a $ref
+    # (meaning no $type key).
+    if element_property.get("anyOf"):
+        # Separate non null variants. Note a $ref variant would have type None NOT null. A {"type": "null"} variant is
+        # filtered out.
+        variants = element_property["anyOf"]
+        non_null = [v for v in variants if v.get("type") != "null"]
+        # Detect if this is an Optional pattern here
+        has_null = len(non_null) < len(variants)
+        # This recurses on the first non-null variant which would handle the $ref, nested_arrays, nested_objects...
+        # anything. Wrap it in Optional if has_null.
+        if non_null:
+            inner_type = _get_element_type(non_null[0], schema)
+            return typing.Optional[inner_type] if has_null else inner_type  # type: ignore
+        # return None if all types are None
+        return type(None)
+
+    element_type = element_property.get("type", "string")
+    element_format = element_property.get("format")
+
+    # Handle marshmallow-style Optional types where "type" is a list e.g. ["integer", "null"]
+    # This pattern is not used by Pydantic (which uses anyOf) but is used by the marshmallow/DataClassJsonMixin path in v1
     if isinstance(element_type, list):
-        # Element type of Optional[int] is [integer, None]
         return typing.Optional[_get_element_type({"type": element_type[0]}, schema)]  # type: ignore
 
     if element_type == "string":
@@ -2570,6 +2640,16 @@ def _get_element_type(
             return int
         else:
             return float
+    # Recursively discover the types when an array or object element type is discovered
+    elif element_type == "array":
+        return typing.List[_get_element_type(element_property.get("items", {}), schema)]  # type: ignore
+    elif element_type == "object":
+        if element_property.get("additionalProperties"):
+            return typing.Dict[str, _get_element_type(element_property["additionalProperties"], schema)]  # type: ignore
+        return dict
+    # Corner case - practically useless but List[None] is a legal Python type
+    elif element_type == "null":
+        return type(None)
     return str