feat(sessions): Improved control over sessions (getsentry#646)

mitsuhiko · web-flow · commit 8d475f9aa27e · 2020-03-12T15:51:16.000+01:00
diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py
@@ -320,7 +320,7 @@ def capture_event(
 
         # whenever we capture an event we also check if the session needs
         # to be updated based on that information.
-        session = scope.session if scope else None
+        session = scope._session if scope else None
         if session:
             self._update_session_from_event(session, event)
 
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -567,27 +567,47 @@ def inner():
 
         return inner()
 
+    def start_session(self):
+        # type: (...) -> None
+        """Starts a new session."""
+        self.end_session()
+        client, scope = self._stack[-1]
+        scope._session = Session(
+            release=client.options["release"] if client else None,
+            environment=client.options["environment"] if client else None,
+            user=scope._user,
+        )
+
     def end_session(self):
         # type: (...) -> None
         """Ends the current session if there is one."""
         client, scope = self._stack[-1]
-        session = scope.session
+        session = scope._session
         if session is not None:
             session.close()
             if client is not None:
                 client.capture_session(session)
-        self._stack[-1][1].session = None
+        self._stack[-1][1]._session = None
 
-    def start_session(self):
+    def stop_auto_session_tracking(self):
         # type: (...) -> None
-        """Starts a new session."""
+        """Stops automatic session tracking.
+
+        This temporarily session tracking for the current scope when called.
+        To resume session tracking call `resume_auto_session_tracking`.
+        """
         self.end_session()
         client, scope = self._stack[-1]
-        scope.session = Session(
-            release=client.options["release"] if client else None,
-            environment=client.options["environment"] if client else None,
-            user=scope._user,
-        )
+        scope._force_auto_session_tracking = False
+
+    def resume_auto_session_tracking(self):
+        # type: (...) -> None
+        """Resumes automatic session tracking for the current scope if
+        disabled earlier.  This requires that generally automatic session
+        tracking is enabled.
+        """
+        client, scope = self._stack[-1]
+        scope._force_auto_session_tracking = None
 
     def flush(
         self,
diff --git a/sentry_sdk/integrations/atexit.py b/sentry_sdk/integrations/atexit.py
@@ -54,6 +54,9 @@ def _shutdown():
             if integration is not None:
                 logger.debug("atexit: shutting down client")
 
+                # If there is a session on the hub, close it now.
+                hub.end_session()
+
                 # If an integration is there, a client has to be there.
                 client = hub.client  # type: Any
                 client.close(callback=integration.callback)
diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
@@ -87,6 +87,7 @@ class Scope(object):
         "_should_capture",
         "_span",
         "_session",
+        "_force_auto_session_tracking",
     )
 
     def __init__(self):
@@ -114,6 +115,7 @@ def clear(self):
 
         self._span = None  # type: Optional[Span]
         self._session = None  # type: Optional[Session]
+        self._force_auto_session_tracking = None  # type: Optional[bool]
 
     @_attr_setter
     def level(self, value):
@@ -169,17 +171,6 @@ def span(self, span):
             if span_transaction:
                 self._transaction = span_transaction
 
-    @property
-    def session(self):
-        # type: () -> Optional[Session]
-        """Get/set current tracing session."""
-        return self._session
-
-    @session.setter
-    def session(self, session):
-        # type: (Optional[Session]) -> None
-        self._session = session
-
     def set_tag(
         self,
         key,  # type: str
@@ -404,6 +395,7 @@ def __copy__(self):
         rv._should_capture = self._should_capture
         rv._span = self._span
         rv._session = self._session
+        rv._force_auto_session_tracking = self._force_auto_session_tracking
 
         return rv
 
diff --git a/sentry_sdk/sessions.py b/sentry_sdk/sessions.py
@@ -20,11 +20,25 @@
     from sentry_sdk._types import SessionStatus
 
 
+def is_auto_session_tracking_enabled(hub=None):
+    # type: (Optional[sentry_sdk.Hub]) -> bool
+    """Utility function to find out if session tracking is enabled."""
+    if hub is None:
+        hub = sentry_sdk.Hub.current
+    should_track = hub.scope._force_auto_session_tracking
+    if should_track is None:
+        exp = hub.client.options["_experiments"] if hub.client else {}
+        should_track = exp.get("auto_session_tracking")
+    return should_track
+
+
 @contextmanager
-def auto_session_tracking(hub):
-    # type: (sentry_sdk.Hub) -> Generator[None, None, None]
-    exp = hub.client.options["_experiments"] if hub.client else {}
-    should_track = exp.get("auto_session_tracking")
+def auto_session_tracking(hub=None):
+    # type: (Optional[sentry_sdk.Hub]) -> Generator[None, None, None]
+    """Starts and stops a session automatically around a block."""
+    if hub is None:
+        hub = sentry_sdk.Hub.current
+    should_track = is_auto_session_tracking_enabled(hub)
     if should_track:
         hub.start_session()
     try:
