Implement heartbeats for Python

.changesets/implement-heartbeats.md

@@ -0,0 +1,32 @@
+---
+bump: "minor"
+type: "add"
+---
+
+_Heartbeats are currently only available to beta testers. If you are interested in trying it out, [send an email to [email protected]](mailto:[email protected]?subject=Heartbeat%20beta)!_
+
+---
+
+Add heartbeats support. You can use the `heartbeat` function to send heartbeats directly from your code, to track the execution of certain processes:
+
+```python
+from appsignal import heartbeat
+
+def send_invoices():
+  # ... your code here ...
+  heartbeat("send_invoices")
+```
+
+It is also possible to pass a defined function as an argument to the `heartbeat`
+function:
+
+```python
+def send_invoices():
+  # ... your code here ...
+
+heartbeat("send_invoices", send_invoices)
+```
+
+If an exception is raised within the function, the finish event will not be
+reported to AppSignal, triggering a notification about the missing heartbeat.
+The exception will be raised outside of the heartbeat function.

conftest.py

@@ -17,6 +17,7 @@
 
 from appsignal import probes
 from appsignal.agent import agent
+from appsignal.client import _reset_client
 from appsignal.opentelemetry import METRICS_PREFERRED_TEMPORALITY
 
 
@@ -100,6 +101,11 @@ def reset_agent_active_state():
     agent.active = False
 
 
+@pytest.fixture(scope="function", autouse=True)
+def reset_global_client():
+    _reset_client()
+
+
 @pytest.fixture(scope="function", autouse=True)
 def stop_agent():
     tmp_path = "/tmp" if platform.system() == "Darwin" else tempfile.gettempdir()

pyproject.toml

@@ -91,6 +91,7 @@ dependencies = [
   "types-deprecated",
   "types-requests",
   "django-stubs",
+  "pytest",
 ]
 
 [tool.hatch.envs.lint.scripts]

src/appsignal/__init__.py

@@ -4,6 +4,7 @@
 from runpy import run_path
 
 from .client import Client as Appsignal
+from .heartbeat import Heartbeat, heartbeat
 from .metrics import add_distribution_value, increment_counter, set_gauge
 from .tracing import (
     send_error,
@@ -42,6 +43,8 @@
     "increment_counter",
     "set_gauge",
     "add_distribution_value",
+    "Heartbeat",
+    "heartbeat",
     "start",
 ]
 

src/appsignal/cli/diagnose.py

@@ -5,18 +5,16 @@
 import json
 import os
 import platform
-import urllib
 from argparse import ArgumentParser
 from pathlib import Path
 from sys import stderr
 from typing import Any
 
-import requests
-
 from ..__about__ import __version__
 from ..agent import Agent
 from ..config import Config
 from ..push_api_key_validator import PushApiKeyValidator
+from ..transmitter import transmit
 from .command import AppsignalCLICommand
 
 
@@ -377,19 +375,9 @@ def _report_prompt(self) -> bool | None:
         return None
 
     def _send_diagnose_report(self) -> None:
-        params = urllib.parse.urlencode(
-            {
-                "api_key": self.config.option("push_api_key"),
-                "name": self.config.option("name"),
-                "environment": self.config.option("environment"),
-                "hostname": self.config.option("hostname") or "",
-            }
-        )
-
-        endpoint = self.config.option("diagnose_endpoint")
-        url = f"{endpoint}?{params}"
+        url = self.config.option("diagnose_endpoint")
 
-        response = requests.post(url, json={"diagnose": self.report})
+        response = transmit(url, json={"diagnose": self.report}, config=self.config)
 
         status = response.status_code
         if status == 200:

src/appsignal/client.py

@@ -15,6 +15,14 @@
     from typing_extensions import Unpack
 
 
+_client: Client | None = None
+
+
+def _reset_client() -> None:
+    global _client
+    _client = None
+
+
 class Client:
     _logger: Logger
     _config: Config
@@ -28,12 +36,22 @@ class Client:
     }
 
     def __init__(self, **options: Unpack[Options]) -> None:
+        global _client
+
         self._config = Config(options)
         self._start_logger()
+        _client = self
 
         if not self._config.is_active():
             self._logger.info("AppSignal not starting: no active config found")
 
+    @classmethod
+    def config(cls) -> Config | None:
+        if _client is None:
+            return None
+
+        return _client._config
+
     def start(self) -> None:
         if self._config.is_active():
             self._logger.info("Starting AppSignal")

src/appsignal/config.py

@@ -3,24 +3,11 @@
 import os
 import platform
 import tempfile
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    ClassVar,
-    List,
-    Literal,
-    TypedDict,
-    cast,
-    get_args,
-)
+from typing import Any, ClassVar, List, Literal, TypedDict, cast, get_args
 
 from .__about__ import __version__
 
 
-if TYPE_CHECKING:
-    pass
-
-
 class Options(TypedDict, total=False):
     active: bool | None
     app_path: str | None
@@ -50,6 +37,7 @@ class Options(TypedDict, total=False):
     log: str | None
     log_level: str | None
     log_path: str | None
+    logging_endpoint: str | None
     opentelemetry_port: str | int | None
     name: str | None
     push_api_key: str | None
@@ -90,6 +78,7 @@ class Config:
         files_world_accessible=True,
         log="file",
         log_level="info",
+        logging_endpoint="https://appsignal-endpoint.net",
         opentelemetry_port=8099,
         send_environment_metadata=True,
         send_params=True,
@@ -185,6 +174,7 @@ def load_from_environment() -> Options:
             log=os.environ.get("APPSIGNAL_LOG"),
             log_level=os.environ.get("APPSIGNAL_LOG_LEVEL"),
             log_path=os.environ.get("APPSIGNAL_LOG_PATH"),
+            logging_endpoint=os.environ.get("APPSIGNAL_LOGGING_ENDPOINT"),
             opentelemetry_port=os.environ.get("APPSIGNAL_OPENTELEMETRY_PORT"),
             name=os.environ.get("APPSIGNAL_APP_NAME"),
             push_api_key=os.environ.get("APPSIGNAL_PUSH_API_KEY"),
@@ -252,6 +242,7 @@ def set_private_environ(self) -> None:
             "_APPSIGNAL_LOG": options.get("log"),
             "_APPSIGNAL_LOG_LEVEL": options.get("log_level"),
             "_APPSIGNAL_LOG_FILE_PATH": self.log_file_path(),
+            "_APPSIGNAL_LOGGING_ENDPOINT": options.get("logging_endpoint"),
             "_APPSIGNAL_OPENTELEMETRY_PORT": options.get("opentelemetry_port"),
             "_APPSIGNAL_PUSH_API_KEY": options.get("push_api_key"),
             "_APPSIGNAL_PUSH_API_ENDPOINT": options.get("endpoint"),

src/appsignal/heartbeat.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import logging
+from binascii import hexlify
+from os import urandom
+from time import time
+from typing import Any, Callable, Literal, TypedDict, TypeVar, Union
+
+from .client import Client
+from .config import Config
+from .transmitter import transmit
+
+
+T = TypeVar("T")
+
+EventKind = Union[Literal["start"], Literal["finish"]]
+
+
+class Event(TypedDict):
+    name: str
+    id: str
+    kind: EventKind
+    timestamp: int
+
+
+class Heartbeat:
+    name: str
+    id: str
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.id = hexlify(urandom(8)).decode("utf-8")
+
+    def _event(self, kind: EventKind) -> Event:
+        return Event(name=self.name, id=self.id, kind=kind, timestamp=int(time()))
+
+    def _transmit(self, event: Event) -> None:
+        config = Client.config() or Config()
+        logger = logging.getLogger("appsignal")
+
+        if not config.is_active():
+            logger.debug("AppSignal not active, not transmitting heartbeat event")
+            return
+
+        url = f"{config.option('logging_endpoint')}/heartbeats/json"
+        try:
+            response = transmit(url, json=event)
+            if 200 <= response.status_code <= 299:
+                logger.debug(
+                    f"Transmitted heartbeat `{event['name']}` ({event['id']}) "
+                    f"{event['kind']} event"
+                )
+            else:
+                logger.error(
+                    "Failed to transmit heartbeat event: "
+                    f"status code was {response.status_code}"
+                )
+        except Exception as e:
+            logger.error(f"Failed to transmit heartbeat event: {e}")
+
+    def start(self) -> None:
+        self._transmit(self._event("start"))
+
+    def finish(self) -> None:
+        self._transmit(self._event("finish"))
+
+    def __enter__(self) -> None:
+        self.start()
+
+    def __exit__(
+        self, exc_type: Any = None, exc_value: Any = None, traceback: Any = None
+    ) -> Literal[False]:
+        if exc_type is None:
+            self.finish()
+
+        return False
+
+
+def heartbeat(name: str, fn: Callable[[], T] | None = None) -> None | T:
+    heartbeat = Heartbeat(name)
+    output = None
+
+    if fn is not None:
+        heartbeat.start()
+        output = fn()
+
+    heartbeat.finish()
+    return output

src/appsignal/push_api_key_validator.py

@@ -1,31 +1,14 @@
-import urllib
-
-import requests
-
 from appsignal.config import Config
+from appsignal.transmitter import transmit
 
 
 class PushApiKeyValidator:
     @staticmethod
     def validate(config: Config) -> str:
         endpoint = config.option("endpoint")
-        params = urllib.parse.urlencode(
-            {
-                "api_key": config.option("push_api_key"),
-                "name": config.option("name"),
-                "environment": config.option("environment"),
-                "hostname": config.option("hostname") or "",
-            }
-        )
-
-        url = f"{endpoint}/1/auth?{params}"
-        proxies = {}
-        if config.option("http_proxy"):
-            proxies["http"] = config.option("http_proxy")
-            proxies["https"] = config.option("http_proxy")
+        url = f"{endpoint}/1/auth"
 
-        cert = config.option("ca_file_path")
-        response = requests.post(url, proxies=proxies, verify=cert)
+        response = transmit(url, config=config)
 
         if response.status_code == 200:
             return "valid"

src/appsignal/transmitter.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import urllib
+from typing import TYPE_CHECKING, Any
+
+import requests
+
+from appsignal.client import Client
+from appsignal.config import Config
+
+
+if TYPE_CHECKING:
+    from requests import Response
+
+
+def transmit(
+    url: str, json: Any | None = None, config: Config | None = None
+) -> Response:
+    if config is None:
+        config = Client.config() or Config()
+
+    params = urllib.parse.urlencode(
+        {
+            "api_key": config.option("push_api_key") or "",
+            "name": config.option("name") or "",
+            "environment": config.option("environment") or "",
+            "hostname": config.option("hostname") or "",
+        }
+    )
+
+    url = f"{url}?{params}"
+
+    proxies = {}
+    if config.option("http_proxy"):
+        proxies["http"] = config.option("http_proxy")
+        proxies["https"] = config.option("http_proxy")
+
+    cert = config.option("ca_file_path")
+
+    return requests.post(url, json=json, proxies=proxies, verify=cert)