chalkpy 2.89.22__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/parsed/user_types_to_json.py

@@ -110,6 +110,22 @@ def project_settings_to_gql(config: ProjectSettings) -> ProjectSettingsGQL:
 
 
 def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGraphGQL:
+    # Validate registries BEFORE conversion to catch errors early
+    # This ensures parity with Proto validation path
+    from chalk.parsed.validation_from_registries import validate_all_from_registries
+
+    try:
+        validate_all_from_registries(
+            features_registry=FeatureSetBase.registry,
+            resolver_registry=RESOLVER_REGISTRY,
+        )
+    except Exception as e:
+        # If validation fails, add to failed but continue
+        # to allow other validation to complete
+        if not LSPErrorBuilder.promote_exception(e):
+            # Not an LSP error, so log it as a failed import
+            failed.append(build_failed_import(e, "validation"))
+
     features = []
     feature_classes: list[FeatureClassGQL] = []
     for x in FeatureSetBase.registry.values():
@@ -173,6 +189,14 @@ def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGr
     named_queries: list[UpsertNamedQueryGQL] = []
     for named_query in NAMED_QUERY_REGISTRY.values():
         if named_query.filename is None or _is_relative_to(Path(named_query.filename), scope_to):
+            # only try to convert if there are no errors
+            if not named_query.errors:
+                try:
+                    named_queries.append(convert_type_to_gql(named_query, path_prefix=path_prefix_to_remove))
+                except Exception as e:
+                    failed.append(build_failed_import(e, f"named query '{named_query.name}'"))
+
+            # named_query.errors can go from empty to non-empty after conversion
             if named_query.errors:
                 code_object_diagnostics[named_query.filename or ""].extend(
                     [
@@ -185,15 +209,17 @@ def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGr
                         for error in named_query.errors
                     ]
                 )
-                continue
-            try:
-                named_queries.append(convert_type_to_gql(named_query, path_prefix=path_prefix_to_remove))
-            except Exception as e:
-                failed.append(build_failed_import(e, f"named query '{named_query.name}'"))
 
     model_references: list[UpsertModelReferenceGQL] = []
     for mr in MODEL_REFERENCE_REGISTRY.values():
         if mr.filename is None or _is_relative_to(Path(mr.filename), scope_to):
+            # similar logic to named_queries above
+            if not mr.errors:
+                try:
+                    model_references.append(convert_type_to_gql(mr, path_prefix=path_prefix_to_remove))
+                except Exception as e:
+                    failed.append(build_failed_import(e, f"model reference '{mr.name}'"))
+
             if mr.errors:
                 code_object_diagnostics[mr.filename or ""].extend(
                     [
@@ -206,11 +232,6 @@ def get_registered_types(scope_to: Path, failed: List[FailedImport]) -> UpsertGr
                         for error in mr.errors
                     ]
                 )
-                continue
-            try:
-                model_references.append(convert_type_to_gql(mr, path_prefix=path_prefix_to_remove))
-            except Exception as e:
-                failed.append(build_failed_import(e, f"model reference '{mr.name}'"))
 
     # online store configs
     for osc in ONLINE_STORE_CONFIG_REGISTRY.values():

chalk/parsed/validation_from_registries.py

@@ -0,0 +1,182 @@
+"""
+Unified validation layer that operates directly on registries.
+
+This module provides validation that works for BOTH GQL and Proto conversion paths.
+It triggers lazy validations by accessing properties on registry objects and performs
+explicit validation checks.
+
+This ensures validation parity between:
+- GQL path: get_registered_types() → validate_graph()
+- Proto path: ToProtoConverter.convert_graph()
+
+By calling validate_all_from_registries() before conversion in both paths, we ensure
+developers cannot add validation to one path and forget the other.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from chalk._lsp.error_builder import LSPErrorBuilder
+
+if TYPE_CHECKING:
+    from chalk.features.feature_field import Feature
+    from chalk.features.feature_set import Features
+    from chalk.features.resolver import ResolverRegistry
+
+
+def validate_all_from_registries(
+    features_registry: dict[str, type["Features"]],
+    resolver_registry: "ResolverRegistry",
+) -> None:
+    """
+    Trigger all validations by accessing properties on registry objects.
+    This can be run multiple times and not show duplicates.
+
+    This function should be called by BOTH GQL and Proto conversion paths BEFORE
+    they perform their conversions. It validates by triggering lazy validations:
+
+    - Error[24]: Feature names with protected prefixes
+    - Error[25]: Namespace names with protected prefixes
+    - Error[32]: Invalid join syntax (composite joins must use & not and)
+    - Error[35]: Missing has-one join definition
+    - Error[37]: Join filter with incorrect type annotation
+    - Error[40]: Invalid join lambda
+    - Error[42]: Bad foreign key types (type mismatch)
+    - Error[43]: Multi-namespace joins
+    - Error[51]: Multiple primary features (versioned primary keys)
+    - Error[119]: State type validation (stream resolvers)
+    - Error[135]: Unrecognized feature reference
+
+    Parameters
+    ----------
+    features_registry : dict[str, type[Features]]
+        The feature registry to validate (FeatureSetBase.registry).
+    resolver_registry : ResolverRegistry
+        The resolver registry to validate (RESOLVER_REGISTRY).
+
+    Returns
+    -------
+    None
+        Validation errors are accumulated in LSPErrorBuilder and raised as exceptions.
+    """
+
+    # ========================================================================
+    # FEATURE VALIDATION
+    # ========================================================================
+
+    for _, features_cls in features_registry.items():
+        # --------------------------------------------------------------------
+        # Error[51]: Multiple primary features (versioned primary keys)
+        # --------------------------------------------------------------------
+        # Accessing __chalk_primary__ triggers _discover_feature() which
+        # validates that there's only one primary key. Versioned features
+        # create multiple primary keys (e.g., id, id@2, id@3, id@4) which
+        # triggers Error[51].
+        try:
+            _ = features_cls.__chalk_primary__
+        except Exception as e:
+            # LSPErrorBuilder.promote_exception() re-raises LSP errors
+            if not LSPErrorBuilder.promote_exception(e):
+                # If it's not an LSP error, something else went wrong
+                raise
+
+        # --------------------------------------------------------------------
+        # Iterate through all features in this feature set
+        # --------------------------------------------------------------------
+        for feature in features_cls.features:
+            # Skip autogenerated and no-display features (same as user_types_to_json.py:138)
+            # This prevents validating internal features like __chalk_* that are allowed
+            # to have protected names
+            if feature.is_autogenerated or feature.no_display:
+                continue
+
+            # ----------------------------------------------------------------
+            # Error[32,35,37,40,42,43]: Join validation
+            # ----------------------------------------------------------------
+            # Accessing the .join property triggers:
+            # - _validate_join() in feature_field.py (Error[32,37])
+            # - _validate_filter() in feature_field.py (Error[40,42,43])
+            #
+            # During GQL conversion, convert_type_to_gql() also checks:
+            # - Error[35]: if t.is_has_one and t.join is None
+            try:
+                _ = feature.join
+            except Exception as e:
+                if not LSPErrorBuilder.promote_exception(e):
+                    raise
+
+            # ----------------------------------------------------------------
+            # Error[24,25]: Feature and namespace name validation
+            # ----------------------------------------------------------------
+            try:
+                _validate_feature_names_from_registry(feature)
+            except Exception as e:
+                if not LSPErrorBuilder.promote_exception(e):
+                    raise
+
+    # ========================================================================
+    # RESOLVER VALIDATION
+    # ========================================================================
+
+    for resolver in resolver_registry.get_all_resolvers():
+        # --------------------------------------------------------------------
+        # Error[135]: Unrecognized feature reference
+        # --------------------------------------------------------------------
+        # Accessing resolver.inputs triggers _do_parse() which validates
+        # that all input features are recognized and exist in the registry.
+        try:
+            _ = resolver.inputs
+        except Exception as e:
+            if not LSPErrorBuilder.promote_exception(e):
+                raise
+
+        # --------------------------------------------------------------------
+        # Error[119]: State type validation (stream resolvers)
+        # --------------------------------------------------------------------
+        # Accessing resolver state and default_args triggers validation
+        # that default state values match their type annotations.
+        try:
+            _ = resolver.state
+            _ = resolver.default_args
+        except Exception as e:
+            if not LSPErrorBuilder.promote_exception(e):
+                raise
+
+
+def _validate_feature_names_from_registry(feature: "Feature") -> None:
+    """
+    Validate that feature names and namespace names don't use protected prefixes.
+
+    This performs the same validation as _validate_feature_names() in
+    _graph_validation.py, but operates on Feature objects from the registry
+    rather than UpsertFeatureGQL objects.
+
+    Parameters
+    ----------
+    feature : Feature
+        The feature to validate from FeatureSetBase.registry
+
+    Raises
+    ------
+    Exception
+        If feature or namespace name starts with '_chalk' or '__'
+    """
+    # Error[24]: Feature names cannot begin with '_chalk' or '__'
+    if feature.name.startswith("__") or feature.name.startswith("_chalk"):
+        feature.lsp_error_builder.add_diagnostic(
+            message="Feature names cannot begin with '_chalk' or '__'.",
+            range=feature.lsp_error_builder.property_range(feature.attribute_name or feature.name),
+            label="protected name",
+            code="24",
+        )
+
+    # Error[25]: Namespace names cannot begin with '_chalk' or '__'
+    if feature.namespace.startswith("__") or feature.namespace.startswith("_chalk"):
+        feature.lsp_error_builder.add_diagnostic(
+            message="Feature classes cannot have names that begin with '_chalk' or '__'.",
+            label="protected namespace",
+            range=feature.lsp_error_builder.decorator_kwarg_value_range("name")
+            or feature.lsp_error_builder.class_definition_range(),
+            code="25",
+        )

chalk/queries/named_query.py

@@ -4,6 +4,7 @@ import inspect
 import traceback
 from typing import TYPE_CHECKING, Mapping, Sequence
 
+from chalk._lsp.error_builder import LSPErrorBuilder
 from chalk.features import unwrap_feature
 from chalk.utils.object_inspect import get_source_object_starting
 from chalk.utils.source_parsing import should_skip_source_code_parsing
@@ -157,9 +158,12 @@ class NamedQuery:
                 self._input = [str(f) for f in self._input_raw]
             elif self._output_raw is not None:
                 self._input = [str(unwrap_feature(o).primary_feature) for o in self._output_raw]
-        except Exception:
+        except Exception as e:
             self._input = None
-            self.errors.append(traceback.format_exc())
+            if not LSPErrorBuilder.promote_exception(e):
+                self.errors.append(
+                    f"Error creating NamedQuery '{self.name} ({self.version})': {traceback.format_exc()}"
+                )
 
         return self._input
 
@@ -170,9 +174,12 @@ class NamedQuery:
         try:
             if self._output_raw is not None:
                 self._output = [str(o) for o in self._output_raw]
-        except Exception:
+        except Exception as e:
             self._output = None
-            self.errors.append(f"Error creating NamedQuery '{self.name} ({self.version})': {traceback.format_exc()}")
+            if not LSPErrorBuilder.promote_exception(e):
+                self.errors.append(
+                    f"Error creating NamedQuery '{self.name} ({self.version})': {traceback.format_exc()}"
+                )
 
         return self._output
 
@@ -183,9 +190,12 @@ class NamedQuery:
         try:
             if self._additional_logged_features_raw is not None:
                 self._additional_logged_features = [str(alf) for alf in self._additional_logged_features_raw]
-        except Exception:
+        except Exception as e:
             self._additional_logged_features = None
-            self.errors.append(f"Error creating NamedQuery '{self.name} ({self.version})': {traceback.format_exc()}")
+            if not LSPErrorBuilder.promote_exception(e):
+                self.errors.append(
+                    f"Error creating NamedQuery '{self.name} ({self.version})': {traceback.format_exc()}"
+                )
 
         return self._additional_logged_features
 

chalk/queries/scheduled_query.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Collection
 
 from chalk.utils.duration import CronTab, Duration
@@ -21,12 +21,16 @@ class ScheduledQuery:
         lower_bound: datetime | None = None,
         upper_bound: datetime | None = None,
         tags: Collection[str] | None = None,
+        dataset_name: str | None = None,
         required_resolver_tags: Collection[str] | None = None,
         store_online: bool = True,
         store_offline: bool = True,
         incremental_resolvers: Collection[str] | None = None,
         planner_options: dict[str, str] | None = None,
         resource_group: str | None = None,
+        completion_deadline: timedelta | None = None,
+        num_shards: int | None = None,
+        num_workers: int | None = None,
     ):
         """Create an offline query which runs on a schedule.
 
@@ -64,6 +68,8 @@ class ScheduledQuery:
             incrementalization.
         tags
             Allows selecting resolvers with these tags.
+        dataset_name
+            Associated dataset name for the scheduled query.
         required_resolver_tags
             Requires that resolvers have these tags.
         store_online
@@ -143,6 +149,7 @@ class ScheduledQuery:
         self.lower_bound = lower_bound
         self.upper_bound = upper_bound
         self.tags = tags
+        self.dataset_name = dataset_name
         self.required_resolver_tags = required_resolver_tags
         self.filename = caller_filename
         self.store_online = store_online
@@ -155,6 +162,11 @@ class ScheduledQuery:
         self.planner_options = {k: str(v) for k, v in planner_options.items()} if planner_options else None
         self.resource_group = resource_group
 
+        self.completion_deadline = completion_deadline
+
+        self.num_shards = num_shards
+        self.num_workers = num_workers
+
         CRON_QUERY_REGISTRY[name] = self
 
 

chalk/serialization/parsed_annotation.py

@@ -111,7 +111,7 @@ class ParsedAnnotation:
             typ = self._features_cls.__annotations__[self._attribute_name]
         else:
             typ = self._parsed_annotation or self._underlying or self._unparsed_underlying
-        if sys.version_info >= (3, 9) and get_origin(typ) == typing.Annotated:
+        if get_origin(typ) == typing.Annotated:
             typ = typ.__origin__  # pyright: ignore
         if isinstance(typ, type):
             typ = typ.__name__
@@ -350,6 +350,7 @@ class ParsedAnnotation:
             if "__chalk_document__" in args:
                 self._is_document = True
             origin = get_origin(annotation)
+            self._parsed_annotation = cast(type, annotation)
 
         # The only allowed collections here are Set, List, or DataFrame
         if origin in (set, Set):
@@ -372,17 +373,6 @@ class ParsedAnnotation:
             annotation = args[0]
         if origin in (tuple, Tuple):
             args = get_args(annotation)
-            if len(args) != 2 or args[1] is not ... or args[0] is ...:
-                self._type_error(
-                    (
-                        "Tuple should be given exactly two type parameters. "
-                        "The first should be the type of the elements, and the second should be '...', "
-                        "which indicates that the tuple is of variable length. "
-                        "For example, 'Tuple[int, ...]' is a tuple of ints of variable length."
-                    ),
-                    code="74",
-                    label="invalid tuple",
-                )
             annotation = args[0]
         if origin in (list, List):
             args = get_args(annotation)
@@ -567,3 +557,26 @@ class ParsedAnnotation:
         if self._parsed_annotation is None:
             self._parse_annotation()
         return self._is_feature_time
+
+    def is_dataframe_annotation(self) -> bool:
+        """
+        Check if the annotation represents a DataFrame type, even if validation failed.
+        This checks the raw parsed annotation without triggering full validation,
+        useful for preventing false positive errors when DataFrame validation fails.
+        """
+        from typing import get_args
+
+        from chalk.features.dataframe import DataFrameMeta
+
+        if self._parsed_annotation is None:
+            self._parse_annotation()
+
+        # Check if directly a DataFrame
+        if isinstance(self.parsed_annotation, DataFrameMeta):
+            return True
+
+        # Check if wrapped in Optional, Union, etc.
+        if any(isinstance(x, DataFrameMeta) for x in get_args(self.parsed_annotation)):
+            return True
+
+        return False