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

chalk/sql/_internal/integrations/mssql.py

@@ -1,15 +1,20 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union
+import struct
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, Optional, Union
 
 from chalk.integrations.named import create_integration_variable, load_integration_variable
+from chalk.sql._internal.query_execution_parameters import QueryExecutionParameters
 from chalk.sql._internal.sql_source import BaseSQLSource, SQLSourceKind, TableIngestMixIn
+from chalk.sql.finalized_query import FinalizedChalkQuery
 from chalk.sql.protocols import SQLSourceWithTableIngestProtocol
+from chalk.utils.environment_parsing import env_var_bool
 from chalk.utils.missing_dependency import missing_dependency_exception
 
 if TYPE_CHECKING:
-    from sqlalchemy.engine import URL
+    import pyarrow as pa
+    from sqlalchemy.engine import URL, Connection
 
 _MSSQL_HOST_NAME = "MSSQL_HOST"
 _MSSQL_TCP_PORT_NAME = "MSSQL_TCP_PORT"
@@ -40,10 +45,10 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
         integration_variable_override: Optional[Mapping[str, str]] = None,
     ):
         try:
-            import pymssql
-        except ImportError:
-            raise missing_dependency_exception("chalkpy[mssql]")
-        del pymssql
+            import pyodbc
+        except ImportError as e:
+            raise missing_dependency_exception("chalkpy[mssql]", original_error=e)
+        del pyodbc
 
         self.name = name
         self.host = host or load_integration_variable(
@@ -112,16 +117,46 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
 
         BaseSQLSource.__init__(self, name=name, engine_args=engine_args, async_engine_args=async_engine_args)
 
+        # Register event listener for managed identity token injection
+        if not self.client_id and not self.user:
+            from sqlalchemy import event
+
+            event.listens_for(self.get_engine(), "do_connect")(self._inject_azure_token)
+
+    def _inject_azure_token(self, _dialect: Any, _conn_rec: Any, _cargs: Any, cparams: Dict[str, Any]) -> None:
+        """SQLAlchemy event handler to inject Azure AD token on each connection."""
+        try:
+            from azure.identity import DefaultAzureCredential
+        except ImportError:
+            raise missing_dependency_exception("chalkpy[mssql]")
+
+        try:
+            credential = DefaultAzureCredential()
+            token = credential.get_token("https://database.windows.net/.default")
+        except Exception as e:
+            raise Exception(f"Failed to acquire Azure AD token for MSSQL connection: {e}") from e
+
+        token_bytes = token.token.encode("utf-16-le")
+        token_struct = struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)
+        cparams["attrs_before"] = {1256: token_struct}  # SQL_COPT_SS_ACCESS_TOKEN
+
     def get_sqlglot_dialect(self) -> str | None:
         return "tsql"
 
     def local_engine_url(self) -> "URL":
         from sqlalchemy.engine.url import URL
 
-        # Determine authentication method
+        trust_server_cert = env_var_bool("CHALK_MSSQL_TRUST_SERVER_CERTIFICATE", default=False)
+
         if self.client_id and self.client_secret and self.tenant_id:
             # Service Principal authentication
             # Use pyodbc driver for Azure AD support
+            query_params = {
+                "driver": "ODBC Driver 18 for SQL Server",
+                "Authentication": "ActiveDirectoryServicePrincipal",
+            }
+            if trust_server_cert:
+                query_params["TrustServerCertificate"] = "yes"
             return URL.create(
                 drivername="mssql+pyodbc",
                 username=self.client_id,
@@ -129,74 +164,32 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
                 host=self.host,
                 port=self.port,
                 database=self.db,
-                query={
-                    "driver": "ODBC Driver 18 for SQL Server",
-                    "Authentication": "ActiveDirectoryServicePrincipal",
-                },
+                query=query_params,
             )
         elif self.user and self.password:
             # SQL authentication
-            return URL.create(
-                drivername="mssql+pymssql",
-                username=self.user,
-                password=self.password,
-                host=self.host,
-                port=self.port,
-                database=self.db,
-            )
-        else:
-            # Managed Identity authentication
+            query_params = {"driver": "ODBC Driver 18 for SQL Server"}
+            if trust_server_cert:
+                query_params["TrustServerCertificate"] = "yes"
             return URL.create(
                 drivername="mssql+pyodbc",
-                host=self.host,
-                port=self.port,
-                database=self.db,
-                query={
-                    "driver": "ODBC Driver 18 for SQL Server",
-                    "Authentication": "ActiveDirectoryMsi",
-                },
-            )
-
-    def async_local_engine_url(self) -> "URL":
-        from sqlalchemy.engine.url import URL
-
-        # Determine authentication method
-        if self.client_id and self.client_secret and self.tenant_id:
-            # Service Principal authentication
-            return URL.create(
-                drivername="mssql+aioodbc",
-                username=self.client_id,
-                password=self.client_secret,
-                host=self.host,
-                port=self.port,
-                database=self.db,
-                query={
-                    "driver": "ODBC Driver 18 for SQL Server",
-                    "Authentication": "ActiveDirectoryServicePrincipal",
-                },
-            )
-        elif self.user and self.password:
-            # SQL authentication
-            return URL.create(
-                drivername="mssql+aioodbc",
                 username=self.user,
                 password=self.password,
                 host=self.host,
                 port=self.port,
                 database=self.db,
-                query={"driver": "ODBC Driver 18 for SQL Server"},
+                query=query_params,
             )
         else:
-            # Managed Identity authentication
+            # Managed Identity: token injected via event listener
+            connection_string = (
+                f"DRIVER={{ODBC Driver 18 for SQL Server}};SERVER={self.host},{self.port};DATABASE={self.db}"
+            )
+            if trust_server_cert:
+                connection_string += ";TrustServerCertificate=yes"
             return URL.create(
-                drivername="mssql+aioodbc",
-                host=self.host,
-                port=self.port,
-                database=self.db,
-                query={
-                    "driver": "ODBC Driver 18 for SQL Server",
-                    "Authentication": "ActiveDirectoryMsi",
-                },
+                drivername="mssql+pyodbc",
+                query={"odbc_connect": connection_string},
             )
 
     def _recreate_integration_variables(self) -> dict[str, str]:
@@ -214,3 +207,75 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
             ]
             if v is not None
         }
+
+    def execute_query_efficient_raw(
+        self,
+        finalized_query: FinalizedChalkQuery,
+        expected_output_schema: "pa.Schema",
+        connection: Optional["Connection"],
+        query_execution_parameters: QueryExecutionParameters,
+    ) -> Iterable["pa.RecordBatch"]:
+        """Execute query efficiently for MSSQL and return raw PyArrow RecordBatches."""
+        import contextlib
+
+        import pyarrow as pa
+        import pyarrow.compute as pc
+
+        # Get the compiled query
+        _, _, _ = self.compile_query(finalized_query)
+
+        # Use existing connection or create new one
+        with (self.get_engine().connect() if connection is None else contextlib.nullcontext(connection)) as cnx:
+            with cnx.begin():
+                # Handle temp tables
+                with contextlib.ExitStack() as exit_stack:
+                    for (
+                        _,
+                        temp_value,
+                        create_temp_table,
+                        temp_table,
+                        drop_temp_table,
+                    ) in finalized_query.temp_tables.values():
+                        exit_stack.enter_context(
+                            self._create_temp_table(create_temp_table, temp_table, drop_temp_table, cnx, temp_value)
+                        )
+
+                    # Execute query
+                    result = cnx.execute(finalized_query.query, finalized_query.params)
+
+                    # Convert result to PyArrow
+                    rows = result.fetchall()
+                    column_names = result.keys()
+
+                    if not rows:
+                        # Return empty batch with expected schema
+                        arrays = [pa.nulls(0, field.type) for field in expected_output_schema]
+                        batch = pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
+                        if query_execution_parameters.yield_empty_batches:
+                            yield batch
+                        return
+
+                    # Convert rows to column arrays
+                    data: dict[str, list[Any]] = {}
+                    for i, col_name in enumerate(column_names):
+                        col_data = [row[i] for row in rows]
+                        data[col_name] = col_data
+
+                    # Create PyArrow table
+                    table = pa.table(data)
+
+                    # Map columns to expected schema
+                    arrays: list[pa.Array] = []
+                    for field in expected_output_schema:
+                        if field.name in table.column_names:
+                            col = table.column(field.name)
+                            # Cast to expected type if needed
+                            if col.type != field.type:
+                                col = pc.cast(col, field.type)
+                            arrays.append(col)
+                        else:
+                            # Column not found, create null array
+                            arrays.append(pa.nulls(len(table), field.type))
+
+                    batch = pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
+                    yield batch

chalk/sql/_internal/integrations/redshift.py

@@ -260,6 +260,8 @@ class RedshiftSourceImpl(BaseSQLSource):
                 temp_table_name = f"query_{str(uuid.uuid4()).replace('-', '_')}"
                 try:
                     _logger.debug(f"Executing query & creating temp table '{temp_table_name}'")
+                    _public_logger.info(f"Executing Redshift query [{temp_query_id}]: {operation}")
+                    _public_logger.debug(f"Query parameters [{temp_query_id}]: {params}")
                     cursor.execute(f"CREATE TEMP TABLE {temp_table_name} AS ({operation})", params)
                 except Exception as e:
                     _public_logger.error(f"Failed to create temp table for operation: {operation}", exc_info=e)
@@ -366,6 +368,8 @@ class RedshiftSourceImpl(BaseSQLSource):
                 temp_table_name = f"query_{str(uuid.uuid4()).replace('-', '_')}"
                 try:
                     _logger.debug(f"Executing query & creating temp table '{temp_table_name}'")
+                    _public_logger.info(f"Executing Redshift query [{temp_query_id}]: {operation}")
+                    _public_logger.debug(f"Query parameters [{temp_query_id}]: {params}")
                     cursor.execute(f"CREATE TEMP TABLE {temp_table_name} AS ({operation})", params)
                 except Exception as e:
                     _public_logger.error(f"Failed to create temp table for operation: {operation}", exc_info=e)

chalk/sql/_internal/sql_file_resolver.py

@@ -33,7 +33,7 @@ from yaml.scanner import ScannerError
 
 from chalk import Environments, OfflineResolver, OnlineResolver, Tags
 from chalk._lsp.error_builder import SQLFileResolverErrorBuilder
-from chalk.features import DataFrame, Feature, FeatureNotFoundException, Features
+from chalk.features import DataFrame, Feature, FeatureNotFoundException, Features, Underscore
 from chalk.features.feature_set import CURRENT_FEATURE_REGISTRY
 from chalk.features.namespace_context import build_namespaced_name
 from chalk.features.namespace_context import namespace as namespace_ctx
@@ -42,6 +42,8 @@ from chalk.features.resolver import Cron, ResolverArgErrorHandler, StreamResolve
 from chalk.sql._internal.incremental import IncrementalSettings
 from chalk.sql._internal.integrations.bigquery import BigQuerySourceImpl
 from chalk.sql._internal.integrations.cloudsql import CloudSQLSourceImpl
+from chalk.sql._internal.integrations.databricks import DatabricksSourceImpl
+from chalk.sql._internal.integrations.mssql import MSSQLSourceImpl
 from chalk.sql._internal.integrations.mysql import MySQLSourceImpl
 from chalk.sql._internal.integrations.postgres import PostgreSQLSourceImpl
 from chalk.sql._internal.integrations.redshift import RedshiftSourceImpl
@@ -85,8 +87,10 @@ _SOURCES: Mapping[str, Union[Type[BaseSQLSource], Type[StreamSource]]] = {
     "postgres": PostgreSQLSourceImpl,
     "postgresql": PostgreSQLSourceImpl,
     "mysql": MySQLSourceImpl,
+    "mssql": MSSQLSourceImpl,
     "bigquery": BigQuerySourceImpl,
     "cloudsql": CloudSQLSourceImpl,
+    "databricks": DatabricksSourceImpl,
     "redshift": RedshiftSourceImpl,
     "sqlite": SQLiteSourceImpl,
     "kafka": KafkaSource,
@@ -242,6 +246,7 @@ class SQLStringResult:
     override_comment_dict: Optional[CommentDict] = None
     override_name: Optional[str] = None
     autogenerated: bool = False
+    postprocessing_expr: Underscore | None = None
 
     def __post_init__(self):
         # Validation: if autogenerated is True, override_name must not be None
@@ -303,7 +308,7 @@ def get_sql_file_resolvers(
     """Iterate through all `.chalk.sql` filepaths, gather the sql strings, and get a resolver hopefully for each."""
     for dp, dn, fn in os.walk(os.path.expanduser(sql_file_resolve_location)):
         del dn  # unused
-        for f in fn:
+        for f in sorted(fn):  # Sort filenames for deterministic ordering
             filepath = os.path.join(dp, f)
             if not filepath.endswith(CHALK_SQL_FILE_RESOLVER_FILENAME_SUFFIX):
                 continue
@@ -316,7 +321,10 @@ def get_sql_file_resolvers(
                     continue
                 _filepath_to_sql_string[filepath] = sql_string_result.sql_string
             yield get_sql_file_resolver(sources, sql_string_result, has_import_errors)
-    for sql_string_result in _GENERATED_SQL_FILE_RESOLVER_REGISTRY.get_generated_sql_file_resolvers():
+    # Only yield generated resolvers whose filepath is under the directory being scanned
+    for sql_string_result in _GENERATED_SQL_FILE_RESOLVER_REGISTRY.get_generated_sql_file_resolvers(
+        filter_by_directory=sql_file_resolve_location
+    ):
         yield get_sql_file_resolver(sources, sql_string_result, has_import_errors)
 
 
@@ -339,8 +347,12 @@ def get_sql_file_resolvers_from_paths(
             sql_string_result=sql_string_result,
             has_import_errors=has_import_errors,
         )
+    # Only yield generated resolvers whose filepath is in the paths list
+    # If paths is empty, yield all generated resolvers (no filtering)
     for sql_string_result in _GENERATED_SQL_FILE_RESOLVER_REGISTRY.get_generated_sql_file_resolvers():
-        yield get_sql_file_resolver(sources, sql_string_result, has_import_errors)
+        # Check if this generated resolver's filepath is in the provided paths
+        if not paths or sql_string_result.path in paths:
+            yield get_sql_file_resolver(sources, sql_string_result, has_import_errors)
 
 
 def get_sql_file_resolver(
@@ -693,6 +705,7 @@ def get_sql_file_resolver(
                     incremental_settings=incremental_settings,
                     params_to_root_fqn=glot_result.args,
                 ),
+                postprocessing=sql_string_result.postprocessing_expr,
             )
         except Exception as e:
             raise e
@@ -1573,6 +1586,7 @@ class GeneratedSQLFileResolverInfo:
     filepath: str
     sql_string: str
     comment_dict: CommentDict
+    postprocessing_expr: Underscore | None
 
 
 class GeneratedSQLFileResolverRegistry:
@@ -1580,17 +1594,43 @@ class GeneratedSQLFileResolverRegistry:
         super().__init__()
         self.resolver_name_to_generated_infos: Dict[str, GeneratedSQLFileResolverInfo] = {}
 
-    def add_sql_file_resolver(self, name: str, filepath: str, sql_string: str, comment_dict: CommentDict):
+    def add_sql_file_resolver(
+        self,
+        name: str,
+        filepath: str,
+        sql_string: str,
+        comment_dict: CommentDict,
+        postprocessing_expr: Underscore | None = None,
+    ):
         if name in self.resolver_name_to_generated_infos and filepath != "<notebook>":
             raise ValueError(f"A SQL file resolver already exists with name '{name}'. They must have unique names.")
         self.resolver_name_to_generated_infos[name] = GeneratedSQLFileResolverInfo(
-            filepath=filepath,
-            sql_string=sql_string,
-            comment_dict=comment_dict,
+            filepath=filepath, sql_string=sql_string, comment_dict=comment_dict, postprocessing_expr=postprocessing_expr
         )
 
-    def get_generated_sql_file_resolvers(self) -> Iterable[SQLStringResult]:
+    def get_generated_sql_file_resolvers(self, filter_by_directory: Path | None = None) -> Iterable[SQLStringResult]:
+        """
+        Yield generated SQL file resolvers, optionally filtered by directory.
+
+        Args:
+            filter_by_directory: If provided, only yield resolvers whose filepath is under this directory.
+                                 If None, yield all generated resolvers (legacy behavior).
+        """
         for name, generated_info in self.resolver_name_to_generated_infos.items():
+            # If filtering by directory is requested, check if the resolver's filepath is under that directory
+            if filter_by_directory is not None:
+                # Special case: notebook resolvers (filepath == "<notebook>") should never be auto-yielded
+                # when scanning directories, only when explicitly requested
+                if generated_info.filepath == "<notebook>":
+                    continue
+
+                # Convert to absolute paths for comparison and check if resolver path is under filter directory
+                resolver_path = Path(generated_info.filepath).resolve()
+                filter_path = Path(filter_by_directory).resolve()
+
+                if not resolver_path.is_relative_to(filter_path):
+                    continue
+
             yield SQLStringResult(
                 path=generated_info.filepath,
                 sql_string=generated_info.sql_string,
@@ -1598,6 +1638,7 @@ class GeneratedSQLFileResolverRegistry:
                 override_comment_dict=generated_info.comment_dict,
                 override_name=name,
                 autogenerated=True,
+                postprocessing_expr=generated_info.postprocessing_expr,
             )
 
 
@@ -1626,6 +1667,7 @@ def make_sql_file_resolver(
     partitioned_by: Collection[Any] | None = None,
     total: Optional[bool] = None,
     skip_sql_validation: Optional[bool] = None,
+    postprocessing_expression: Optional[Underscore] = None,
 ):
     """Generate a Chalk SQL file resolver from a filepath and a sql string.
     This will generate a resolver in your web dashboard that can be queried,
@@ -1809,6 +1851,7 @@ def make_sql_file_resolver(
         sql_string=sql,
         comment_dict=comment_dict,
         name=name,
+        postprocessing_expr=postprocessing_expression,
     )
     if is_defined_in_notebook:
         from chalk.sql import SQLSourceGroup
@@ -1845,6 +1888,7 @@ def make_sql_file_resolver(
             override_comment_dict=generated_info.comment_dict,
             override_name=name,
             autogenerated=True,
+            postprocessing_expr=postprocessing_expression,
         )
         resolver_result = get_sql_file_resolver(
             sources=current_sql_sources, sql_string_result=info, has_import_errors=False

chalk/sql/_internal/sql_source.py

@@ -214,11 +214,32 @@ class BaseSQLSource(BaseSQLSourceProtocol):
         if getattr(self, "kind", None) != SQLSourceKind.trino:
             engine_args.setdefault("pool_pre_ping", env_var_bool("USE_CLIENT_POOL_PRE_PING"))
             async_engine_args.setdefault("pool_pre_ping", env_var_bool("USE_CLIENT_POOL_PRE_PING"))
-        self.engine_args = engine_args
-        self.async_engine_args = async_engine_args
+        # Store raw args internally, expose filtered versions via properties
+        self._raw_engine_args = engine_args
+        self._raw_async_engine_args = async_engine_args
         self._engine = None
         self._async_engine = None
 
+    @property
+    def engine_args(self) -> Dict[str, Any]:
+        """Engine arguments with native_args filtered out for SQLAlchemy."""
+        return {k: v for k, v in self._raw_engine_args.items() if k != "native_args"}
+
+    @engine_args.setter
+    def engine_args(self, args: dict[str, Any]):
+        """Set raw engine args (for backward compatibility)."""
+        self._raw_engine_args = args
+
+    @property
+    def async_engine_args(self) -> Dict[str, Any]:
+        """Async engine arguments with native_args filtered out for SQLAlchemy."""
+        return {k: v for k, v in self._raw_async_engine_args.items() if k != "native_args"}
+
+    @async_engine_args.setter
+    def async_engine_args(self, args: dict[str, Any]):
+        """Set raw async engine args (for backward compatibility)."""
+        self._raw_async_engine_args = args
+
     @property
     def _engine_args(self):
         """Backcompat support for private subclassing of BaseSQLSource"""
@@ -239,6 +260,16 @@ class BaseSQLSource(BaseSQLSourceProtocol):
         """Backcompat support for private subclassing of BaseSQLSource"""
         self.async_engine_args = args
 
+    @property
+    def native_args(self) -> Dict[str, Any]:
+        """Native arguments to be passed to the underlying database driver.
+
+        These arguments are extracted from engine_args and async_engine_args
+        and are not passed to SQLAlchemy's create_engine or create_async_engine.
+        Instead, they should be used by subclasses to configure native driver connections.
+        """
+        return self._raw_engine_args.get("native_args", {})
+
     def get_sqlglot_dialect(self) -> Union[str, None]:
         """Returns the name of the SQL dialect (if it has one) for `sqlglot` to parse the SQL string.
         This allows for use of dialect-specific syntax while parsing and modifying queries."""
@@ -832,6 +863,7 @@ class BaseSQLSource(BaseSQLSourceProtocol):
         if self._engine is None:
             self.register_sqlalchemy_compiler_overrides()
             self._check_engine_isolation_level()
+            # engine_args property already filters out native_args
             self._engine = create_engine(url=self.local_engine_url(), **self.engine_args)
         return self._engine
 
@@ -841,6 +873,7 @@ class BaseSQLSource(BaseSQLSourceProtocol):
         if self._async_engine is None:
             self.register_sqlalchemy_compiler_overrides()
             self._check_engine_isolation_level()
+            # async_engine_args property already filters out native_args
             self._async_engine = create_async_engine(url=self.async_local_engine_url(), **self.async_engine_args)
         return self._async_engine
 

chalk/streams/_kafka_source.py

@@ -27,6 +27,7 @@ _KAFKA_SASL_MECHANISM_NAME = "KAFKA_SASL_MECHANISM"
 _KAFKA_SASL_USERNAME_NAME = "KAFKA_SASL_USERNAME"
 _KAFKA_SASL_PASSWORD_NAME = "KAFKA_SASL_PASSWORD"
 _KAFKA_ADDITIONAL_KAFKA_ARGS_NAME = "KAFKA_ADDITIONAL_KAFKA_ARGS"
+_KAFKA_DEAD_LETTER_QUEUE_TOPIC = "KAFKA_DEAD_LETTER_QUEUE_TOPIC"
 
 
 class KafkaSource(StreamSource, SinkIntegrationProtocol, BaseModel, frozen=True):
@@ -165,7 +166,10 @@ class KafkaSource(StreamSource, SinkIntegrationProtocol, BaseModel, frozen=True)
             ),
             name=name,
             late_arrival_deadline=late_arrival_deadline,
-            dead_letter_queue_topic=dead_letter_queue_topic,
+            dead_letter_queue_topic=dead_letter_queue_topic
+            or load_integration_variable(
+                name=_KAFKA_DEAD_LETTER_QUEUE_TOPIC, integration_name=name, override=integration_variable_override
+            ),
             ssl_ca_file=ssl_ca_file
             or load_integration_variable(
                 name=_KAFKA_SSL_CA_FILE_NAME, integration_name=name, override=integration_variable_override

chalk/streams/_windows.py

@@ -70,6 +70,7 @@ class WindowedMeta(type, Generic[TRich]):
             validations=None,
             offline_ttl=None,
             expression=None,
+            offline_expression=None,
             materialization=None,
         )  # noqa
 
@@ -79,9 +80,14 @@ JsonValue = Any
 
 def get_name_with_duration(name_or_fqn: str, duration: Union[str, int, timedelta]) -> str:
     duration_secs = parse_chalk_duration_s(duration)
+    name_or_fqn_components = name_or_fqn.split("@")
+    assert len(name_or_fqn_components) <= 2, f"Received invalid fqn format.\nfqn={name_or_fqn}"
+    unversioned_fqn = name_or_fqn_components[0]
+    version = None if len(name_or_fqn_components) != 2 else name_or_fqn_components[1]
+
     if duration_secs >= CHALK_MAX_TIMEDELTA.total_seconds():
-        return f"{name_or_fqn}__all__"
-    return f"{name_or_fqn}__{duration_secs}__"
+        return f"{unversioned_fqn}__all__" + ("" if version is None else f"@{version}")
+    return f"{unversioned_fqn}__{duration_secs}__" + ("" if version is None else f"@{version}")
 
 
 if TYPE_CHECKING:
@@ -175,6 +181,7 @@ class Windowed(Generic[TRich], metaclass=_WINDOWED_METACLASS):
             window_durations=tuple(self.buckets_seconds) if bucket is None else tuple(),
             window_duration=window_duration,
             underscore_expression=self._expression,
+            offline_underscore_expression=self._offline_expression,
             window_materialization=(
                 MaterializationWindowConfig(bucket_duration=timedelta(seconds=window_duration))
                 if self._materialization is True and window_duration is not None
@@ -208,6 +215,7 @@ class Windowed(Generic[TRich], metaclass=_WINDOWED_METACLASS):
         kind: Type[TRich] | None,
         offline_ttl: Duration | ellipsis | None,
         expression: Underscore | None,
+        offline_expression: Underscore | None,
         materialization: MaterializationWindowConfig | Literal[True] | None,
     ):
         super().__init__()
@@ -234,6 +242,7 @@ class Windowed(Generic[TRich], metaclass=_WINDOWED_METACLASS):
         self._validations = validations
         self._dtype = dtype
         self._expression = expression
+        self._offline_expression = offline_expression
         self._materialization = materialization
 
 
@@ -574,6 +583,7 @@ def windowed(
     validations: List[Validation] | None = None,
     dtype: pa.DataType | None = None,
     expression: Underscore | None = None,
+    offline_expression: Underscore | None = None,
     materialization: MaterializationWindowConfig | Literal[True] | None = None,
 ) -> Windowed[TRich]:
     """Create a windowed feature.
@@ -650,6 +660,8 @@ def windowed(
         the feature value and will treat it as failed.
     expression
         The expression to compute the feature. This is an underscore expression, like `_.transactions[_.amount].sum()`.
+    offline_expression
+        Defines an alternate expression to compute the feature during offline queries.
     validations
         A list of Validations to apply to this feature.
 
@@ -748,5 +760,6 @@ def windowed(
         validations=validations,
         offline_ttl=offline_ttl,
         expression=expression,
+        offline_expression=offline_expression,
         materialization=materialization,
     )

chalk/utils/_otel_version.py

@@ -0,0 +1,13 @@
+def _can_use_otel_trace() -> bool:
+    try:
+        from opentelemetry import context  # pyright: ignore[reportUnusedImport]
+        from opentelemetry import trace  # pyright: ignore[reportUnusedImport]
+        from opentelemetry.sdk.resources import Resource  # pyright: ignore[reportUnusedImport]
+        from opentelemetry.sdk.trace import TracerProvider  # pyright: ignore[reportUnusedImport]
+
+        return True
+    except Exception:
+        return False
+
+
+can_use_otel_trace = _can_use_otel_trace()

chalk/utils/async_helpers.py

@@ -61,8 +61,8 @@ def run_coroutine_fn_threadsafe(
             @functools.wraps(coro_fn)
             async def wrapped_with_context(*args: P.args, **kwargs: P.kwargs):
                 try:
-                    safe_activate_trace_context(current_trace_context)
-                    res = await coro_fn(*args, **kwargs)
+                    with safe_activate_trace_context(current_trace_context):
+                        res = await coro_fn(*args, **kwargs)
                 except BaseException as exc:
                     ans.set_exception(exc)
                 else:

chalk/utils/missing_dependency.py

@@ -5,7 +5,8 @@ class MissingDependencyException(ImportError):
     ...
 
 
-def missing_dependency_exception(name: str):
-    return MissingDependencyException(
-        f"Missing pip dependency '{name}' for chalkpy=={chalk.__version__}. Please add this to your requirements.txt file and pip install."
-    )
+def missing_dependency_exception(name: str, original_error: Exception | None = None):
+    msg = f"Missing pip dependency '{name}' for chalkpy=={chalk.__version__}. Please add this to your requirements.txt file and pip install."
+    if original_error:
+        msg += f"\n\n{original_error}"
+    return MissingDependencyException(msg)