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

chalk/utils/tracing.py

@@ -8,37 +8,85 @@ import types
 from typing import TYPE_CHECKING, Any, 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
 from chalk.utils.environment_parsing import env_var_bool
+from chalk.utils.log_with_context import get_logger
 
 if TYPE_CHECKING:
     import ddtrace.context
+    from opentelemetry import trace as otel_trace
 
-if can_use_ddtrace and can_use_datadog_statsd:
-    import ddtrace
-    from datadog.dogstatsd.base import statsd
-    from ddtrace.propagation.http import HTTPPropagator
+_logger = get_logger(__name__)
 
-    def safe_set_gauge(gauge: str, value: int | float):
-        statsd.gauge(gauge, value)
+if can_use_otel_trace:
+    from opentelemetry import context as otel_context
+    from opentelemetry import trace as otel_trace
+    from opentelemetry.propagate import inject as otel_inject
 
-    def safe_incr(counter: str, value: int | float, tags: list[str] | None = None):
-        statsd.increment(counter, value, tags)
+    _logger.debug("OTEL trace packages installed, otel tracing is available")
 
-    def safe_distribution(counter: str, value: int | float, tags: list[str] | None = None):
-        statsd.distribution(counter, value, tags)
+    @contextlib.contextmanager
+    def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):  # pyright: ignore[reportRedeclaration]
+        if attributes is None:
+            attributes = {}
+        attributes = dict(attributes)
+        attributes["thread_id"] = str(threading.get_native_id())
+        with otel_trace.get_tracer("chalk").start_as_current_span(span_id) as span:
+            span.set_attributes(attributes)
+            yield span
+
+    def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):  # pyright: ignore[reportRedeclaration]
+        current_span = otel_trace.get_current_span()
+        current_span.set_attributes(dict(metrics))
+
+    def safe_add_tags(tags: Mapping[str, str]):
+        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]
+        return otel_trace.get_current_span().get_span_context()
 
     @contextlib.contextmanager
-    def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):
-        if not ddtrace.tracer.enabled:
+    def safe_activate_trace_context(  # pyright: ignore[reportRedeclaration]
+        ctx: ddtrace.context.Context
+        | ddtrace.Span
+        | otel_trace.SpanContext
+        | None,  # pyright: ignore[reportPrivateImportUsage]
+    ):
+        if isinstance(ctx, otel_trace.SpanContext):
+            new_span = otel_trace.NonRecordingSpan(ctx)
+            new_context = otel_trace.set_span_in_context(new_span)
+            token = otel_context.attach(new_context)
             yield
-            return
-        if (current_ctx := ddtrace.tracer.current_trace_context()) is None:
+            otel_context.detach(token)
+        else:
             yield
-            return
-        if (priority := current_ctx.sampling_priority) is not None and priority <= 0:
-            # If a priority is negative, then it won't be sampled
-            # See https://github.com/DataDog/dd-trace-py/blob/09edef713bf9f0ab30f554bf7765d7a7c2ed6f30/ddtrace/constants.py#L74
-            # Not sure what a priority=None means
+
+    def add_trace_headers(  # pyright: ignore[reportRedeclaration]
+        input_headers: None | dict[str, str]
+    ) -> dict[str, str]:
+        current_span_ctx = otel_trace.get_current_span().get_span_context()
+        new_span_ctx = otel_trace.SpanContext(
+            trace_id=current_span_ctx.trace_id,
+            span_id=current_span_ctx.span_id,
+            is_remote=current_span_ctx.is_remote,
+            trace_flags=otel_trace.TraceFlags(otel_trace.TraceFlags.SAMPLED),
+            trace_state=current_span_ctx.trace_state,
+        )
+        ctx = otel_trace.set_span_in_context(otel_trace.NonRecordingSpan(new_span_ctx))
+        headers: dict[str, str] = dict(input_headers if input_headers is not None else {})
+        otel_inject(headers, context=ctx)
+        return headers
+
+elif can_use_ddtrace:
+    import ddtrace
+    from ddtrace.propagation.http import HTTPPropagator
+
+    _logger.debug("ddtrace installed and available, using it to trace")
+
+    @contextlib.contextmanager
+    def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):  # pyright: ignore[reportRedeclaration]
+        if not ddtrace.tracer.enabled:
             yield
             return
         if attributes is None:
@@ -57,69 +105,95 @@ if can_use_ddtrace and can_use_datadog_statsd:
                 span.set_tags(cast(Any, attributes))
             yield
 
-    def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):
+    def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):  # pyright: ignore[reportRedeclaration]
         span = ddtrace.tracer.current_span()
         if span:
             span.set_metrics(cast(Any, metrics))
 
-    def safe_add_tags(tags: Mapping[str, str]):
+    def safe_add_tags(tags: Mapping[str, str]):  # pyright: ignore[reportRedeclaration]
         span = ddtrace.tracer.current_span()
         if span:
             span.set_tags(cast(Any, tags))
 
-    def safe_current_trace_context():  # pyright: ignore[reportRedeclaration]
+    def safe_current_trace_context() -> ddtrace.context.Context | otel_trace.SpanContext | None:  # pyright: ignore[reportRedeclaration]
         return ddtrace.tracer.current_trace_context()
 
-    def safe_activate_trace_context(
-        ctx: ddtrace.context.Context | ddtrace.Span | None,  # pyright: ignore[reportPrivateImportUsage]
-    ) -> None:
-        ddtrace.tracer.context_provider.activate(ctx)
+    @contextlib.contextmanager
+    def safe_activate_trace_context(  # pyright: ignore[reportRedeclaration]
+        ctx: ddtrace.context.Context
+        | ddtrace.Span
+        | otel_trace.SpanContext
+        | None,  # pyright: ignore[reportPrivateImportUsage]
+    ):
+        if isinstance(ctx, ddtrace.context.Context) or isinstance(ctx, ddtrace.Span):
+            ddtrace.tracer.context_provider.activate(ctx)
+        yield
 
     def add_trace_headers(  # pyright: ignore[reportRedeclaration]
         input_headers: None | dict[str, str]
     ) -> dict[str, str]:
-        if input_headers is None:
-            input_headers = dict[str, str]()
-        headers = dict(input_headers)
+        headers: dict[str, str] = dict(input_headers if input_headers is not None else {})
         span = ddtrace.tracer.current_span()
         if span:
+            span.context.sampling_priority = 2
             span.set_tags({ddtrace.constants.SAMPLING_PRIORITY_KEY: 2})  # Ensure that sampling is enabled
             HTTPPropagator.inject(span.context, headers)
         return headers
 
 else:
+    _logger.debug("no trace packages found, tracing will not work")
 
-    def safe_set_gauge(gauge: str, value: int | float):
+    @contextlib.contextmanager
+    def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):  # pyright: ignore[reportRedeclaration]
+        yield
+
+    def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):  # pyright: ignore[reportRedeclaration]
         pass
 
-    def safe_incr(counter: str, value: int | float, tags: list[str] | None = None):
+    def safe_add_tags(tags: Mapping[str, str]):  # pyright: ignore[reportRedeclaration]
         pass
 
+    def safe_current_trace_context() -> ddtrace.context.Context | otel_trace.SpanContext | None:  # pyright: ignore[reportRedeclaration]
+        return
+
     @contextlib.contextmanager
-    def safe_trace(span_id: str, attributes: Mapping[str, str] | None = None):
+    def safe_activate_trace_context(  # pyright: ignore[reportRedeclaration]
+        ctx: ddtrace.context.Context
+        | ddtrace.Span
+        | otel_trace.Context
+        | otel_trace.SpanContext
+        | None,  # pyright: ignore[reportPrivateImportUsage]
+    ):
         yield
 
-    def safe_add_metrics(metrics: Mapping[str, Union[int, float]]):
-        pass
+    def add_trace_headers(headers: None | dict[str, str]) -> dict[str, str]:  # pyright: ignore[reportRedeclaration]
+        if headers is None:
+            return {}
+        return headers
 
-    def safe_add_tags(tags: Mapping[str, str]):
-        pass
 
-    def safe_current_trace_context():
-        return
+if can_use_datadog_statsd:
+    from datadog.dogstatsd.base import statsd
 
-    def safe_activate_trace_context(
-        ctx: ddtrace.context.Context | ddtrace.Span | None,  # pyright: ignore[reportPrivateImportUsage]
-    ) -> None:
-        pass
+    def safe_set_gauge(gauge: str, value: int | float):
+        statsd.gauge(gauge, value)
+
+    def safe_incr(counter: str, value: int | float, tags: list[str] | None = None):
+        statsd.increment(counter, value, tags)
 
     def safe_distribution(counter: str, value: int | float, tags: list[str] | None = None):
+        statsd.distribution(counter, value, tags)
+
+else:
+
+    def safe_set_gauge(gauge: str, value: int | float):
         pass
 
-    def add_trace_headers(headers: None | dict[str, str]) -> dict[str, str]:
-        if headers is None:
-            return {}
-        return headers
+    def safe_incr(counter: str, value: int | float, tags: list[str] | None = None):
+        pass
+
+    def safe_distribution(counter: str, value: int | float, tags: list[str] | None = None):
+        pass
 
 
 class PerfTimer:
@@ -153,54 +227,70 @@ def configure_tracing(default_service_name: str):
 
     _logger = get_logger(__name__)
 
-    if not can_use_ddtrace:
-        _logger.warning("ddtrace is not installed")
-        return
-
-    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,
+    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]
-    )
+configure_tracing("chalkpy")

{chalkpy-2.90.1.dist-info → chalkpy-2.95.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.90.1
+Version: 2.95.3
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai
@@ -24,7 +24,7 @@ Requires-Dist: cattrs<25,>=22.1.0
 Requires-Dist: dataclasses_json>=0.5.7
 Requires-Dist: executing<3,>=1.2.0
 Requires-Dist: googleapis-common-protos>=1.56.0
-Requires-Dist: grpcio<2,>=1.59.2
+Requires-Dist: grpcio<2,>=1.63.0
 Requires-Dist: ipywidgets>=8.0.6
 Requires-Dist: isodate<0.8,>=0.6.1
 Requires-Dist: numpy<3
@@ -211,9 +211,8 @@ Requires-Dist: s3fs; extra == "redshift"
 Requires-Dist: sqlalchemy[asyncio]<2,>=1.4.26; extra == "redshift"
 Requires-Dist: sqlglot<21.2.0,>=19.0.0; extra == "redshift"
 Provides-Extra: mssql
-Requires-Dist: pymssql<3,>=2.2.0; extra == "mssql"
+Requires-Dist: azure-identity<2,>=1.12.0; extra == "mssql"
 Requires-Dist: pyodbc<6,>=4.0.0; extra == "mssql"
-Requires-Dist: aioodbc<1,>=0.5.0; extra == "mssql"
 Requires-Dist: adlfs; extra == "mssql"
 Requires-Dist: duckdb<1.2.0,>=0.6; extra == "mssql"
 Requires-Dist: fsspec; extra == "mssql"
@@ -284,8 +283,8 @@ Provides-Extra: all
 Requires-Dist: PyAthena>=3.0.0; extra == "all"
 Requires-Dist: adlfs; extra == "all"
 Requires-Dist: aiomysql<0.3,>=0.1.1; extra == "all"
-Requires-Dist: aioodbc<1,>=0.5.0; extra == "all"
 Requires-Dist: aiosqlite<0.21,>=0.19.0; extra == "all"
+Requires-Dist: azure-identity<2,>=1.12.0; extra == "all"
 Requires-Dist: boto3; extra == "all"
 Requires-Dist: chalk-sqlalchemy-redshift<0.9,>=0.8.11; extra == "all"
 Requires-Dist: clickhouse-driver; extra == "all"
@@ -309,7 +308,6 @@ Requires-Dist: polars[timezone]!=1.0,!=1.1,!=1.10,!=1.11,!=1.12,!=1.13,!=1.14,!=
 Requires-Dist: psycopg2<3,>=2.9.4; extra == "all"
 Requires-Dist: psycopg[binary]<3.3,>=3.1.9; extra == "all"
 Requires-Dist: pydynamodb<0.8,>=0.6; extra == "all"
-Requires-Dist: pymssql<3,>=2.2.0; extra == "all"
 Requires-Dist: pymysql<2,>=1.0.2; extra == "all"
 Requires-Dist: pyodbc<6,>=4.0.0; extra == "all"
 Requires-Dist: python-json-logger<4.0.0,>=3.0.0; extra == "all"