246 lines
8.0 KiB
Python
246 lines
8.0 KiB
Python
import contextvars
|
|
from contextlib import contextmanager
|
|
from typing import Optional, Any, Callable, Dict, TypeVar, cast
|
|
|
|
|
|
class ContextScope:
|
|
def __init__(
|
|
self,
|
|
parent=None,
|
|
fresh: bool = False,
|
|
capture_exceptions: bool = True,
|
|
):
|
|
self.parent = parent
|
|
self.fresh = fresh
|
|
self.capture_exceptions = capture_exceptions
|
|
self.session_id: Optional[str] = None
|
|
self.distinct_id: Optional[str] = None
|
|
self.tags: Dict[str, Any] = {}
|
|
|
|
def set_session_id(self, session_id: str):
|
|
self.session_id = session_id
|
|
|
|
def set_distinct_id(self, distinct_id: str):
|
|
self.distinct_id = distinct_id
|
|
|
|
def add_tag(self, key: str, value: Any):
|
|
self.tags[key] = value
|
|
|
|
def get_parent(self):
|
|
return self.parent
|
|
|
|
def get_session_id(self) -> Optional[str]:
|
|
if self.session_id is not None:
|
|
return self.session_id
|
|
if self.parent is not None and not self.fresh:
|
|
return self.parent.get_session_id()
|
|
return None
|
|
|
|
def get_distinct_id(self) -> Optional[str]:
|
|
if self.distinct_id is not None:
|
|
return self.distinct_id
|
|
if self.parent is not None and not self.fresh:
|
|
return self.parent.get_distinct_id()
|
|
return None
|
|
|
|
def collect_tags(self) -> Dict[str, Any]:
|
|
tags = self.tags.copy()
|
|
if self.parent and not self.fresh:
|
|
# We want child tags to take precedence over parent tags,
|
|
# so we can't use a simple update here, instead collecting
|
|
# the parent tags and then updating with the child tags.
|
|
new_tags = self.parent.collect_tags()
|
|
tags.update(new_tags)
|
|
return tags
|
|
|
|
|
|
_context_stack: contextvars.ContextVar[Optional[ContextScope]] = contextvars.ContextVar(
|
|
"posthog_context_stack", default=None
|
|
)
|
|
|
|
|
|
def _get_current_context() -> Optional[ContextScope]:
|
|
return _context_stack.get()
|
|
|
|
|
|
@contextmanager
|
|
def new_context(fresh=False, capture_exceptions=True):
|
|
"""
|
|
Create a new context scope that will be active for the duration of the with block.
|
|
Any tags set within this scope will be isolated to this context. Any exceptions raised
|
|
or events captured within the context will be tagged with the context tags.
|
|
|
|
Args:
|
|
fresh: Whether to start with a fresh context (default: False).
|
|
If False, inherits tags, identity and session id's from parent context.
|
|
If True, starts with no state
|
|
capture_exceptions: Whether to capture exceptions raised within the context (default: True).
|
|
If True, captures exceptions and tags them with the context tags before propagating them.
|
|
If False, exceptions will propagate without being tagged or captured.
|
|
|
|
Examples:
|
|
# Inherit parent context tags
|
|
with posthog.new_context():
|
|
posthog.tag("request_id", "123")
|
|
# Both this event and the exception will be tagged with the context tags
|
|
posthog.capture("event_name", {"property": "value"})
|
|
raise ValueError("Something went wrong")
|
|
|
|
# Start with fresh context (no inherited tags)
|
|
with posthog.new_context(fresh=True):
|
|
posthog.tag("request_id", "123")
|
|
# Both this event and the exception will be tagged with the context tags
|
|
posthog.capture("event_name", {"property": "value"})
|
|
raise ValueError("Something went wrong")
|
|
|
|
"""
|
|
from posthog import capture_exception
|
|
|
|
current_context = _get_current_context()
|
|
new_context = ContextScope(current_context, fresh, capture_exceptions)
|
|
_context_stack.set(new_context)
|
|
|
|
try:
|
|
yield
|
|
except Exception as e:
|
|
if new_context.capture_exceptions:
|
|
capture_exception(e)
|
|
raise
|
|
finally:
|
|
_context_stack.set(new_context.get_parent())
|
|
|
|
|
|
def tag(key: str, value: Any) -> None:
|
|
"""
|
|
Add a tag to the current context.
|
|
|
|
Args:
|
|
key: The tag key
|
|
value: The tag value
|
|
|
|
Example:
|
|
posthog.tag("user_id", "123")
|
|
"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
current_context.add_tag(key, value)
|
|
|
|
|
|
# NOTE: we should probably also remove this - there's no reason for the user to ever
|
|
# need to manually interact with the current tag set
|
|
def get_tags() -> Dict[str, Any]:
|
|
"""
|
|
Get all tags from the current context. Note, modifying
|
|
the returned dictionary will not affect the current context.
|
|
|
|
Returns:
|
|
Dict of all tags in the current context
|
|
"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
return current_context.collect_tags()
|
|
return {}
|
|
|
|
|
|
# NOTE: We should probably remove this function - the way to clear scope context
|
|
# is by entering a new, fresh context, rather than by clearing the tags or other
|
|
# scope data directly.
|
|
def clear_tags() -> None:
|
|
"""Clear all tags in the current context. Does not clear parent tags"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
current_context.tags.clear()
|
|
|
|
|
|
def identify_context(distinct_id: str) -> None:
|
|
"""
|
|
Identify the current context with a distinct ID, associating all events captured in this or
|
|
child contexts with the given distinct ID (unless identify_context is called again). This is overridden by
|
|
distinct id's passed directly to posthog.capture and related methods (identify, set etc). Entering a
|
|
fresh context will clear the context-level distinct ID.
|
|
|
|
Args:
|
|
distinct_id: The distinct ID to associate with the current context and its children.
|
|
"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
current_context.set_distinct_id(distinct_id)
|
|
|
|
|
|
def set_context_session(session_id: str) -> None:
|
|
"""
|
|
Set the session ID for the current context, associating all events captured in this or
|
|
child contexts with the given session ID (unless set_context_session is called again).
|
|
Entering a fresh context will clear the context-level session ID.
|
|
|
|
Args:
|
|
session_id: The session ID to associate with the current context and its children. See https://posthog.com/docs/data/sessions
|
|
"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
current_context.set_session_id(session_id)
|
|
|
|
|
|
def get_context_session_id() -> Optional[str]:
|
|
"""
|
|
Get the session ID for the current context.
|
|
|
|
Returns:
|
|
The session ID if set, None otherwise
|
|
"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
return current_context.get_session_id()
|
|
return None
|
|
|
|
|
|
def get_context_distinct_id() -> Optional[str]:
|
|
"""
|
|
Get the distinct ID for the current context.
|
|
|
|
Returns:
|
|
The distinct ID if set, None otherwise
|
|
"""
|
|
current_context = _get_current_context()
|
|
if current_context:
|
|
return current_context.get_distinct_id()
|
|
return None
|
|
|
|
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
|
|
|
|
def scoped(fresh=False, capture_exceptions=True):
|
|
"""
|
|
Decorator that creates a new context for the function. Simply wraps
|
|
the function in a with posthog.new_context(): block.
|
|
|
|
Args:
|
|
fresh: Whether to start with a fresh context (default: False)
|
|
capture_exceptions: Whether to capture and track exceptions with posthog error tracking (default: True)
|
|
|
|
Example:
|
|
@posthog.scoped()
|
|
def process_payment(payment_id):
|
|
posthog.tag("payment_id", payment_id)
|
|
posthog.tag("payment_method", "credit_card")
|
|
|
|
# This event will be captured with tags
|
|
posthog.capture("payment_started")
|
|
# If this raises an exception, it will be captured with tags
|
|
# and then re-raised
|
|
some_risky_function()
|
|
"""
|
|
|
|
def decorator(func: F) -> F:
|
|
from functools import wraps
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
with new_context(fresh=fresh, capture_exceptions=capture_exceptions):
|
|
return func(*args, **kwargs)
|
|
|
|
return cast(F, wrapper)
|
|
|
|
return decorator
|