chalkpy 2.95.2__tar.gz → 2.95.4__tar.gz

{chalkpy-2.95.2 → chalkpy-2.95.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.95.2
+Version: 2.95.4
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

chalkpy-2.95.4/chalk/_version.py

@@ -0,0 +1 @@
+__version__ = "2.95.4"

{chalkpy-2.95.2 → chalkpy-2.95.4}/chalk/features/feature_set_decorator.py

@@ -1311,7 +1311,7 @@ def _process_class(
                 )
                 cls_cast = cast(Type, cls)
                 if not f_i.is_typ_set():
-                    f_i.typ = ParsedAnnotation(cls_cast, f.name)
+                    f_i.typ = ParsedAnnotation(cls_cast, f.attribute_name)
 
                 f_i.features_cls = cls_cast
                 f_i.attribute_name = f"{f.attribute_name}_v{i}"
@@ -1329,9 +1329,18 @@ def _process_class(
                     class_cache_strategy=cache_strategy,
                     error_builder=error_builder,
                 )
+
+                # copy over expressions for default version
+                if i == f.version.default:
+                    if f.underscore_expression is None and f_i.underscore_expression is not None:
+                        f.underscore_expression = f_i.underscore_expression
+                    if f.offline_underscore_expression is None and f_i.offline_underscore_expression is not None:
+                        f.offline_underscore_expression = f_i.offline_underscore_expression
+
                 f.version.reference[i] = f_i
                 # The default feature already exists.
                 f_i.no_display = i == f.version.default
+                cls_fields.append(f_i)
             continue
 
         for i in range(1, f.version.maximum + 1):

{chalkpy-2.95.2 → chalkpy-2.95.4}/chalk/sql/_internal/integrations/mssql.py

@@ -2,15 +2,19 @@ from __future__ import annotations
 
 import os
 import struct
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Union
+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"
@@ -42,8 +46,8 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
     ):
         try:
             import pyodbc
-        except ImportError:
-            raise missing_dependency_exception("chalkpy[mssql]")
+        except ImportError as e:
+            raise missing_dependency_exception("chalkpy[mssql]", original_error=e)
         del pyodbc
 
         self.name = name
@@ -142,9 +146,17 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
     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,
@@ -152,13 +164,13 @@ 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
+            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,
@@ -166,13 +178,15 @@ class MSSQLSourceImpl(BaseSQLSource, TableIngestMixIn, SQLSourceWithTableIngestP
                 host=self.host,
                 port=self.port,
                 database=self.db,
-                query={"driver": "ODBC Driver 18 for SQL Server"},
+                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},
@@ -193,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

chalkpy-2.95.4/chalk/utils/missing_dependency.py

@@ -0,0 +1,12 @@
+import chalk
+
+
+class MissingDependencyException(ImportError):
+    ...
+
+
+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)

{chalkpy-2.95.2 → chalkpy-2.95.4}/chalk/utils/tracing.py

@@ -5,7 +5,7 @@ import os
 import threading
 import time
 import types
-from typing import TYPE_CHECKING, Any, Mapping, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Mapping, Union, cast
 
 from chalk.utils._ddtrace_version import can_use_datadog_statsd, can_use_ddtrace
 from chalk.utils._otel_version import can_use_otel_trace
@@ -16,6 +16,41 @@ if TYPE_CHECKING:
     import ddtrace.context
     from opentelemetry import trace as otel_trace
 
+
+class Once:
+    """Execute a function exactly once and block all callers until the function returns
+
+    Same as golang's `sync.Once <https://pkg.go.dev/sync#Once>`_
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._done = False
+        super().__init__()
+
+    def do_once(self, func: Callable[[], None]) -> bool:
+        """Execute ``func`` if it hasn't been executed or return.
+
+        Will block until ``func`` has been called by one thread.
+
+        Returns:
+            Whether or not ``func`` was executed in this call
+        """
+
+        # fast path, try to avoid locking
+        if self._done:
+            return False
+
+        with self._lock:
+            if not self._done:
+                func()
+                self._done = True
+                return True
+        return False
+
+
+_TRACING_CONFIGURED = Once()
+
 _logger = get_logger(__name__)
 
 if can_use_otel_trace:
@@ -27,6 +62,7 @@ if can_use_otel_trace:
 
     @contextlib.contextmanager
     def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         if attributes is None:
             attributes = {}
         attributes = dict(attributes)
@@ -36,14 +72,17 @@ if can_use_otel_trace:
             yield span
 
     def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         current_span = otel_trace.get_current_span()
         current_span.set_attributes(dict(metrics))
 
     def safe_add_tags(tags: Mapping[str, str]):
+        configure_tracing("chalkpy")
         current_span = otel_trace.get_current_span()
         current_span.set_attributes(dict(tags))
 
     def safe_current_trace_context() -> ddtrace.context.Context | otel_trace.SpanContext | None:  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         return otel_trace.get_current_span().get_span_context()
 
     @contextlib.contextmanager
@@ -53,6 +92,7 @@ if can_use_otel_trace:
         | otel_trace.SpanContext
         | None,  # pyright: ignore[reportPrivateImportUsage]
     ):
+        configure_tracing("chalkpy")
         if isinstance(ctx, otel_trace.SpanContext):
             new_span = otel_trace.NonRecordingSpan(ctx)
             new_context = otel_trace.set_span_in_context(new_span)
@@ -65,6 +105,7 @@ if can_use_otel_trace:
     def add_trace_headers(  # pyright: ignore[reportRedeclaration]
         input_headers: None | dict[str, str]
     ) -> dict[str, str]:
+        configure_tracing("chalkpy")
         current_span_ctx = otel_trace.get_current_span().get_span_context()
         new_span_ctx = otel_trace.SpanContext(
             trace_id=current_span_ctx.trace_id,
@@ -86,6 +127,7 @@ elif can_use_ddtrace:
 
     @contextlib.contextmanager
     def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         if not ddtrace.tracer.enabled:
             yield
             return
@@ -106,16 +148,19 @@ elif can_use_ddtrace:
             yield
 
     def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         span = ddtrace.tracer.current_span()
         if span:
             span.set_metrics(cast(Any, metrics))
 
     def safe_add_tags(tags: Mapping[str, str]):  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         span = ddtrace.tracer.current_span()
         if span:
             span.set_tags(cast(Any, tags))
 
     def safe_current_trace_context() -> ddtrace.context.Context | otel_trace.SpanContext | None:  # pyright: ignore[reportRedeclaration]
+        configure_tracing("chalkpy")
         return ddtrace.tracer.current_trace_context()
 
     @contextlib.contextmanager
@@ -125,6 +170,7 @@ elif can_use_ddtrace:
         | otel_trace.SpanContext
         | None,  # pyright: ignore[reportPrivateImportUsage]
     ):
+        configure_tracing("chalkpy")
         if isinstance(ctx, ddtrace.context.Context) or isinstance(ctx, ddtrace.Span):
             ddtrace.tracer.context_provider.activate(ctx)
         yield
@@ -132,6 +178,7 @@ elif can_use_ddtrace:
     def add_trace_headers(  # pyright: ignore[reportRedeclaration]
         input_headers: None | dict[str, str]
     ) -> dict[str, str]:
+        configure_tracing("chalkpy")
         headers: dict[str, str] = dict(input_headers if input_headers is not None else {})
         span = ddtrace.tracer.current_span()
         if span:
@@ -223,74 +270,74 @@ class PerfTimer:
 
 
 def configure_tracing(default_service_name: str):
-    from chalk.utils.log_with_context import get_logger
-
-    _logger = get_logger(__name__)
-
-    if can_use_otel_trace:
-        from opentelemetry import trace as otel_trace
-        from opentelemetry.sdk.resources import Resource
-        from opentelemetry.sdk.trace import TracerProvider
-
-        provider = TracerProvider(
-            resource=Resource.create(
-                {
-                    "service.name": default_service_name,
-                }
-            ),
-        )
-        otel_trace.set_tracer_provider(provider)
-
-    elif can_use_ddtrace:
-        import ddtrace
-        from ddtrace.filters import FilterRequestsOnUrl
-
-        if ddtrace.config.service is None:
-            ddtrace.config.service = default_service_name
-        # Re-configuring the global tracer to capture any setting changes from environs from a .dotenv file
-        # which might be loaded after the first ddtrace import
-
-        ddtrace.tracer.configure(
-            enabled=None if "DD_TRACE_ENABLED" not in os.environ else env_var_bool("DD_TRACE_ENABLED"),
-            hostname=os.getenv("DD_AGENT_HOST") or os.getenv("DD_TRACE_AGENT_URL"),
-            uds_path=os.getenv("DD_TRACE_AGENT_URL"),
-            dogstatsd_url=os.getenv("DD_DOGSTATSD_URL"),
-            api_version=os.getenv("DD_TRACE_API_VERSION"),
-            compute_stats_enabled=env_var_bool("DD_TRACE_COMPUTE_STATS"),
-            iast_enabled=None if "DD_IAST_ENABLED" not in os.environ else env_var_bool("DD_IAST_ENABLED"),
-            # exclude healthcheck url from apm trace collection
-            settings={
-                "FILTERS": [
-                    FilterRequestsOnUrl(
-                        [
-                            r"^http://.*/healthcheck$",
-                            r"^http://.*/ready$",
-                            r"^http://[^/]*/$",  # exclude "/"
-                        ]
-                    )
-                ]
-            },
-        )
-        if ddtrace.tracer.enabled:
-            ddtrace.patch(
-                asyncio=True,
-                databricks=False,
-                fastapi=True,
-                futures=True,
-                httplib=True,
-                httpx=True,
-                psycopg=True,
-                redis=True,
-                requests=True,
-                sqlalchemy=False,
-                urllib3=True,
+    def do_configure_tracing():
+        from chalk.utils.log_with_context import get_logger
+
+        _logger = get_logger(__name__)
+
+        if can_use_otel_trace:
+            from opentelemetry import trace as otel_trace
+            from opentelemetry.sdk.resources import Resource
+            from opentelemetry.sdk.trace import TracerProvider
+
+            provider = TracerProvider(
+                resource=Resource.create(
+                    {
+                        "service.name": default_service_name,
+                    }
+                ),
             )
+            otel_trace.set_tracer_provider(provider)
+
+        elif can_use_ddtrace:
+            import ddtrace
+            from ddtrace.filters import FilterRequestsOnUrl
+
+            if ddtrace.config.service is None:
+                ddtrace.config.service = default_service_name
+            # Re-configuring the global tracer to capture any setting changes from environs from a .dotenv file
+            # which might be loaded after the first ddtrace import
+
+            ddtrace.tracer.configure(
+                enabled=None if "DD_TRACE_ENABLED" not in os.environ else env_var_bool("DD_TRACE_ENABLED"),
+                hostname=os.getenv("DD_AGENT_HOST") or os.getenv("DD_TRACE_AGENT_URL"),
+                uds_path=os.getenv("DD_TRACE_AGENT_URL"),
+                dogstatsd_url=os.getenv("DD_DOGSTATSD_URL"),
+                api_version=os.getenv("DD_TRACE_API_VERSION"),
+                compute_stats_enabled=env_var_bool("DD_TRACE_COMPUTE_STATS"),
+                iast_enabled=None if "DD_IAST_ENABLED" not in os.environ else env_var_bool("DD_IAST_ENABLED"),
+                # exclude healthcheck url from apm trace collection
+                settings={
+                    "FILTERS": [
+                        FilterRequestsOnUrl(
+                            [
+                                r"^http://.*/healthcheck$",
+                                r"^http://.*/ready$",
+                                r"^http://[^/]*/$",  # exclude "/"
+                            ]
+                        )
+                    ]
+                },
+            )
+            if ddtrace.tracer.enabled:
+                ddtrace.patch(
+                    asyncio=True,
+                    databricks=False,
+                    fastapi=True,
+                    futures=True,
+                    httplib=True,
+                    httpx=True,
+                    psycopg=True,
+                    redis=True,
+                    requests=True,
+                    sqlalchemy=False,
+                    urllib3=True,
+                )
 
-        _logger.info(
-            f"Configuring DDtrace tracing: enabled={ddtrace.tracer.enabled}, service={ddtrace.config.service}, env={ddtrace.config.env}, trace_agent_url: {ddtrace.config._trace_agent_url}, effective trace agent: {ddtrace.tracer._agent_url}"  # pyright: ignore [reportAttributeAccessIssue, reportPrivateUsage]
-        )
-    else:
-        _logger.warning("neither opentelemetry nor ddtrace are installed")
-
+            _logger.info(
+                f"Configuring DDtrace tracing: enabled={ddtrace.tracer.enabled}, service={ddtrace.config.service}, env={ddtrace.config.env}, trace_agent_url: {ddtrace.config._trace_agent_url}, effective trace agent: {ddtrace.tracer._agent_url}"  # pyright: ignore [reportAttributeAccessIssue, reportPrivateUsage]
+            )
+        else:
+            _logger.warning("neither opentelemetry nor ddtrace are installed")
 
-configure_tracing("chalkpy")
+    _TRACING_CONFIGURED.do_once(do_configure_tracing)

chalkpy-2.95.2/chalk/_version.py

@@ -1 +0,0 @@
-__version__ = "2.95.2"

chalkpy-2.95.2/chalk/utils/missing_dependency.py

@@ -1,11 +0,0 @@
-import chalk
-
-
-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."
-    )