warpzone-sdk 15.1.1__py3-none-any.whl → 15.1.1.dev1__py3-none-any.whl

warpzone/db/client.py

@@ -195,10 +195,12 @@ class WarpzoneDatabaseClient:
             include_validity_period_columns (bool, optional): Whether to include
             validity period columns in the result;
             (`valid_from_time_utc`, `valid_to_time_utc`).
-            Defaults to False.
+            Defaults to False. If set to True while using `columns`-argument, make sure
+            to include these columns in the `columns`-list.
             include_generated_columns (bool, optional): Whether to include generated
             columns in the result; (e.g. `valid_from_time_utc`, `valid_to_time_utc`).
-            Defaults to False.
+            Defaults to False. If set to True while using `columns`-argument, make sure
+            to include these columns in the `columns`-list.
 
         Returns:
             pd.DataFrame: The result of the query.
@@ -250,12 +252,10 @@ class WarpzoneDatabaseClient:
             for field in table.schema().fields:
                 if field.generated_as is not None:
                     generated_cols.append(field.column_name)
-            pd_df = pd_df.drop(columns=generated_cols, errors="ignore")
+            pd_df = pd_df.drop(columns=generated_cols)
 
         # Drop valid-from/to columns
         if not include_validity_period_columns:
-            pd_df = pd_df.drop(
-                columns=["valid_from_time_utc", "valid_to_time_utc"], errors="ignore"
-            )
+            pd_df = pd_df.drop(columns=["valid_from_time_utc", "valid_to_time_utc"])
 
         return pd_df

warpzone/monitor/traces.py

@@ -1,13 +1,16 @@
+import inspect
 import logging
 import os
 import threading
 from contextlib import contextmanager
+from functools import wraps
 from logging import StreamHandler
+from typing import Callable
 
 from azure.monitor.opentelemetry.exporter import AzureMonitorTraceExporter
 from opentelemetry import context, trace
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
-from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace import Tracer, TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.sdk.trace.sampling import ALWAYS_ON
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
@@ -15,7 +18,104 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapProp
 logger = logging.getLogger(__name__)
 logger.addHandler(StreamHandler())
 
-tracer = trace.get_tracer(__name__)
+
+class WarpzoneTracer:
+    """Wrapper around OpenTelemetry tracer with additional trace decorator method"""
+
+    def __init__(self, otel_tracer: Tracer):
+        # Store the original tracer instead of calling super().__init__
+        # because we want to wrap an existing tracer, not create a new one
+        self._tracer = otel_tracer
+
+    def __getattr__(self, name):
+        """Delegate all attributes to the underlying tracer"""
+        return getattr(self._tracer, name)
+
+    def trace_function(
+        self,
+        name: str = None,
+        set_args_as_attributes: bool = False,
+        on_input: Callable = None,
+        on_output: Callable = None,
+    ):
+        """
+        Decorator to trace a function using this tracer object.
+
+        This decorator wraps functions with OpenTelemetry tracing, allowing for:
+        - Automatically create spans for function execution
+        - Customize the span name
+        - Add custom attributes to the span
+        - Add custom logic for inputs and outputs
+
+        Args:
+            name: Optional name for the span. If not provided, uses the function name.
+            set_args_as_attributes: If True, sets function arguments as attributes.
+            on_input: Optional callback called with (span, *args, **kwargs) before.
+            on_output: Optional callback called with (span, result) after.
+
+        Example:
+            # Simple tracing with function name
+            tracer = get_tracer(__name__)
+
+            @tracer.trace_function()
+            def my_function():
+                pass
+
+            # Custom tracing with input/output callbacks
+            @tracer.trace_function(
+                on_input=lambda span, new_data, existing_data, now: (
+                    span.set_attribute("new_records", len(new_data)),
+                    span.set_attribute("existing_records", len(existing_data)),
+                ),
+                on_output=lambda span, result: (
+                    span.set_attribute("merged_records", len(result))
+                ),
+            )
+            def merge_new_and_existing(new_data, existing_data, now):
+                pass
+        """
+
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                span_name = name or func.__name__
+                with self._tracer.start_as_current_span(span_name) as span:
+                    if set_args_as_attributes:
+                        # Get parameter names from function signature
+                        sig = inspect.signature(func)
+                        param_names = list(sig.parameters.keys())
+
+                        # Set positional arguments with their parameter names
+                        for i, arg in enumerate(args):
+                            if i < len(param_names):
+                                span.set_attribute(param_names[i], str(arg))
+                            else:
+                                # Fallback for *args if there are more args than params
+                                span.set_attribute(f"arg_{i}", str(arg))
+
+                        # Set keyword arguments
+                        for key, value in kwargs.items():
+                            span.set_attribute(str(key), str(value))
+
+                    # Call on_input callback if provided
+                    if on_input:
+                        on_input(span, *args, **kwargs)
+
+                    result = func(*args, **kwargs)
+
+                    # Call on_output callback if provided
+                    if on_output:
+                        on_output(span, result)
+
+                    return result
+
+            return wrapper
+
+        return decorator
+
+
+tracer = WarpzoneTracer(trace.get_tracer(__name__))
+
 
 _TRACING_LOCK = threading.Lock()
 TRACING_IS_CONFIGURED = False
@@ -81,8 +181,8 @@ def set_trace_context(trace_parent: str, trace_state: str = ""):
 
 
 def get_tracer(name: str):
-    tracer = trace.get_tracer(name)
-    return tracer
+    otel_tracer = trace.get_tracer(name)
+    return WarpzoneTracer(otel_tracer)
 
 
 def get_current_diagnostic_id() -> str:
@@ -111,14 +211,14 @@ _SB_TRACE_NAMESPACE = "Microsoft.ServiceBus"
 
 
 @contextmanager
-def servicebus_send_span(subject: str) -> trace.Span:
+def servicebus_send_span(subject: str):
     """Start span for Service Bus message tracing.
 
     Args:
         subject: The message subject (used as span name for easy identification)
 
     Yields:
-        trace.Span: the span
+        Span: the span
     """
     with tracer.start_as_current_span(
         subject, kind=trace.SpanKind.PRODUCER

{warpzone_sdk-15.1.1.dist-info → warpzone_sdk-15.1.1.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: warpzone-sdk
-Version: 15.1.1
+Version: 15.1.1.dev1
 Summary: The main objective of this package is to centralize logic used to interact with Azure Functions, Azure Service Bus and Azure Table Storage
 Author: Team Enigma
 Author-email: enigma@energinet.dk

{warpzone_sdk-15.1.1.dist-info → warpzone_sdk-15.1.1.dev1.dist-info}/RECORD

@@ -2,7 +2,7 @@ warpzone/__init__.py,sha256=Ay7znIadokCdHHqsPlWAqhdg-zdWhXHNtxVinViYm7o,1533
 warpzone/blobstorage/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
 warpzone/blobstorage/client.py,sha256=YwDV83acoCeHS_D_ydsTYwnf56rSTy9CKpsxqeoXmBs,4638
 warpzone/db/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
-warpzone/db/client.py,sha256=TdzfpgC81floe0WM7JGNbQpVz1esj1uXjnIZ9kTZmlY,9518
+warpzone/db/client.py,sha256=tqY_ANThnO72rBaB4sydigLhlSFwCmadg_6rVZwjO8M,9688
 warpzone/deltastorage/__init__.py,sha256=cV8sGT2N_N5Z-E179NMW5O7q3FUDrb3j5f-yVNlNPv0,152
 warpzone/deltastorage/data_types.py,sha256=tWjLO_0ig7-tYxSHShvrd0znA7FFDnS-wBuFClUQG2U,2059
 warpzone/deltastorage/generated_columns.py,sha256=Dr_bihM7v9JKCgBXxc3JQC2P2mUGobXsReRKOl6jDO4,4765
@@ -29,7 +29,7 @@ warpzone/healthchecks/__init__.py,sha256=9gc_Mt2szs8sDSwy0V4l3JZ6d9hX41xTpZCkDP2
 warpzone/healthchecks/model.py,sha256=mM7DnrirLbUpBPPfi82MUPP654D0eOR2_F65TmzsPD0,1187
 warpzone/monitor/__init__.py,sha256=gXT2cxR4tlZER54zd7D49ZQBVyitLaqj13_cUoILuyM,109
 warpzone/monitor/logs.py,sha256=q3SUQCtG1ii0B9GkVs2l8kgZ5b5bI2qI6L0SoOW2QTY,3147
-warpzone/monitor/traces.py,sha256=Xc_po1LxJFy5jtNWxIVphIInl_d89Zw3Rb21PsdQhQA,4170
+warpzone/monitor/traces.py,sha256=IJf5CjWeh3Th2GPwvmdlhJRdVoPHgHFuZeOw2VeNu1A,7962
 warpzone/servicebus/data/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
 warpzone/servicebus/data/client.py,sha256=zECS3JwedhYnDk8PntYgIYpBF_uu9YN38KzpPFK7CKs,6511
 warpzone/servicebus/events/__init__.py,sha256=lnc0uiaGLF0qMi_rWhCpRSFvaj0CJEiMCAl6Yqn1ZiA,21
@@ -52,6 +52,6 @@ warpzone/tools/copy.py,sha256=5fddotMZkXZO8avzUbGOhvs0cp8mce95pNpy0oPVjnQ,2596
 warpzone/transform/__init__.py,sha256=ruGa7tl-v4ndlWpULE1jSGU_a4_iRc3V6eyNr5xKP9E,27
 warpzone/transform/data.py,sha256=Abb8PcrgMbbNCJkkIUdtrTHdlY0OfXid387qw1nDpFY,2362
 warpzone/transform/schema.py,sha256=nbSQtDMvXkyqGKuwhuFCF0WsEDsaNyoPYpMKvbsKlv8,2423
-warpzone_sdk-15.1.1.dist-info/METADATA,sha256=6WeGlAAYB3qDhdzNl4bS3GFiEzmAasRu-kakPJ3LqSM,7279
-warpzone_sdk-15.1.1.dist-info/WHEEL,sha256=kJCRJT_g0adfAJzTx2GUMmS80rTJIVHRCfG0DQgLq3o,88
-warpzone_sdk-15.1.1.dist-info/RECORD,,
+warpzone_sdk-15.1.1.dev1.dist-info/METADATA,sha256=cGg0TkCeji-1k162MizWD0MvGbG2k5xBY42NkeXol-c,7284
+warpzone_sdk-15.1.1.dev1.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
+warpzone_sdk-15.1.1.dev1.dist-info/RECORD,,

{warpzone_sdk-15.1.1.dist-info → warpzone_sdk-15.1.1.dev1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.3.1
+Generator: poetry-core 2.3.0
 Root-Is-Purelib: true
 Tag: py3-none-any