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

chalk/sql/_internal/integrations/mssql.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+import os
+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:
+    import pyarrow as pa
+    from sqlalchemy.engine import URL, Connection
+
+_MSSQL_HOST_NAME = "MSSQL_HOST"
+_MSSQL_TCP_PORT_NAME = "MSSQL_TCP_PORT"
+_MSSQL_DATABASE_NAME = "MSSQL_DATABASE"
+_MSSQL_USER_NAME = "MSSQL_USER"
+_MSSQL_PWD_NAME = "MSSQL_PWD"
+_MSSQL_CLIENT_ID_NAME = "MSSQL_CLIENT_ID"
+_MSSQL_CLIENT_SECRET_NAME = "MSSQL_CLIENT_SECRET"
+_MSSQL_TENANT_ID_NAME = "MSSQL_TENANT_ID"
+
+
+class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestProtocol):
+    kind = SQLSourceKind.mssql
+
+    def __init__(
+        self,
+        host: Optional[str] = None,
+        port: Optional[Union[int, str]] = None,
+        db: Optional[str] = None,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        name: Optional[str] = None,
+        engine_args: Optional[Dict[str, Any]] = None,
+        async_engine_args: Optional[Dict[str, Any]] = None,
+        integration_variable_override: Optional[Mapping[str, str]] = None,
+    ):
+        try:
+            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(
+            integration_name=name, name=_MSSQL_HOST_NAME, override=integration_variable_override
+        )
+        self.port = (
+            int(port)
+            if port is not None
+            else load_integration_variable(
+                integration_name=name, name=_MSSQL_TCP_PORT_NAME, parser=int, override=integration_variable_override
+            )
+        )
+        self.db = db or load_integration_variable(
+            integration_name=name, name=_MSSQL_DATABASE_NAME, override=integration_variable_override
+        )
+        self.user = user or load_integration_variable(
+            integration_name=name,
+            name=_MSSQL_USER_NAME,
+            override=integration_variable_override,
+        )
+        self.password = password or load_integration_variable(
+            integration_name=name,
+            name=_MSSQL_PWD_NAME,
+            override=integration_variable_override,
+        )
+        self.client_id = client_id or load_integration_variable(
+            integration_name=name,
+            name=_MSSQL_CLIENT_ID_NAME,
+            override=integration_variable_override,
+        )
+        self.client_secret = client_secret or load_integration_variable(
+            integration_name=name,
+            name=_MSSQL_CLIENT_SECRET_NAME,
+            override=integration_variable_override,
+        )
+        self.tenant_id = tenant_id or load_integration_variable(
+            integration_name=name,
+            name=_MSSQL_TENANT_ID_NAME,
+            override=integration_variable_override,
+        )
+        self.ingested_tables: Dict[str, Any] = {}
+
+        if engine_args is None:
+            engine_args = {}
+        if async_engine_args is None:
+            async_engine_args = {}
+
+        if name:
+            engine_args_from_ui = self._load_env_engine_args(name, override=integration_variable_override)
+            for k, v in engine_args_from_ui.items():
+                engine_args.setdefault(k, v)
+                async_engine_args.setdefault(k, v)
+
+        chalk_default_engine_args = {
+            "pool_size": 20,
+            "max_overflow": 60,
+            "pool_recycle": 90,
+        }
+        for k, v in chalk_default_engine_args.items():
+            engine_args.setdefault(k, v)
+            async_engine_args.setdefault(k, v)
+
+        # Set isolation level for read-only operations
+        engine_args.setdefault("isolation_level", os.environ.get("CHALK_SQL_ISOLATION_LEVEL", "AUTOCOMMIT"))
+        async_engine_args.setdefault("isolation_level", os.environ.get("CHALK_SQL_ISOLATION_LEVEL", "AUTOCOMMIT"))
+
+        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
+
+        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,
+                password=self.client_secret,
+                host=self.host,
+                port=self.port,
+                database=self.db,
+                query=query_params,
+            )
+        elif self.user and self.password:
+            # SQL authentication
+            query_params = {"driver": "ODBC Driver 18 for SQL Server"}
+            if trust_server_cert:
+                query_params["TrustServerCertificate"] = "yes"
+            return URL.create(
+                drivername="mssql+pyodbc",
+                username=self.user,
+                password=self.password,
+                host=self.host,
+                port=self.port,
+                database=self.db,
+                query=query_params,
+            )
+        else:
+            # 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+pyodbc",
+                query={"odbc_connect": connection_string},
+            )
+
+    def _recreate_integration_variables(self) -> dict[str, str]:
+        return {
+            k: v
+            for k, v in [
+                create_integration_variable(_MSSQL_HOST_NAME, self.name, self.host),
+                create_integration_variable(_MSSQL_TCP_PORT_NAME, self.name, self.port),
+                create_integration_variable(_MSSQL_DATABASE_NAME, self.name, self.db),
+                create_integration_variable(_MSSQL_USER_NAME, self.name, self.user),
+                create_integration_variable(_MSSQL_PWD_NAME, self.name, self.password),
+                create_integration_variable(_MSSQL_CLIENT_ID_NAME, self.name, self.client_id),
+                create_integration_variable(_MSSQL_CLIENT_SECRET_NAME, self.name, self.client_secret),
+                create_integration_variable(_MSSQL_TENANT_ID_NAME, self.name, self.tenant_id),
+            ]
+            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/postgres.py

@@ -28,6 +28,7 @@ from chalk.sql.protocols import SQLSourceWithTableIngestProtocol
 from chalk.utils.environment_parsing import env_var_bool
 from chalk.utils.log_with_context import get_logger
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import polars_uses_schema_overrides
 from chalk.utils.tracing import safe_add_metrics, safe_add_tags, safe_trace
 
 if TYPE_CHECKING:
@@ -262,8 +263,11 @@ class PostgreSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIn
             # pl.read_csv(use_pyarrow=True) has the same performance degradation,
             # UNLESS a `dtypes` arg is provided.
 
-            # 'dtypes' deprecated for 'schema_overrides' in polars 0.20+, but parameter renamed without breaking
-            pl_table = pl.read_csv(buffer, dtypes=parse_dtypes)  # pyright: ignore[reportCallIssue]
+            # 'dtypes' deprecated for 'schema_overrides' in polars 0.20.31+
+            if polars_uses_schema_overrides:
+                pl_table = pl.read_csv(buffer, schema_overrides=parse_dtypes)  # pyright: ignore[reportCallIssue]
+            else:
+                pl_table = pl.read_csv(buffer, dtypes=parse_dtypes)  # pyright: ignore[reportCallIssue]
             if boolean_columns:
                 # DO NOT use map_dict. Causes a segfault when multiple uvicorn workers are handling
                 # requests in parallel.
@@ -498,7 +502,11 @@ class PostgreSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIn
                 else:
                     parse_dtypes[field.name] = pl.Utf8
 
-            pl_table = pl.read_csv(buffer, dtypes=parse_dtypes)  # pyright: ignore[reportCallIssue]
+            # 'dtypes' deprecated for 'schema_overrides' in polars 0.20.31+
+            if polars_uses_schema_overrides:
+                pl_table = pl.read_csv(buffer, schema_overrides=parse_dtypes)  # pyright: ignore[reportCallIssue]
+            else:
+                pl_table = pl.read_csv(buffer, dtypes=parse_dtypes)  # pyright: ignore[reportCallIssue]
 
             # Convert to arrow and map to expected schema
             arrow_table = pl_table.to_arrow()

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/integrations/snowflake.py

@@ -31,6 +31,7 @@ from chalk.sql.finalized_query import FinalizedChalkQuery
 from chalk.utils.df_utils import is_list_like, pa_array_to_pl_series
 from chalk.utils.environment_parsing import env_var_bool
 from chalk.utils.missing_dependency import missing_dependency_exception
+from chalk.utils.pl_helpers import str_json_decode_compat
 from chalk.utils.threading import DEFAULT_IO_EXECUTOR, MultiSemaphore
 from chalk.utils.tracing import safe_incr, safe_set_gauge
 
@@ -398,11 +399,19 @@ class SnowflakeSourceImpl(BaseSQLSource):
                 if pa.types.is_list(expected_type) or pa.types.is_large_list(expected_type):
                     if pa.types.is_string(actual_type) or pa.types.is_large_string(actual_type):
                         series = pa_array_to_pl_series(tbl[col_name])
-                        column = series.str.json_extract(feature.converter.polars_dtype).to_arrow().cast(expected_type)
+                        column = (
+                            str_json_decode_compat(series, feature.converter.polars_dtype)
+                            .to_arrow()
+                            .cast(expected_type)
+                        )
                 if pa.types.is_struct(expected_type):
                     if pa.types.is_string(actual_type):
                         series = pa_array_to_pl_series(tbl[col_name])
-                        column = series.str.json_extract(feature.converter.polars_dtype).to_arrow().cast(expected_type)
+                        column = (
+                            str_json_decode_compat(series, feature.converter.polars_dtype)
+                            .to_arrow()
+                            .cast(expected_type)
+                        )
                 if actual_type != expected_type:
                     column = column.cast(options=pc.CastOptions(target_type=expected_type, allow_time_truncate=True))
                 if isinstance(column, pa.ChunkedArray):

chalk/sql/_internal/integrations/util.py

@@ -4,6 +4,7 @@ import pyarrow as pa
 
 from chalk.features._encoding.pyarrow import pyarrow_to_polars
 from chalk.utils.df_utils import pa_cast, pa_table_to_pl_df
+from chalk.utils.pl_helpers import str_json_decode_compat
 
 
 def convert_hex_to_binary(table: pa.Table, cols_to_convert: List[str]) -> pa.Table:
@@ -60,7 +61,7 @@ def json_parse_and_cast(tbl: pa.Table, schema: Mapping[str, pa.DataType]) -> pa.
             expr = pl.col(col_name)
             if pl_df.schema[col_name] == pl.Binary():
                 expr = expr.cast(pl.Utf8())
-            expr = expr.str.json_extract(pl_dtype).alias(col_name)
+            expr = str_json_decode_compat(expr, pl_dtype).alias(col_name)
             pl_exprs.append(expr)
 
         pl_df = pl_df.with_columns(pl_exprs)

chalk/sql/_internal/sql_file_resolver.py

@@ -19,20 +19,21 @@ from typing import (
     Literal,
     Mapping,
     Optional,
+    ParamSpec,
     Sequence,
     Type,
+    TypeAlias,
     TypeVar,
     Union,
     cast,
 )
 
 import yaml
-from typing_extensions import ParamSpec, TypeAlias
 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
@@ -41,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
@@ -84,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,
@@ -241,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
@@ -302,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
@@ -315,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)
 
 
@@ -338,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(
@@ -692,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
@@ -1572,6 +1586,7 @@ class GeneratedSQLFileResolverInfo:
     filepath: str
     sql_string: str
     comment_dict: CommentDict
+    postprocessing_expr: Underscore | None
 
 
 class GeneratedSQLFileResolverRegistry:
@@ -1579,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,
@@ -1597,6 +1638,7 @@ class GeneratedSQLFileResolverRegistry:
                 override_comment_dict=generated_info.comment_dict,
                 override_name=name,
                 autogenerated=True,
+                postprocessing_expr=generated_info.postprocessing_expr,
             )
 
 
@@ -1625,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,
@@ -1808,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
@@ -1844,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

@@ -163,6 +163,7 @@ class SQLSourceKind(str, Enum):
     athena = "athena"
     duckdb = "duckdb"
     dynamodb = "dynamodb"
+    mssql = "mssql"
     mysql = "mysql"
     postgres = "postgres"
     redshift = "redshift"
@@ -213,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"""
@@ -238,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."""
@@ -831,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
 
@@ -840,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/__init__.py

@@ -1,7 +1,5 @@
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, TypeVar, Union
-
-from typing_extensions import ParamSpec
+from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, ParamSpec, TypeVar, Union
 
 from chalk._lsp.error_builder import get_resolver_error_builder
 from chalk.features.tag import Environments

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