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

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
@@ -28,6 +28,9 @@ class ScheduledQuery:
         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.
 
@@ -159,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

@@ -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

chalk/sql/__init__.py

@@ -1437,6 +1437,8 @@ def DatabricksSource(
     access_token: str = ...,
     db: str = ...,
     port: str = ...,
+    client_id: str = ...,
+    client_secret: str = ...,
     engine_args: Optional[Dict[str, Any]] = ...,
 ) -> BaseSQLSourceProtocol:
     """You can also configure the integration directly using environment
@@ -1457,6 +1459,10 @@ def DatabricksSource(
         Database to use.
     port
         Port number to use.
+    client_id
+        OAuth service principal client ID (alternative to access_token).
+    client_secret
+        OAuth service principal client secret (alternative to access_token).
     engine_args
         Additional arguments to use when constructing the SQLAlchemy engine.
 
@@ -1475,6 +1481,14 @@ def DatabricksSource(
     ...     db=os.getenv("DATABRICKS_DATABASE"),
     ...     port=os.getenv("DATABRICKS_PORT"),
     ... )
+    >>> databricks_with_oauth = DatabricksSource(
+    ...     host=os.getenv("DATABRICKS_HOST"),
+    ...     http_path=os.getenv("DATABRICKS_HTTP_PATH"),
+    ...     client_id=os.getenv("DATABRICKS_CLIENT_ID"),
+    ...     client_secret=os.getenv("DATABRICKS_CLIENT_SECRET"),
+    ...     db=os.getenv("DATABRICKS_DATABASE"),
+    ...     port=os.getenv("DATABRICKS_PORT"),
+    ... )
     """
     ...
 
@@ -1487,6 +1501,8 @@ def DatabricksSource(
     access_token: Optional[str] = None,
     db: Optional[str] = None,
     port: Optional[Union[str, int]] = None,
+    client_id: Optional[str] = None,
+    client_secret: Optional[str] = None,
     engine_args: Optional[Dict[str, Any]] = None,
 ) -> BaseSQLSourceProtocol:
     """Create a Databricks data source. SQL-based data sources
@@ -1505,6 +1521,8 @@ def DatabricksSource(
         db=db,
         port=port,
         name=name,
+        client_id=client_id,
+        client_secret=client_secret,
         engine_args=engine_args,
     )
 

chalk/sql/_internal/integrations/databricks.py

@@ -33,6 +33,8 @@ _DATABRICKS_HTTP_PATH_NAME = "DATABRICKS_HTTP_PATH"
 _DATABRICKS_TOKEN_NAME = "DATABRICKS_TOKEN"
 _DATABRICKS_DATABASE_NAME = "DATABRICKS_DATABASE"
 _DATABRICKS_PORT_NAME = "DATABRICKS_PORT"
+_DATABRICKS_CLIENT_ID_NAME = "DATABRICKS_CLIENT_ID"
+_DATABRICKS_CLIENT_SECRET_NAME = "DATABRICKS_CLIENT_SECRET"
 
 
 class DatabricksSourceImpl(BaseSQLSource):
@@ -46,6 +48,8 @@ class DatabricksSourceImpl(BaseSQLSource):
         db: Optional[str] = None,
         port: Optional[Union[int, str]] = None,
         name: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
         engine_args: Optional[Dict[str, Any]] = None,
         executor: Optional[concurrent.futures.ThreadPoolExecutor] = None,
         integration_variable_override: Optional[Mapping[str, str]] = None,
@@ -74,21 +78,53 @@ class DatabricksSourceImpl(BaseSQLSource):
                 name=_DATABRICKS_PORT_NAME, integration_name=name, parser=int, override=integration_variable_override
             )
         )
+        self.client_id = client_id or load_integration_variable(
+            name=_DATABRICKS_CLIENT_ID_NAME, integration_name=name, override=integration_variable_override
+        )
+        self.client_secret = client_secret or load_integration_variable(
+            name=_DATABRICKS_CLIENT_SECRET_NAME, integration_name=name, override=integration_variable_override
+        )
         self.executor = executor or DEFAULT_IO_EXECUTOR
 
+        has_token = self.access_token is not None
+        has_oauth = self.client_id is not None and self.client_secret is not None
+
+        if has_token and has_oauth:
+            chalk_logger.warning(
+                "Both OAuth credentials and a personal access token were provided. Using OAuth authentication."
+            )
+            self.access_token = None
+
+        self._credentials_provider = None
+        if has_oauth:
+            try:
+                from databricks.sdk.core import Config, oauth_service_principal
+            except ImportError:
+                raise missing_dependency_exception("chalkpy[databricks]")
+
+            def credentials_provider():
+                config = Config(host=self.host, client_id=self.client_id, client_secret=self.client_secret)
+                return oauth_service_principal(config)
+
+            self._credentials_provider = credentials_provider
+
         if engine_args is None:
             engine_args = {}
+
+        connect_args: dict[str, Any] = {
+            "keepalives": 1,
+            "keepalives_idle": 30,
+            "keepalives_interval": 10,
+            "keepalives_count": 5,
+        }
+
+        if self._credentials_provider:
+            connect_args["credentials_provider"] = self._credentials_provider
+
         engine_args.setdefault("pool_size", 20)
         engine_args.setdefault("max_overflow", 60)
-        engine_args.setdefault(
-            "connect_args",
-            {
-                "keepalives": 1,
-                "keepalives_idle": 30,
-                "keepalives_interval": 10,
-                "keepalives_count": 5,
-            },
-        )
+        engine_args.setdefault("connect_args", connect_args)
+
         BaseSQLSource.__init__(self, name=name, engine_args=engine_args, async_engine_args={})
 
     def supports_inefficient_fallback(self) -> bool:
@@ -97,6 +133,12 @@ class DatabricksSourceImpl(BaseSQLSource):
     def get_sqlglot_dialect(self) -> str | None:
         return "databricks"
 
+    def _get_client_auth(self) -> Dict[str, str | Callable | None]:
+        if self._credentials_provider:
+            return {"credentials_provider": self._credentials_provider}
+        else:
+            return {"access_token": self.access_token}
+
     @contextlib.contextmanager
     def _create_temp_table(
         self,
@@ -168,10 +210,7 @@ class DatabricksSourceImpl(BaseSQLSource):
 
                     # Connect using databricks-sql-python for efficient Arrow fetching
                     with sql.connect(
-                        server_hostname=self.host,
-                        http_path=self.http_path,
-                        access_token=self.access_token,
-                        catalog=self.db,
+                        server_hostname=self.host, http_path=self.http_path, catalog=self.db, **self._get_client_auth()
                     ) as databricks_conn:
                         chalk_logger.info("Established connection with Databricks using databricks-sql-python")
 
@@ -273,10 +312,7 @@ class DatabricksSourceImpl(BaseSQLSource):
 
             # Connect using databricks-sql-python for efficient Arrow fetching
             with sql.connect(
-                server_hostname=self.host,
-                http_path=self.http_path,
-                access_token=self.access_token,
-                catalog=self.db,
+                server_hostname=self.host, http_path=self.http_path, catalog=self.db, **self._get_client_auth()
             ) as databricks_cnx:
                 with databricks_cnx.cursor() as cursor:
                     formatted_op, positional_params, named_params = self.compile_query(finalized_query)
@@ -353,6 +389,8 @@ class DatabricksSourceImpl(BaseSQLSource):
                 create_integration_variable(_DATABRICKS_TOKEN_NAME, self.name, self.access_token),
                 create_integration_variable(_DATABRICKS_DATABASE_NAME, self.name, self.db),
                 create_integration_variable(_DATABRICKS_PORT_NAME, self.name, self.port),
+                create_integration_variable(_DATABRICKS_CLIENT_ID_NAME, self.name, self.client_id),
+                create_integration_variable(_DATABRICKS_CLIENT_SECRET_NAME, self.name, self.client_secret),
             ]
             if v is not None
         }