2025-07-01

This commit is contained in:
2026-03-17 14:30:01 -06:00
parent f9a22056dd
commit 62b5978595
4579 changed files with 1257472 additions and 0 deletions
@@ -0,0 +1,4 @@
[flake8]
extend-ignore = E501
exclude =
__pycache__
@@ -0,0 +1,815 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
from concurrent.futures import Future
from datetime import datetime
from functools import lru_cache
from typing import Any, Callable, Dict, List, Optional, Tuple
from errno import EACCES, ENOSPC
import functools
import os
import json
import webbrowser
from .assets import (AssetType,
AssetData,
ModelType,
SIZES)
from .user import (PoliigonUser,
PoliigonSubscription)
from .plan_manager import SubscriptionState, PoliigonPlanUpgradeManager
from . import api
from . import asset_index
from . import env
from .logger import (DEBUG, # noqa F401, allowing downstream const usage
ERROR,
INFO,
get_addon_logger,
NOT_SET,
WARNING)
from .notifications import NotificationSystem
from . import settings
from . import updater
from .multilingual import Multilingual
from . import thread_manager as tm
DIR_PATH = os.path.dirname(os.path.abspath(__file__))
RESOURCES_PATH = os.path.join(DIR_PATH, "resources")
class PoliigonAddon():
"""Poliigon addon used for creating base singleton in DCC applications."""
addon_name: str # e.g. poliigon-addon-blender
addon_version: tuple # Current addon version
software_source: str # e.g. blender
software_version: tuple # DCC software version, e.g. (3, 0)
addon_convention: int # Maximum convention supported by DCC implementation
library_paths: List = []
def __init__(self,
addon_name: str,
addon_version: tuple,
software_source: str,
software_version: tuple,
addon_env: env.PoliigonEnvironment,
addon_settings: settings.PoliigonSettings,
addon_convention: int,
addon_supported_model: List[ModelType] = [ModelType.FBX],
language: str = "en-US",
# See ThreadManager.__init__ for signature below,
# e.g. print_exc(fut: Future, key_pool: PoolKeys)
callback_print_exc: Optional[Callable] = None):
self.log_manager = get_addon_logger(env=addon_env)
if addon_env.env_name == "prod":
have_filehandler = False
else:
have_filehandler = True
self.logger = self.log_manager.initialize_logger(
have_filehandler=have_filehandler)
self.logger_api = self.log_manager.initialize_logger(
"API", have_filehandler=have_filehandler)
self.logger_dl = self.log_manager.initialize_logger(
"DL", have_filehandler=have_filehandler)
self.language = language
self.multilingual = Multilingual()
self.multilingual.install_domain(language=self.language,
dir_lang=os.path.join(RESOURCES_PATH, "lang"),
domain="addon-core")
self.addon_name = addon_name
self.addon_version = addon_version
self.software_source = software_source
self.software_version = software_version
self.addon_convention = addon_convention
self.user = None
self.login_error = None
self.api_rc = None # To be set on the DCC side
self.upgrade_manager = PoliigonPlanUpgradeManager(self)
self._env = addon_env
self.set_logger_verbose(verbose=False)
self._settings = addon_settings
self._api = api.PoliigonConnector(
env=self._env,
software=software_source,
logger=self.logger_api
)
self.logger.debug(f"API URL V1: {self._api.api_url}")
self.logger.debug(f"API URL V2: {self._api.api_url_v2}")
if "v1" in self._api.api_url and "apiv1" not in self._api.api_url:
self.logger.warning("Likely you are running with an outdated API V1 URL")
self._api.register_update(
".".join([str(x) for x in addon_version]),
".".join([str(x) for x in software_version])
)
self._tm = tm.ThreadManager(callback_print_exc=callback_print_exc)
self.notify = NotificationSystem(self)
self._api.notification_system = self.notify
self._updater = updater.SoftwareUpdater(
addon_name=addon_name,
addon_version=addon_version,
software_version=software_version,
notification_system=self.notify,
local_json=self._env.local_updater_json
)
self.settings_config = self._settings.config
self.user_addon_dir = os.path.join(
os.path.expanduser("~"),
"Poliigon"
)
self.setup_libraries()
self.categories_path = os.path.join(self.user_addon_dir, "categories.json")
default_asset_index_path = os.path.join(
self.user_addon_dir,
"AssetIndex",
"asset_index.json",
)
self._asset_index = asset_index.AssetIndex(
addon=self,
addon_convention=addon_convention,
path_cache=default_asset_index_path,
addon_supported_model=addon_supported_model,
log=None
)
self.online_previews_path = self.setup_temp_previews_folder()
# TODO(Andreas): Could well be done in constructor itself.
# Yet, it would break DCC implementations, atm.
def init_addon_parameters(
self,
*,
get_optin: Callable,
callback_on_invalidated_token: Callable,
report_message: Callable,
report_exception: Callable,
report_thread: Callable,
status_listener: Callable,
urls_dcc: Dict[str, str],
notify_icon_info: Any,
notify_icon_no_connection: Any,
notify_icon_survey: Any,
notify_icon_warn: Any,
notify_update_body: str
# TODO(Andreas): Once API RC gets instanced here, add:
# page_size_online_assets: int,
# page_size_my_assets: int,
# callback_get_categories_done: Callable,
# callback_get_asset_done: Callable,
# callback_get_user_data_done: Callable,
# callback_get_download_prefs_done: Callable
) -> None:
"""Initializes all parameters of PoliigonAddon."""
self._api.get_optin = get_optin
self._api.set_on_invalidated(callback_on_invalidated_token)
self._api._status_listener = status_listener
self._api.add_poliigon_urls(urls_dcc)
self._api._report_message = report_message
self._api._report_exception = report_exception
self._tm.reporting_callable = report_thread
self.notify.init_icons(
icon_info=notify_icon_info,
icon_no_connection=notify_icon_no_connection,
icon_survey=notify_icon_survey,
icon_warn=notify_icon_warn)
self.notify.addon_params.update_body = notify_update_body
# TODO(Andreas): Once API RC gets instanced in constructor,
# add the following here:
# params = self.api_rc._addon_params
# params.online_assets_chunk_size = page_size_online_assets
# params.my_assets_chunk_size = page_size_my_assets
# params.callback_get_categories_done = callback_get_categories_done
# params.callback_get_asset_done = callback_get_asset_done
# params.callback_get_user_data_done = callback_get_user_data_done
# params.callback_get_download_prefs_done = callback_get_download_prefs_done
# Decorator copied from comment in thread_manager.py
def run_threaded(key_pool: tm.PoolKeys,
max_threads: Optional[int] = None,
foreground: bool = False) -> Callable:
"""Schedule a function to run in a thread of a chosen pool"""
def wrapped_func(func: Callable) -> Callable:
@functools.wraps(func)
def wrapped_func_call(self, *args, **kwargs):
args = (self, ) + args
return self._tm.queue_thread(func, key_pool,
max_threads, foreground,
*args, **kwargs)
return wrapped_func_call
return wrapped_func
def setup_libraries(self):
default_lib_path = os.path.join(self.user_addon_dir, "Library")
multi_dir = self.settings_config["directories"]
primary_lib_path = self.settings_config.get(
"library", "primary", fallback=None)
# If primary lib is not found in settings, set the default path as
# primary, the DCC side should handle the value missing in settings
# (e.g. choose main lib screen)
if primary_lib_path not in [None, ""]:
self.library_paths.append(primary_lib_path)
else:
self.library_paths.append(default_lib_path)
for dir_idx in multi_dir.keys():
path = self.settings_config.get("directories", str(dir_idx))
self.library_paths.append(path)
# TODO(Andreas): why is it called temp?
def setup_temp_previews_folder(self) -> str:
previews_dir = os.path.join(self.user_addon_dir, "OnlinePreviews")
try:
os.makedirs(previews_dir, exist_ok=True)
except Exception:
self.logger.exception(
f"Failed to create directory: {previews_dir}")
# Removing lock temp files for thumbs
for _file in os.listdir(previews_dir):
file_path = os.path.join(previews_dir, _file)
if os.path.isfile(file_path) and _file.endswith("_temp"):
os.remove(file_path)
return previews_dir
def load_categories_from_disk(self) -> Optional[Dict]:
"""Loads categories from disk."""
if not os.path.exists(self.categories_path):
return None
try:
with open(self.categories_path, "r") as file_categories:
category_json = json.load(file_categories)
if not isinstance(category_json, List):
return None
# TODO(Andreas): error handling
# Whatever error we encounter, worst outcome is no cached categories
except OSError as e:
if e.errno == EACCES:
return None
else:
return None
except Exception:
return None
return category_json
def save_categories_to_disk(self, category_json: List) -> None:
"""Stores categories (as received from API) to disk."""
try:
with open(self.categories_path, "w") as file_categories:
json.dump(category_json, file_categories, indent=4)
# TODO(Andreas): error handling
# Whatever error we encounter, worst outcome is no cached categories
except OSError as e:
if e.errno == ENOSPC:
return
elif e.errno == EACCES:
return
else:
return
except Exception:
return
def set_logger_verbose(self, verbose: bool) -> None:
"""To be used by DCC side to set main logger verbosity."""
log_lvl_from_env = NOT_SET
if self._env.config is not None:
log_lvl_from_env = self._env.config.getint(
"DEFAULT", "log_lvl", fallback=NOT_SET)
if log_lvl_from_env != NOT_SET:
self.logger.info(f"Log level forced by env: {log_lvl_from_env}")
return
log_lvl = INFO if verbose else ERROR
self.logger.setLevel(log_lvl)
def is_logged_in(self) -> bool:
"""Returns whether or not the user is currently logged in."""
return self._api.token is not None and not self._api.invalidated
def is_user_invalidated(self) -> bool:
"""Returns whether or not the user token was invalidated."""
return self._api.invalidated
def clear_user_invalidated(self):
"""Clears any invalidation flag for a user."""
self._api.invalidated = False
@run_threaded(tm.PoolKeys.INTERACTIVE)
def log_in_with_credentials(self,
email: str,
password: str,
*,
wait_for_user: bool = False) -> Future:
self.clear_user_invalidated()
req = self._api.log_in(
email,
password
)
if req.ok:
user_data = req.body.get("user", {})
fut = self.create_user(user_data.get("name"), user_data.get("id"))
if wait_for_user:
fut.result(timeout=api.TIMEOUT)
self.login_error = None
else:
self.login_error = req.error
return req
def log_in_with_website(self):
pass
def check_for_survey_notice(
self,
free_user_url: str,
plan_user_url: str,
interval: int,
label: str,
tooltip: str = "",
auto_enqueue: bool = True) -> None:
already_shown = self.settings_config.get(
"user", "survey_notice_shown", fallback=None)
if already_shown not in [None, ""]:
# Never notify again if already did once
return
first_local_asset = self.settings_config.get(
"user", "first_local_asset", fallback=None)
if first_local_asset in ["", None]:
return
def set_user_survey_flag() -> None:
self.settings_config.set(
"user", "survey_notice_shown", str(datetime.now()))
self._settings.save_settings()
first_asset_dl = datetime.strptime(first_local_asset, "%Y-%m-%d %H:%M:%S.%f")
difference = datetime.now() - first_asset_dl
if difference.days >= interval:
self.notify.create_survey(
is_free_user=self.is_free_user(),
tooltip=tooltip,
free_survey_url=free_user_url,
active_survey_url=plan_user_url,
label=label,
auto_enqueue=auto_enqueue,
on_dismiss_callable=set_user_survey_flag
)
@run_threaded(tm.PoolKeys.INTERACTIVE)
def log_out(self):
req = self._api.log_out()
if req.ok:
print("Logout success")
else:
print(req.error)
self._api.token = None
# Clear out user on logout.
self.user = None
def add_library_path(self,
path: str,
primary: bool = True,
update_local_assets: bool = True
) -> None:
if not os.path.isdir(path):
self.logger.info(f"Library Path to be added is not a directory: {path}")
return
elif path in self.library_paths:
if primary:
self.remove_library_path(path)
else:
self.logger.info(f"Library Path to be added is already in the list: {path}")
return
if primary:
if len(self.library_paths) == 0:
self.library_paths = [path]
else:
self.library_paths[0] = path
self.settings_config.set("library", "primary", path)
else:
self.library_paths.append(path)
idx = 0
list_directory_idxs = list(self.settings_config["directories"].keys())
if len(list_directory_idxs) > 0:
idx = int(list_directory_idxs[-1]) + 1
self.settings_config.set("directories", str(idx), path)
self._settings.save_settings()
if update_local_assets:
self._asset_index.update_all_local_assets(self.library_paths)
def remove_library_path(self,
path: str,
update_local_assets: bool = True
) -> None:
if path not in self.library_paths:
self.logger.info(f"Library Path to be removed is not in the list: {path}")
return
self.library_paths.remove(path)
for dir_idx in self.settings_config["directories"].keys():
dir_path = self.settings_config.get("directories", dir_idx)
if dir_path == path:
self.settings_config.remove_option("directories", dir_idx)
self._settings.save_settings()
if update_local_assets:
self._asset_index.flush_is_local()
self._asset_index.update_all_local_assets(self.library_paths)
def replace_library_path(self,
path_old: str,
path_new: str,
primary: bool = True,
update_local_assets: bool = True
) -> None:
self.remove_library_path(path_old, update_local_assets=False)
self._asset_index.flush_is_local()
self.add_library_path(path_new,
primary=primary,
update_local_assets=update_local_assets)
def get_library_paths(self):
return self.library_paths
def get_library_path(self, primary: bool = True):
if self.library_paths and primary:
return self.library_paths[0]
elif len(self.library_paths) > 1:
# TODO(Mitchell): Return the most relevant lib path based on some input (?)
return None
else:
return None
def _get_user_info(self) -> Tuple:
req = self._api.get_user_info()
user_name = None
user_id = None
if req.ok:
data = req.body
user_name = data["user"]["name"]
user_id = data["user"]["id"]
self.login_error = None
else:
# TODO(SOFT-1029): Create an error log for fail in get user info
self.login_error = req.error
return user_name, user_id
def _get_credits(self):
if self.user is None:
msg = "_get_credits() called without user."
self._api.report_message(
"addon_get_credits", msg, "error")
return
req = self._api.get_user_balance()
if req.ok:
data = req.body
self.user.credits = data.get("subscription_balance")
self.user.credits_od = data.get("ondemand_balance")
else:
self.user.credits = None
self.user.credits_od = None
msg = f"ERROR: {req.error}"
self._api.report_message(
"addon_get_credits", msg, "error")
def _get_subscription_details(self):
"""Fetches the current user's subscription status."""
req = self._api.get_subscription_details()
if req.ok:
plan = req.body
self.user.plan.update_from_dict(plan)
@run_threaded(tm.PoolKeys.INTERACTIVE)
def update_plan_data(self, done_callback: Optional[Callable] = None) -> None:
# TODO(Joao): sub thread the two private functions
self._get_credits()
self._get_subscription_details()
if done_callback is not None:
done_callback()
def create_user(
self,
user_name: Optional[str] = None,
user_id: Optional[int] = None,
done_callback: Optional[Callable] = None) -> Optional[Future]:
if user_name is None or user_id is None:
user_name, user_id = self._get_user_info()
if user_name is None or user_id is None:
return None
self.user = PoliigonUser(
user_name=user_name,
user_id=user_id,
plan=PoliigonSubscription(
subscription_state=SubscriptionState.NOT_POPULATED)
)
future = self.update_plan_data(done_callback)
return future
def is_free_user(self) -> bool:
"""Identifies a free user which neither
has a plan nor on demand credits."""
if self.user is None:
# Should not happen in practice with a Poliigon addon
return False
sub_state = self.user.plan.subscription_state
free_plan = sub_state == SubscriptionState.FREE
no_credits = self.user.credits in [0, None]
no_od_credits = self.user.credits_od in [0, None]
return free_plan and no_credits and no_od_credits
def is_unlimited_user(self) -> bool:
if self.user is None:
return False
elif self.user.plan in [None, SubscriptionState.NOT_POPULATED]:
return False
elif self.user.plan.is_unlimited is None:
return False
return self.user.plan.is_unlimited
def is_paused_subscription(self) -> Optional[bool]:
"""Return True, if the Subscription is in paused state.
Return value may be None, if there is no plan.
"""
if self.user is None or self.user.plan is None:
return None
return self.user.plan.subscription_state == SubscriptionState.PAUSED
def get_user_credits(self, incl_od: bool = True) -> int:
"""Returns the number of _spendable_ credits."""
subscr_paused = self.is_paused_subscription()
credits = self.user.credits
credits_od = self.user.credits_od
if not incl_od and credits_od is not None:
credits_od = 0
if credits is None and credits_od is None:
return 0
elif credits_od is None:
return credits if not subscr_paused else 0
elif credits is None:
return credits_od
else:
if subscr_paused:
return credits_od
else:
return credits + credits_od
def get_thumbnail_path(self, asset_name, index):
"""Return the best fitting thumbnail preview for an asset.
The primary grid UI preview will be named asset_preview1.png,
all others will be named such as asset_preview1_1K.png
"""
if index == 0:
# 0 is the small grid preview version of _preview1.
# Fallback to legacy option of .jpg files if .png not found.
thumb = os.path.join(
self.online_previews_path,
asset_name + "_preview1.png"
)
if not os.path.exists(thumb):
thumb = os.path.join(
self.online_previews_path,
asset_name + "_preview1.jpg"
)
else:
thumb = os.path.join(
self.online_previews_path,
asset_name + f"_preview{index}_1K.png")
return thumb
def get_type_default_size(self, asset_data: AssetData) -> Optional[str]:
"""Returns a list of sizes valid for download."""
type_data = asset_data.get_type_data()
sizes_data = type_data.get_size_list()
size = None
if asset_data.asset_type == AssetType.TEXTURE:
size = self.settings_config.get("download", "tex_res")
elif asset_data.asset_type == AssetType.MODEL:
settings_size = self.settings_config.get(
"download", "model_res")
size_default = asset_data.model.size_default
has_default = size_default is not None
if settings_size in ["", "NONE", None] and has_default:
size = size_default
else:
size = settings_size
elif asset_data.asset_type == AssetType.HDRI:
size = self.settings_config.get("download", "hdri_light")
# TODO(Andreas): what about bg size?
elif asset_data.asset_type == AssetType.BRUSH:
size = self.settings_config.get("download", "brush")
valid_size = size in sizes_data
# If no valid size found, try to find at least one matching asset's
# available size data
if not valid_size:
for _size in reversed(SIZES):
if _size in sizes_data:
size = _size
break
return size
def set_first_local_asset(self, force_update: bool = False) -> None:
"""Conditionally assigns the current date to the settings file.
Meant to be used in conjunction with surveying, this should be called
either on first download or first import, if the value hasn't already
been set or if force_update is true."""
first_asset_timestamp = self.settings_config.get(
"user", "first_local_asset", fallback="")
if first_asset_timestamp == "" or force_update:
time_stamp = datetime.now()
self.settings_config.set(
"user", "first_local_asset", str(time_stamp))
self._settings.save_settings()
def set_first_preview_import(self, force_update: bool = False) -> None:
first_wm_timestamp = self.settings_config.get(
"user", "first_preview_import", fallback="")
if first_wm_timestamp == "" or force_update:
time_stamp = datetime.now()
self.settings_config.set(
"user", "first_preview_import", str(time_stamp))
self._settings.save_settings()
def set_first_purchase(self, force_update: bool = False) -> None:
first_purchase_timestamp = self.settings_config.get(
"user", "first_purchase", fallback="")
if first_purchase_timestamp == "" or force_update:
time_stamp = datetime.now()
self.settings_config.set(
"user", "first_purchase", str(time_stamp))
self._settings.save_settings()
def print_debug(self, *args, dbg=False, bg=True):
"""Print out a debug statement with no separator line.
Cache based on args up to a limit, to avoid excessive repeat prints.
All args must be flat values, such as already casted to strings, else
an error will be thrown.
"""
if dbg:
# Ensure all inputs are hashable, otherwise lru_cache fails.
stringified = [str(arg) for arg in args]
self._cached_print(*stringified, bg=bg)
@lru_cache(maxsize=32)
def _cached_print(self, *args, bg: bool):
"""A safe-to-cache function for printing."""
print(*args)
def open_asset_url(self, asset_id: int) -> None:
asset_data = self._asset_index.get_asset(asset_id)
url = self._api.add_utm_suffix(asset_data.url)
webbrowser.open(url)
def open_poliigon_link(self,
link_type: str,
add_utm_suffix: bool = True
) -> None:
"""Opens a Poliigon URL"""
# TODO(Andreas): As soon as P4B uses PoliigonAddon move code from
# api.open_poliigon_link here and remove function in api
self._api.open_poliigon_link(
link_type, add_utm_suffix, env_name=self._env.env_name)
def get_wm_download_path(self, asset_name: str) -> str:
"""Returns an asset name path inside the OnlinePreviews folder"""
path_poliigon = os.path.dirname(self._settings.base)
path_thumbs = os.path.join(path_poliigon, "OnlinePreviews")
path_wm_previews = os.path.join(path_thumbs, asset_name)
return path_wm_previews
def download_material_wm(
self, files_to_download: List[Tuple[str, str]]) -> None:
"""Synchronous function to download material preview."""
urls = []
files_dl = []
for _url_wm, _filename_wm_dl in files_to_download:
urls.append(_url_wm)
files_dl.append(_filename_wm_dl)
resp = self._api.pooled_preview_download(urls, files_dl)
if not resp.ok:
msg = f"Failed to download WM preview\n{resp}"
self._api.report_message(
"download_mat_preview_dl_failed", msg, "error")
# Continue, as some may have worked.
for _filename_wm_dl in files_dl:
filename_wm = _filename_wm_dl[:-3] # cut of _dl
try:
file_exists = os.path.exists(filename_wm)
dl_exists = os.path.exists(_filename_wm_dl)
if file_exists and dl_exists:
os.remove(filename_wm)
elif not file_exists and not dl_exists:
raise FileNotFoundError
if dl_exists:
os.rename(_filename_wm_dl, filename_wm)
except FileNotFoundError:
msg = f"Neither {filename_wm}, nor {_filename_wm_dl} exist"
self._api.report_message(
"download_mat_existing_file", msg, "error")
except FileExistsError:
msg = f"File {filename_wm} already exists, failed to rename"
self._api.report_message(
"download_mat_rename", msg, "error")
except Exception as e:
self.logger.exception("Unexpected exception while renaming WM preview")
msg = f"Unexpected exception while renaming {_filename_wm_dl}\n{e}"
self._api.report_message(
"download_wm_exception", msg, "error")
return resp
def get_config_param(self,
name_param: str,
name_group: str = "DEFAULT",
fallback: Optional[Any] = None
) -> Any:
"""Safely read a value from config (regardless of setup env or not)."""
if self._env.config is None:
return fallback
return self._env.config.get(name_group, name_param, fallback=fallback)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,870 @@
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
"""This module contains the API Remote Control."""
import concurrent
from concurrent.futures import CancelledError, Future, TimeoutError
from dataclasses import dataclass
from enum import IntEnum, unique
from functools import partial
import os
from queue import Queue
from threading import Event, Lock, Thread
import time
from typing import Callable, Dict, List, Optional, Any
from .addon import PoliigonAddon
from .api import (
ApiResponse,
TIMEOUT,
TIMEOUT_STREAM)
from .api_remote_control_params import (
AddonRemoteControlParams,
ApiJobParams,
ApiJobParamsDownloadAsset,
ApiJobParamsDownloadThumb,
ApiJobParamsDownloadWMPreview,
ApiJobParamsGetCategories,
ApiJobParamsGetUserData,
ApiJobParamsGetDownloadPrefs,
ApiJobParamsGetAvailablePlans,
ApiJobParamsGetUpgradePlan,
ApiJobParamsPutUpgradePlan,
ApiJobParamsResumePlan,
ApiJobParamsGetAssets,
ApiJobParamsLogin,
ApiJobParamsPurchaseAsset,
CmdLoginMode
)
from .assets import AssetData
@dataclass
class ApiResponseNewJob(ApiResponse):
# This class is deliberately empty.
# It only serves the purpose of being able to identify the ApiResponse
# returned from get_new_job_response() via instanceof().
pass
def get_new_job_response() -> ApiResponseNewJob:
resp = ApiResponseNewJob(
body={"data": []},
ok=False,
error="job waiting to execute"
)
return resp
@unique
class JobType(IntEnum):
LOGIN = 0
GET_USER_DATA = 1 # credits, subscription, user info
GET_CATEGORIES = 2
GET_DOWNLOAD_PREFS = 3
GET_AVAILABLE_PLANS = 4
GET_UPGRADE_PLAN = 5
PUT_UPGRADE_PLAN = 6
RESUME_PLAN = 7
GET_ASSETS = 10
DOWNLOAD_THUMB = 11
PURCHASE_ASSET = 12
DOWNLOAD_ASSET = 13
DOWNLOAD_WM_PREVIEW = 14,
UNIT_TEST = 15,
EXIT = 99999
class ApiJob():
"""Describes an ApiJob and gets passed through the queues,
subsequentyly being processed in thread_schedule and thread_collect.
"""
def __init__(
self,
job_type: JobType,
params: Optional[ApiJobParams] = None,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
result: ApiResponse = None,
future: Optional[Future] = None,
timeout: Optional[float] = None
):
self.job_type = job_type
self.params = params
self.callback_cancel = callback_cancel
self.callback_progress = callback_progress
self.callback_done = callback_done
if result is not None:
self.result = result
else:
self.result = get_new_job_response()
self.future = future
self.timeout = timeout
def __eq__(self, other):
return self.job_type == other.job_type and self.params == other.params
def __str__(self) -> str:
return f"{self.job_type.name}\n{str(self.params)}"
class ApiRemoteControl():
def __init__(self, addon: PoliigonAddon):
# Only members defined in addon_core.PoliigonAddon are allowed to be
# used inside this module
self._addon = addon
self._addon_params = AddonRemoteControlParams()
self._tm = addon._tm
self._api = addon._api
self._asset_index = addon._asset_index
is_dev = addon._env.env_name != "prod"
self.logger = self._addon.log_manager.initialize_logger(
"APIRC", have_filehandler=is_dev)
self.queue_jobs = Queue()
self._start_thread_schedule()
self.queue_jobs_done = Queue()
self._start_thread_collect()
self._start_thread_watchdog()
self.lock_jobs_in_flight = Lock()
self.jobs_in_flight = {} # {job_type: [futures]}
self.in_shutdown = False
self.init_stats()
def init_stats(self) -> None:
"""Initializes job statistics counters."""
self.cnt_added = {}
self.cnt_queued = {}
self.cnt_cancelled = {}
self.cnt_exec = {}
self.cnt_done = {}
self.cnt_restart_schedule = 0
self.cnt_restart_collect = 0
for job_type in JobType.__members__.values():
self.cnt_added[job_type] = 0
self.cnt_queued[job_type] = 0
self.cnt_cancelled[job_type] = 0
self.cnt_exec[job_type] = 0
self.cnt_done[job_type] = 0
def get_stats(self) -> Dict:
"""Returns job statistics counters as a dictionary."""
stats = {}
stats["Jobs added"] = self.cnt_added
stats["Jobs queued"] = self.cnt_queued
stats["Jobs cancelled"] = self.cnt_cancelled
stats["Jobs exec"] = self.cnt_exec
stats["Jobs done"] = self.cnt_done
stats["Restart schedule"] = self.cnt_restart_schedule
stats["Restart collect"] = self.cnt_restart_collect
return stats
def _start_thread_schedule(self) -> None:
self.schedule_running = False
thd_schedule_report_wrapped = self._tm.reporting_callable(
self._thread_schedule.__name__,
self._thread_schedule)
self.thd_schedule = Thread(target=thd_schedule_report_wrapped)
self.thd_schedule.name = "API RC Schedule"
self.thd_schedule.start()
def _start_thread_collect(self) -> None:
self.collect_running = False
thd_collect_report_wrapped = self._tm.reporting_callable(
self._thread_collect.__name__,
self._thread_collect)
self.thd_collect = Thread(target=thd_collect_report_wrapped)
self.thd_collect.name = "API RC Collect"
self.thd_collect.start()
def _start_thread_watchdog(self) -> None:
self.watchdog_running = False
self.event_watchdog = Event()
thd_watchdog_report_wrapped = self._tm.reporting_callable(
self._thread_watchdog.__name__,
self._thread_watchdog)
self.thd_watchdog = Thread(target=thd_watchdog_report_wrapped)
self.thd_watchdog.name = "API RC Watchdog"
self.thd_watchdog.start()
def add_job_login(self,
mode: CmdLoginMode = CmdLoginMode.LOGIN_BROWSER,
email: Optional[str] = None,
pwd: Optional[str] = None,
time_since_enable: Optional[int] = None,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to add a login or logout job."""
if mode == CmdLoginMode.LOGOUT:
self.empty_pipeline()
else: # login
self._asset_index.flush()
params = ApiJobParamsLogin(mode, email, pwd, time_since_enable)
self.add_job(
job_type=JobType.LOGIN,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_get_user_data(self,
user_name: str,
user_id: str,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to add a get user data job."""
params = ApiJobParamsGetUserData(user_name, user_id)
self.add_job(
job_type=JobType.GET_USER_DATA,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_get_download_prefs(self,
*,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to get user download preferences."""
params = ApiJobParamsGetDownloadPrefs()
self.add_job(
job_type=JobType.GET_DOWNLOAD_PREFS,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_get_available_plans(self,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
params = ApiJobParamsGetAvailablePlans()
self.add_job(
job_type=JobType.GET_AVAILABLE_PLANS,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_get_upgrade_plan(self,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
params = ApiJobParamsGetUpgradePlan()
self.add_job(
job_type=JobType.GET_UPGRADE_PLAN,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_put_upgrade_plan(self,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
params = ApiJobParamsPutUpgradePlan()
self.add_job(
job_type=JobType.PUT_UPGRADE_PLAN,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_resume_plan(self,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
params = ApiJobParamsResumePlan()
self.add_job(
job_type=JobType.RESUME_PLAN,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_get_categories(self,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to add a get categories job."""
params = ApiJobParamsGetCategories()
self.add_job(
job_type=JobType.GET_CATEGORIES,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_get_assets(self,
library_paths: List[str],
tab: str, # KEY_TAB_ONLINE, KEY_TAB_MY_ASSETS
category_list: List[str] = ["All Assets"],
search: str = "",
idx_page: int = 1,
page_size: int = 10,
force_request: bool = False,
do_get_all: bool = True,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True,
ignore_old_names: bool = True
) -> None:
"""Convenience function to add a get assets job."""
params = ApiJobParamsGetAssets(library_paths,
tab,
category_list,
search,
idx_page,
page_size,
force_request,
do_get_all,
ignore_old_names)
self.add_job(
job_type=JobType.GET_ASSETS,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def add_job_download_thumb(self,
asset_id: int,
url: str,
path: str,
idx_thumb: int = -1,
do_update: bool = False,
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = False
) -> None:
"""Convenience function to add a download thumb job."""
params = ApiJobParamsDownloadThumb(
asset_id, url, path, do_update, idx_thumb=idx_thumb)
temp_path = f"{path}_temp"
if not os.path.isfile(temp_path):
fwriter = open(temp_path, "wb")
fwriter.close()
elif os.path.isfile(path):
job = ApiJob(
job_type=JobType.DOWNLOAD_THUMB,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done)
callback_done(job=job)
return
else:
return
self.add_job(
job_type=JobType.DOWNLOAD_THUMB,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT_STREAM)
def add_job_purchase_asset(
self,
asset_data: AssetData,
category_list: List[str] = ["All Assets"],
search: str = "",
job_download: Optional[Callable] = None, # type: ApiJob
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to add a purchase asset job."""
params = ApiJobParamsPurchaseAsset(asset_data,
category_list,
search,
job_download)
self.add_job(
job_type=JobType.PURCHASE_ASSET,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT)
def create_job_download_asset(self,
asset_data: AssetData,
size: str = "2K",
size_bg: str = "",
type_bg: str = "EXR",
lod: str = "NONE",
variant: str = "",
download_lods: bool = False,
native_mesh: bool = True,
renderer: str = "",
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None
) -> ApiJob:
"""Convenience function to add a download asset job."""
params = ApiJobParamsDownloadAsset(
self._addon, asset_data, size, size_bg, type_bg, lod, variant,
download_lods, native_mesh, renderer)
job = ApiJob(
job_type=JobType.DOWNLOAD_ASSET,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
timeout=TIMEOUT_STREAM
)
# Due to the limitation of the number of threads, the download thread
# may not start immediately. In that case it would seem, as if nothing
# is happening.
asset_data.state.dl.start()
if callback_progress is not None:
callback_progress(job)
return job
def add_job_download_asset(self,
asset_data: AssetData,
size: str = "2K",
size_bg: str = "",
type_bg: str = "EXR",
lod: str = "NONE",
variant: str = "",
download_lods: bool = False,
native_mesh: bool = True,
renderer: str = "",
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to add a download asset job."""
self.cnt_added[JobType.DOWNLOAD_ASSET] += 1
job = self.create_job_download_asset(
asset_data,
size,
size_bg,
type_bg,
lod,
variant,
download_lods,
native_mesh,
renderer,
callback_cancel,
callback_progress,
callback_done
)
self.enqueue_job(job, force)
def add_job_download_wm_preview(
self,
asset_data: AssetData,
renderer: str = "",
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True
) -> None:
"""Convenience function to add a download WM preview job."""
params = ApiJobParamsDownloadWMPreview(asset_data,
renderer)
self.add_job(
job_type=JobType.DOWNLOAD_WM_PREVIEW,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
force=force,
timeout=TIMEOUT_STREAM)
def add_job_exit(self) -> None:
"""Convenience function to add an APIRC exit job."""
job = ApiJob(
job_type=JobType.EXIT,
params={},
callback_cancel=None,
callback_progress=None,
callback_done=None)
# Enqueue directly, as actual enqueue_jobs() gets disabled before
# shutdown
self.queue_jobs.put(job)
def _is_job_already_enqueued(self, job: ApiJob) -> bool:
"""Returns True, if an identical job exists already."""
with self.lock_jobs_in_flight:
jobs_in_flight_copy = self.jobs_in_flight.copy()
try:
return job in jobs_in_flight_copy[job.job_type]
except KeyError:
return False
def enqueue_job(self, job: ApiJob, force: bool = True) -> None:
"""Enqueúes a single ApiJob.
Arguments:
force: Default True, False: Enqueue only, if not queued already
"""
if not force and self._is_job_already_enqueued(job):
return
self.cnt_queued[job.job_type] += 1
self.queue_jobs.put(job)
def add_job(self,
job_type: JobType,
params: Any = {}, # Class from API RC Params
callback_cancel: Optional[Callable] = None,
callback_progress: Optional[Callable] = None,
callback_done: Optional[Callable] = None,
force: bool = True,
timeout: Optional[float] = None
) -> None:
"""Adds a job to be processed by API remote control."""
self.cnt_added[job_type] += 1
job = ApiJob(
job_type=job_type,
params=params,
callback_cancel=callback_cancel,
callback_progress=callback_progress,
callback_done=callback_done,
timeout=timeout)
self.enqueue_job(job, force)
def _release_job(self, job: ApiJob) -> None:
"""Removes a finished job from 'in flight' list."""
try:
with self.lock_jobs_in_flight:
self.jobs_in_flight[job.job_type].remove(job)
except (KeyError, ValueError):
pass # List of job type not found or job not found in list
def enqueue_job_shutdown(self, job: ApiJob, force: bool = True):
"""Used to replace enqueue_job() method during shutdown to avoid any
new jobs being enqueued.
Function is deliberately empty!
"""
pass
def empty_pipeline(self) -> None:
"""Gets rid of any jobs in API RC's pipeline.
In what way is this different to wait_for_all()?
wait_for_all() will get rid of all (or just a single type) jobs
currently in the pipeline. But it does make no attempt to avoid new
jobs being added. For example getting user data spawns five different
follow up jobs, some of those then spawning more follow up jobs (e.g.
the get asstes ones). empty_pipeline()
"""
# Have all jobs already in thread pool exit as early as possible
self.in_shutdown = True
# Prevent any new jobs from being queued
f_enqueue_job_bak = self.enqueue_job
self.enqueue_job = self.enqueue_job_shutdown
# Empty job queue to prevent new jobs from being scheduled in
# thread pool
while not self.queue_jobs.empty():
self.queue_jobs.get_nowait()
# Schedule point to allow threads to run home
# Sleep should be short, but longer than OS's tick
time.sleep(0.050) # 50 ms
self.wait_for_all(timeout=None)
# Re-enable normal operation
self.enqueue_job = f_enqueue_job_bak
self.in_shutdown = False
def shutdown(self) -> None:
"""Stops remote control's threads."""
# Tear watchdog down, first (we do not want it to restart anything)
self.watchdog_running = False
self.event_watchdog.set()
self.thd_watchdog.join()
# Have all jobs already in thread pool exit as early as possible
self.in_shutdown = True
# Prevent any new jobs from being queued
self.enqueue_job = self.enqueue_job_shutdown
# Empty job queue to prevent new jobs from being scheduled in
# thread pool
while not self.queue_jobs.empty():
self.queue_jobs.get_nowait()
# Enqueue the "exit job", which will lead to _thread_schedule and
# _thread_collect to exit
self.add_job_exit()
# Lastly wait for eveything to come to halt.
# timeout=None, use job type specific timeouts
self.wait_for_all(timeout=None)
def _wait_for_type(self,
jobs_in_flight_copy: Dict,
job_type: JobType,
do_wait: bool,
timeout: Optional[int]
) -> None:
"""Cancels all jobs of given type, optionally waits until cancelled."""
for job in jobs_in_flight_copy[job_type]:
try:
with self.lock_jobs_in_flight:
self.jobs_in_flight[job.job_type].remove(job)
except (KeyError, AttributeError):
pass
if job.result is None:
job.result = ApiResponse(ok=True,
body={"data": []},
error="job cancelled")
if job.future is None:
self.logger.warning(f"Future is None: {job.job_type.name}")
continue
elif job.future.cancel():
self.cnt_cancelled[job.job_type] += 1
continue
try:
job.callback_cancel()
except TypeError:
pass # Not every job has a cancel callback
if do_wait:
try:
if timeout is None:
timeout = job.timeout
job.future.result(timeout)
except (CancelledError,
TimeoutError,
concurrent.futures._base.TimeoutError) as e:
msg = ("API RC's job did not return upon cancel. "
f"Timeout: {timeout}\nJob: {str(job)}")
self.logger.exception(msg)
self._addon._api.report_exception(e)
def wait_for_all(self,
job_type: Optional[JobType] = None,
do_wait: bool = True,
timeout: Optional[int] = None
) -> None:
"""Cancels all jobs or just a given type, optionally waits until
cancelled.
Arguments:
job_type: Specify to cancel jobs of a certain type, None for all types
do_wait: Set to True to wait for cancellation, otherwise just cancel
and return immediately.
timeout: Time to wait for futures to finish. If None,
job type specific timeouts will be used (defined in
add_job_xyz() functions below).
"""
with self.lock_jobs_in_flight:
jobs_in_flight_copy = self.jobs_in_flight.copy()
if job_type is None:
for job_type in jobs_in_flight_copy:
self._wait_for_type(
jobs_in_flight_copy, job_type, do_wait, timeout)
elif job_type in jobs_in_flight_copy:
self._wait_for_type(
jobs_in_flight_copy, job_type, do_wait, timeout)
def is_job_type_active(self, job_type: JobType) -> bool:
"""Returns True if there's at least one job of given type in flight."""
return len(self.jobs_in_flight.get(job_type, [])) > 0
def _thread_schedule(self) -> None:
"""Thread waiting on job queue to start jobs in thread pool."""
self.schedule_running = True
while self.schedule_running:
job = self.queue_jobs.get()
self.cnt_exec[job.job_type] += 1
if job.job_type != JobType.EXIT:
with self.lock_jobs_in_flight:
try:
self.jobs_in_flight[job.job_type].append(job)
except KeyError:
self.jobs_in_flight[job.job_type] = [job]
if job.job_type != JobType.EXIT:
job.future = self._tm.queue_thread(
job.params.thread_execute,
job.params.POOL_KEY,
max_threads=None,
foreground=False,
api_rc=self,
job=job
)
else:
# JobType.EXIT
self.queue_jobs_done.put(job) # stop collector
self.schedule_running = False
def callback_enqueue_done(fut, job: ApiJob) -> None:
self.queue_jobs_done.put(job)
cb_done = partial(callback_enqueue_done, job=job)
try:
job.future.add_done_callback(cb_done)
except AttributeError as e:
# JobType.EXIT has no Future
if job.job_type != JobType.EXIT:
msg = f"Job {job.job_type.name} has no Future"
self.logger.exception(msg)
self._addon._api.report_exception(e)
def _thread_collect(self) -> None:
"""Thread awaiting threaded jobs to finish, then executes job's post
processing.
"""
self.collect_running = True
while self.collect_running:
job = self.queue_jobs_done.get()
if job.job_type != JobType.EXIT:
try:
job.params.finish(self, job)
except BaseException as e:
# Finish handlers are not allowed to tear down our
# collect thread...
msg = ("A job's finish function failed unexpectedly: "
f"{str(job)}")
self.logger.exception(msg)
self._addon._api.report_exception(e)
else:
# JobType.EXIT
self.collect_running = False
break
try:
job.callback_done(job=job)
except TypeError:
pass # There is no done callback
except BaseException as e:
# Done callbacksare not allowed to tear down our
# collect thread...
msg = ("A job's done callback failed unexpectedly: "
f"{str(job)}")
self.logger.exception(msg)
self._addon._api.report_exception(e)
self._release_job(job)
self.cnt_done[job.job_type] += 1
def _thread_watchdog(self) -> None:
self.watchdog_running = True
while self.watchdog_running:
# Event used as a sleep (will only get set during shutdown)
self.event_watchdog.wait(1.0)
if not self.watchdog_running:
break
if not self.thd_schedule.is_alive():
self.cnt_restart_schedule += 1
self._start_thread_schedule()
msg = f"API RC's schedule failed ({self.cnt_restart_schedule})"
self._addon._api.report_message(
"apirc_thread_failure_schedule", msg, "error")
self.logger.critical(msg)
if not self.thd_collect.is_alive():
self.cnt_restart_collect += 1
self._start_thread_collect()
msg = f"API RC's collect failed ({self.cnt_restart_collect})"
self._addon._api.report_message(
"apirc_thread_failure_collect", msg, "error")
self.logger.critical(msg)
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,817 @@
# ##### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
import os
from typing import Callable, Dict, List, Optional, Tuple
from dataclasses import dataclass
from concurrent.futures import (Future,
ThreadPoolExecutor)
from xml.etree import ElementTree
import errno
import time
from threading import Lock
from .api import (ApiResponse,
DownloadStatus,
DQStatus,
ERR_OS_NO_PERMISSION,
ERR_URL_EXPIRED,
ERR_OS_NO_SPACE,
ERR_LIMIT_DOWNLOAD_RATE)
from .assets import (AssetType,
AssetData,
PREVIEWS)
DOWNLOAD_TEMP_SUFFIX = "dl"
DOWNLOAD_POLL_INTERVAL = 0.25
MAX_DOWNLOAD_RETRIES = 10
MAX_PARALLEL_ASSET_DOWNLOADS = 2
MAX_PARALLEL_DOWNLOADS_PER_ASSET = 8
SIZE_DEFAULT_POOL = 10
MAX_RETRIES_PER_FILE = 3
MAX_RETRIES_PER_ASSET = 2
# This list defines a priority to fallback available formats
# NOTE: This is only for Convention 1 downloads
SUPPORTED_TEX_FORMATS = ["jpg", "png", "tiff", "exr"]
MODEL_FILE_EXT = ["fbx", "blend", "max", "c4d", "skp", "ma"]
class DownloadTimer():
start_time: float
end_time: float
duration: float
def start(self) -> None:
self.start_time = time.monotonic()
def end(self) -> float:
self.end_time = time.monotonic()
self.duration = self.end_time - self.start_time
return self.duration
@dataclass
class DynamicFile:
name: Optional[str]
contents: Optional[str]
@dataclass
class FileDownload:
asset_id: int
url: str
filename: str
convention: int
size_expected: int
size_downloaded: int = 0
resolution_size: str = None
status: DownloadStatus = DownloadStatus.INITIALIZED
directory: str = ""
fut: Optional[Future] = None
duration: float = -1.0 # avoid div by zero, but result stays clearly wrong
lock: Lock = Lock()
max_retries: int = MAX_RETRIES_PER_FILE
retries: int = 0
error: Optional[str] = None
cf_ray: Optional[str] = None
def do_retry(self) -> bool:
return self.retries < self.max_retries
def get_path(self, temp: bool = False) -> str:
directory = self.directory
return os.path.join(directory, self.get_filename(temp))
def get_filename(self, temp: bool = False) -> str:
if temp:
return self.filename + DOWNLOAD_TEMP_SUFFIX
else:
return self.filename
def set_status_cancelled(self) -> None:
# do not overwrite final states
with self.lock:
is_done = self.status == DownloadStatus.DONE
has_error = self.status == DownloadStatus.ERROR
if not is_done and not has_error:
self.status = DownloadStatus.CANCELLED
def set_status_ongoing(self) -> bool:
res = True
# do not overwrite user cancellation
with self.lock:
if self.status != DownloadStatus.CANCELLED:
self.status = DownloadStatus.ONGOING
else:
res = False
return res
def set_status_error(self) -> None:
with self.lock:
self.status = DownloadStatus.ERROR
def set_status_done(self) -> None:
with self.lock:
self.status = DownloadStatus.DONE
class AssetDownload:
addon = None # PoliigonAddon - No typing due to circular import
data_payload: Dict
tpe: ThreadPoolExecutor
asset_data: AssetData = None
uuid: Optional[str] = None
download_list: Optional[List[FileDownload]] = None
dynamic_files_list: Optional[List[DynamicFile]] = None
size_asset_bytes_expected: int = 0
size_asset_bytes_downloaded: int = 0
# Directory (named after the size) to be used for convention 1 assets
dir_size_target: str = None
# Status flags
max_retries: int = MAX_RETRIES_PER_ASSET
retries: int = 0
stop_files_retry: bool = False
all_done: bool = False
is_cancelled: bool = False
any_error: bool = False
res_error: Optional[str] = None
res_error_message: Optional[str] = None
error_dl_list: List[FileDownload] = list
dl_error: Optional[FileDownload] = None
def __init__(self,
addon, # PoliigonAddon - No typing due to circular import
asset_data: AssetData,
size: str,
dir_target: str,
lod: str = "NONE",
download_lods: bool = False,
native_mesh: bool = False,
renderer: Optional[str] = None,
update_callback: Optional[Callable] = None
) -> None:
self.addon = addon
self.asset_data = asset_data
self.size = size
self.lod = lod
self.download_lods = download_lods
self.native_mesh = native_mesh
self.renderer = renderer
self.update_callback = update_callback
self.dir_target = os.path.join(dir_target, asset_data.asset_name)
self.download_list = []
self.convention = self.asset_data.get_convention()
self.type_data = self.asset_data.get_type_data()
self.workflow = self.type_data.get_workflow()
self.identified_previews = 0
self.previews_reported = False
self.timer = DownloadTimer()
def kickoff_download(self) -> bool:
self.set_data_payload()
self.create_download_folder()
self.tpe = ThreadPoolExecutor(max_workers=MAX_PARALLEL_DOWNLOADS_PER_ASSET)
self.run_asset_download_retries(self.download_asset)
return self.all_done and not self.is_cancelled
def download_asset(self) -> None:
self.get_download_list()
if self.download_list in [None, []]:
self.any_error = True
if self.res_error is not None:
err = self.res_error
else:
err = "Empty download list"
msg = f"AssetId: {self.asset_data.asset_id} Error: {err}"
self.asset_data.state.dl.set_error(error_msg=err)
self.addon.logger_dl.error(msg)
# Only report to Sentry if the error is not Max Download rate
if err != ERR_LIMIT_DOWNLOAD_RATE:
self.addon._api.report_message("download_asset_empty_download_list",
msg,
"error")
return
self.download_loop()
def set_data_payload(self) -> None:
self.data_payload = {
"assets": [
{
"id": self.asset_data.asset_id,
"name": self.asset_data.asset_name
}
]
}
if self.convention == 0:
self.data_payload["assets"][0]["sizes"] = [self.size]
elif self.convention == 1:
self.data_payload["assets"][0]["resolution"] = self.size
if self.asset_data.asset_type in [AssetType.HDRI, AssetType.TEXTURE]:
self.set_texture_payload()
elif self.asset_data.asset_type == AssetType.MODEL:
self.set_model_payload()
def set_texture_payload(self) -> None:
if self.convention == 0:
map_codes = self.type_data.get_map_type_code_list(self.workflow)
self.data_payload["assets"][0]["workflows"] = [self.workflow]
self.data_payload["assets"][0]["type_codes"] = map_codes
elif self.convention == 1:
prefs_available = self.addon.user.map_preferences is not None
use_prefs = self.addon.user.use_preferences_on_download
if prefs_available and use_prefs:
map_descs, _ = self.type_data.get_maps_per_preferences(
self.addon.user.map_preferences,
filter_extensions=True)
else:
map_descs = self.type_data.map_descs[self.workflow]
map_list = []
for _map_desc in map_descs:
file_format = "UNKNOWN"
for _ff in SUPPORTED_TEX_FORMATS:
if _ff in _map_desc.file_formats:
file_format = _ff
break
if file_format == "UNKNOWN":
map_name = _map_desc.display_name
msg = (f"UNKNWOWN file format for download; "
f"Asset Id: {self.asset_data.asset_id} Map: {map_name}")
self.addon._api.report_message(
"download_invalid_format", msg, "error")
self.addon.logger_dl.error(msg)
map_dict = {
"type": _map_desc.map_type_code,
"format": file_format
}
map_list.append(map_dict)
self.data_payload["assets"][0]["maps"] = map_list
def set_model_payload(self) -> None:
self.data_payload["assets"][0]["lods"] = int(self.download_lods)
if self.native_mesh and self.renderer is not None:
self.data_payload["assets"][0]["softwares"] = [self.addon._api.software_dl_dcc]
self.data_payload["assets"][0]["renders"] = [self.renderer]
else:
self.data_payload["assets"][0]["softwares"] = ["ALL_OTHERS"]
def create_download_folder(self) -> bool:
try:
os.makedirs(self.dir_target, exist_ok=True)
except PermissionError:
self.asset_data.state.dl.set_error(error_msg=ERR_OS_NO_PERMISSION)
self.addon.logger_dl.exception(
f"{ERR_OS_NO_PERMISSION}: {self.dir_target}")
return False
except OSError as e:
self.asset_data.state.dl.set_error(error_msg=str(e))
self.addon.logger_dl.exception(f"Download directory: {self.dir_target}")
return False
self.addon.logger_dl.debug(f"Download directory: {self.dir_target}")
self.asset_data.state.dl.set_directory(self.dir_target)
# For convention 1 it should be saved in a size folder
if self.size is not None and self.convention == 1:
self.dir_size_target = os.path.join(self.dir_target, self.size)
if not os.path.isdir(self.dir_size_target):
os.makedirs(self.dir_size_target, exist_ok=True)
return True
def get_download_list(self) -> None:
self.timer.start()
# Getting dl list (FileDownload) and total bytes size
res = self.addon._api.download_asset_get_urls(self.data_payload)
if not res.ok:
self.res_error = res.error
custom_msg = res.body.get("errors", {}).get("message", [])
if len(custom_msg) > 0:
self.res_error_message = custom_msg[0]
return
self.build_download_list(res)
self.addon.logger_dl.info(
f"=== Requesting URLs took {self.timer.end():.3f} s.")
def define_download_folder(self, filename: str) -> str:
if self.convention == 0:
return self.dir_target
dl_folder = self.dir_size_target
base_filename, suffix = os.path.splitext(filename)
base_filename_low = base_filename.lower()
last_preview = None
for _preview in PREVIEWS:
if base_filename_low.endswith(_preview):
dl_folder = self.dir_target
self.identified_previews += 1
last_preview = filename
if self.identified_previews > 3 and not self.previews_reported:
msg = (f"Identifying multiple Preview images for "
f"Asset id: {self.asset_data.asset_id} (e.g {last_preview})")
self.addon._api.report_message(
"multiple_previews", msg, level="info")
self.previews_reported = True
return dl_folder
def build_download_list(self, res: ApiResponse) -> None:
files_list = res.body.get("files", [])
self.uuid = res.body.get("uuid", None)
if self.uuid in [None, ""]:
self.addon.logger_dl.error("No UUID for download")
model_exists = False
filename_model_fbx_source = None
url_model_fbx_source = None
size_expected_model_fbx_source = 0
for url_dict in files_list:
url = url_dict.get("url")
filename = url_dict.get("name")
size_expected = url_dict.get("bytes", 0)
resolution_size = url_dict.get("resolution", None)
self.size_asset_bytes_expected += size_expected
if not url or not filename:
raise RuntimeError(f"Missing url or filename {url}")
elif "_SOURCE" in filename:
if filename.lower().endswith(".fbx"):
filename_model_fbx_source = filename
url_model_fbx_source = url
size_expected_model_fbx_source = size_expected
continue
filename_ext = os.path.splitext(filename)[1].lower()
filename_ext = filename_ext[1:] # get rid of dot
if filename_ext.lower() in MODEL_FILE_EXT:
model_exists = True
dl = FileDownload(
asset_id=self.asset_data.asset_id,
url=url,
filename=filename,
convention=self.convention,
size_expected=size_expected,
resolution_size=resolution_size,
directory=self.define_download_folder(filename))
self.download_list.append(dl)
# Fallback if "xyz_SOURCE.fbx" is the only model file
if filename_model_fbx_source is not None and not model_exists:
dl = FileDownload(asset_id=self.asset_data.asset_id,
url=url_model_fbx_source,
filename=filename_model_fbx_source,
convention=self.convention,
size_expected=size_expected_model_fbx_source,
directory=self.dir_target)
self.download_list.append(dl)
msg = f"Model asset with just SOURCE LOD: {self.asset_data.asset_id}"
self.addon._api.report_message(
"model_with_only_source_lod", msg, level="info")
self.set_dynamic_files(res.body.get("dynamic_files", None))
def set_dynamic_files(self,
dynamic_files_api: Optional[List[Dict]]
) -> None:
"""Reads dynamic file information from server's API response."""
if dynamic_files_api is None:
return
self.dynamic_files_list = []
for _dynamic_file_dict in dynamic_files_api:
name = _dynamic_file_dict.get("name", None)
contents = _dynamic_file_dict.get("contents", None)
dynamic_file = DynamicFile(name=name, contents=contents)
self.dynamic_files_list.append(dynamic_file)
def download_loop(self) -> None:
"""The actual download loop in download_asset_sync()."""
self.all_done = False
self.addon.logger_dl.debug("Download Loop")
self.asset_data.state.dl.set_progress(0.001)
if self.asset_data.state.dl.is_cancelled():
self.is_cancelled = True
self.cancel_downloads()
return
self.schedule_downloads()
self.download_asset_loop_poll()
if self.all_done:
# Consider download failed upon dynamic file error.
#
# ATM we will not expose any issues with dynamic file data from server
# and let the entire download succeed, anyway.
self.all_done = self.store_dynamic_files(expose_api_error=False)
self.rename_downloads()
return
def track_quality(self) -> None:
if self.uuid in [None, ""]:
return
if self.all_done:
# User may still have cancelled download (judging by state in
# asset data), but we suceeded anyway
self.addon._api.track_download_quality(uuid=self.uuid,
status=DQStatus.SUCCESS)
elif self.is_cancelled and not self.any_error:
self.addon._api.track_download_quality(uuid=self.uuid,
status=DQStatus.CANCELED,
error="User cancelled download")
else:
file_dl_error = self.dl_error
if file_dl_error is None:
return
msg = (f"Error: {file_dl_error.error}, "
f"File: {file_dl_error.url}, CF-ray: {file_dl_error.cf_ray}")
self.addon._api.track_download_quality(uuid=self.uuid,
status=DQStatus.FAILED,
error=msg)
def schedule_downloads(self,
download_list: Optional[List[FileDownload]] = None
) -> None:
"""Submits downloads to thread pool."""
if download_list is None:
download_list = self.download_list
self.addon.logger_dl.debug("Scheduling Downloads")
download_list.sort(key=lambda dl: dl.size_expected)
for download in download_list:
# Andreas: Could also check here, if already DONE and not start
# the thread at all.
# Yet, I decided to prefer it handled by the thread itself.
# In this way the flow is always identical.
download.status = DownloadStatus.WAITING
download.retries += 1
download.fut = self.tpe.submit(self.addon._api.download_asset_file,
download=download)
self.addon.logger_dl.debug(f"Submitted {download.filename}. "
f"Retry: {download.retries}")
self.addon.logger_dl.debug("Download Asset Schedule Done")
def check_download_progress(self) -> None:
self.addon.logger_dl.debug(self.download_list)
self.any_error = False
self.error_dl_list = []
self.is_cancelled = self.asset_data.state.dl.is_cancelled()
self.all_done = True
self.size_asset_bytes_downloaded = 0
for download in self.download_list:
self.size_asset_bytes_downloaded += download.size_downloaded
fut = download.fut
if not fut.done():
self.all_done = False
continue
res = fut.result()
exc = fut.exception()
res_error = res.error
had_excp = exc is not None
if not res.ok or had_excp:
if had_excp:
self.addon.logger_dl.error(exc)
self.any_error = True
self.all_done = False
download.error = res_error
self.asset_data.state.dl.set_error(error_msg=res_error)
self.error_dl_list.append(download)
if self.any_error:
self.process_file_retries()
elif self.all_done:
self.addon.logger_dl.debug("All Done :)")
self.set_progress()
def download_asset_loop_poll(self) -> None:
"""Used in download_asset_sync to poll results inside download loop."""
self.addon.logger_dl.debug("Starting Download Poll Loop")
while not self.all_done and not self.stop_files_retry and not self.is_cancelled:
time.sleep(DOWNLOAD_POLL_INTERVAL)
self.check_download_progress()
def process_file_retries(self) -> None:
"""Manages the retries per file."""
for dl_error in self.error_dl_list:
if not dl_error.do_retry():
self.dl_error = dl_error
self.stop_files_retry = True
break
if not self.is_cancelled and not self.stop_files_retry:
self.schedule_downloads(self.error_dl_list)
return
self.cancel_downloads()
def cancel_downloads(self) -> None:
"""Cancels all download threads"""
self.addon.logger_dl.debug("Start cancel")
for download in self.download_list:
download.set_status_cancelled()
if download.fut is not None:
download.fut.cancel()
# Wait for threads to actually return
self.addon.logger_dl.debug("Waiting")
for download in self.download_list:
if download.fut is None:
continue
if download.fut.cancelled():
continue
try:
download.fut.result(timeout=60)
except TimeoutError:
# TODO(Andreas): Now there seems to be some real issue...
msg = (f"Asset id {self.asset_data.asset_id} download file "
"future Timeout error with no result.")
self.addon._api.report_message("download_file_with_no_result",
msg,
"error")
raise
except BaseException:
msg = (f"Asset id {self.asset_data.asset_id} download file "
"exception with no result.")
self.addon._api.report_message("download_file_with_no_result",
msg,
"error")
self.addon.logger_dl.exception(f"Unexpected error: {msg}")
raise
self.addon.logger_dl.debug("Done")
def set_progress(self) -> None:
progress = self.size_asset_bytes_downloaded / max(self.size_asset_bytes_expected, 1)
self.asset_data.state.dl.set_progress(max(progress, 0.001))
self.asset_data.state.dl.set_downloaded_bytes(self.size_asset_bytes_expected)
try: # Init progress bar
self.update_callback()
except TypeError:
pass # No update callback
def do_retry(self) -> bool:
first_attempt = self.retries == 0
if first_attempt:
return True
do_retry = self.retries < self.max_retries
expired_error = False
if self.dl_error is not None:
expired_error = self.dl_error.error == ERR_URL_EXPIRED
# Asset level download only retries in case of Expired URL
return expired_error and do_retry
def run_asset_download_retries(self, method: callable) -> None:
while self.do_retry():
self.retries += 1
method()
self.track_quality()
self.cancel_downloads()
def rename_downloads(self) -> Tuple[bool, str]:
"""Renames dowhloaded temp file."""
self.addon.logger_dl.debug("Start rename")
error_msg = ""
all_successful = True
for download in self.download_list:
if download.status != DownloadStatus.DONE:
self.addon.logger_dl.warning(("File status not done despite "
"all files reported done!"))
path_temp = download.get_path(temp=True)
temp_exists = os.path.exists(path_temp)
path_final = download.get_path(temp=False)
final_exists = os.path.exists(path_final)
if not temp_exists and final_exists:
continue
try:
os.rename(path_temp, path_final)
except FileExistsError:
os.remove(path_temp)
except FileNotFoundError:
download.status = DownloadStatus.ERROR
download.error = f"Missing file: {path_temp}"
self.addon.logger_dl.error(
("Neither temp download file nor target do exist\n"
f" {path_temp}\n"
f" {path_final}"))
all_successful = False
except PermissionError:
# Note from Andreas:
# I am not entirely sure, how this can happen (after all we
# just downloaded the file...).
# My assumption is, that somehow the download thread (while
# already being done) did not actually exit, yet, maybe due to
# some scheduling mishaps and is still keeping a handle to the
# file. If I am correct, maybe a "sleep(0.1 sec)" and another
# attempt to rename could get us out of this.
# But that's of course pretty ugly and we should discuss
# first, if we want to try something like this or just let
# the download fail.
download.status = DownloadStatus.ERROR
download.error = ("Lacking permission to rename downloaded"
f" file: {path_temp}")
self.addon.logger_dl.error(
(f"No permission to rename download:\n from: {path_temp}"
f"\n to: {path_final}"))
all_successful = False
# Gets the first error found to give feedback for the user
if error_msg is not None and download.error not in [None, ""]:
error_msg = download.error
self.addon.logger_dl.debug(f"Done, succeess = {all_successful}")
return all_successful, error_msg
def _check_xml_data(self,
xml_s: str,
expose_api_error: bool = False
) -> bool:
"""Checks an XML string for correct XML structure."""
asset_data = self.asset_data
asset_id = asset_data.asset_id
xml_ok = False
try:
# While we are not really interested in actual contents atm,
# we parse the XML nevertheless to make sure it is "parseable".
xml_root = ElementTree.XML(xml_s)
if xml_root is not None:
xml_ok = True
except ElementTree.ParseError as e:
if expose_api_error:
asset_data.state.dl.set_error(error_msg="Dynamic file error")
msg = (f"Could not save dynamic file for {asset_id}, "
f"XML parsing issue\n{e}")
self.addon.logger_dl.exception(msg)
self.addon._api.report_message("download_df_xml_issue", msg, "error")
if not xml_ok:
return False # NOK reported above in exception
return True
def _check_dynamic_file_data(self,
dynamic_file: DynamicFile,
expose_api_error: bool = False
) -> bool:
"""Checks the dynamic file data (currently expecting XML) received
from API.
"""
asset_data = self.asset_data
asset_id = asset_data.asset_id
if dynamic_file.name is None:
if expose_api_error:
asset_data.state.dl.set_error(error_msg="Dynamic file error")
msg = (f"Could not save dynamic file for {asset_id}, "
"no name provided")
self.addon.logger_dl.error(msg)
self.addon._api.report_message("download_df_no_filename", msg, "error")
return False
contents = dynamic_file.contents
if contents is None:
if expose_api_error:
asset_data.state.dl.set_error(error_msg="Dynamic file error")
msg = (f"Could not save dynamic file for {asset_id}, "
"no contents provided")
self.addon.logger_dl.error(msg)
self.addon._api.report_message("download_df_no_contents", msg, "error")
return False
return self._check_xml_data(contents, expose_api_error)
def _store_single_dynamic_file(self,
dynamic_file: DynamicFile,
expose_api_error: bool = False
) -> bool:
"""Stores a dynamic file (currently only XML data) to disk."""
asset_data = self.asset_data
asset_id = asset_data.asset_id
result = self._check_dynamic_file_data(dynamic_file, expose_api_error)
if not result:
# Here download fails only, if exposure of errors in dynamic_file
# data from server is desired.
return not expose_api_error
# Since we need to store into the correct "size" subfolder and
# also since MaterialX does have little sense, if there were no other
# files downloaded, we'll use the path of the first FileDownload
if len(self.download_list) > 0:
file_download = self.download_list[0]
path_file = file_download.get_path()
dir_asset = os.path.dirname(path_file)
else:
# Without any file downloads, all we have here is asset's path:
dir_asset = self.asset_data.state.dl.get_directory()
path_dynamic_file = os.path.join(dir_asset, dynamic_file.name)
try:
with open(path_dynamic_file, "w") as write_file:
write_file.write(dynamic_file.contents)
except OSError as e:
if e.errno == errno.ENOSPC:
asset_data.state.dl.set_error(error_msg=ERR_OS_NO_SPACE)
msg = f"Asset {asset_id}: No space for dynamic file."
# TODO(Andreas): No logger in PoliigonConnector, yet
self.addon.logger_dl.exception(msg)
self.addon._api.report_message("download_df_no_space", msg, "error")
elif e.errno == errno.EACCES:
asset_data.state.dl.set_error(
error_msg=ERR_OS_NO_PERMISSION)
msg = f"Asset {asset_id}: No permission to write dynamic file."
# TODO(Andreas): No logger in PoliigonConnector, yet
self.addon.logger_dl.exception(msg)
self.addon._api.report_message("download_df_permission", msg, "error")
else:
asset_data.state.dl.set_error(error_msg=str(e))
msg = (f"Asset {asset_id}: Unexpected error "
"upon writing dynamic file.")
# TODO(Andreas): No logger in PoliigonConnector, yet
self.addon.logger_dl.logger.exception(msg)
msg += f"\n{e}"
self.addon._api.report_message("download_df_os_error", msg, "error")
# Note: Even if dynamic file data issue above does not get exposed
# to user, any failure to write the correct MaterialX data
# will still lead to a failed download.
return False
return True
def store_dynamic_files(self,
expose_api_error: bool = False
) -> bool:
"""Stores all dynamic files belonging to an asset download to disk."""
if self.dynamic_files_list is None:
return True
if len(self.dynamic_files_list) == 0:
return True
# Note: We'll get here only after all asset files got downloaded
# successfully. Thus we can store any dynamic files errors in
# AssetData's download status (no need to be afraid of
# overwriting any other error) to present in UI.
for _dynamic_file in self.dynamic_files_list:
result = self._store_single_dynamic_file(_dynamic_file,
expose_api_error)
if not result:
return False
return True
@@ -0,0 +1,116 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
import os
from typing import Optional
try:
import ConfigParser
except Exception:
import configparser as ConfigParser
class PoliigonEnvironment():
"""Poliigon environment used for assisting in program control flow."""
addon_name: str # e.g. poliigon-addon-blender
base: str # Path to base directory of addon or package
env_filename: str
config: ConfigParser.ConfigParser = None
# Required env fields
api_url: str = ""
api_url_v2: str = ""
env_name: str = ""
required_attrs = ["api_url", "api_url_v2", "env_name"]
# Optional env fields
host: str = ""
forced_sampling: bool = False
local_updater_json: Optional[str] = None
def __init__(self,
addon_name: str,
base: str = os.path.dirname(os.path.abspath(__file__)),
env_filename: str = "env.ini"):
self.addon_name = addon_name
self.base = base
self.env_filename = env_filename
self._update_files(base)
self._load_env(base, env_filename)
def _load_env(self, path, filename):
env_file = os.path.join(path, filename)
if os.path.exists(env_file):
try:
# Read .ini here and set values
# https://docs.python.org/3/library/configparser.html#configparser.ConfigParser.optionxform
config = ConfigParser.ConfigParser()
config.optionxform = str
config.read(env_file)
# Required fields
self.config = config
self.api_url = config.get("DEFAULT", "api_url", fallback="")
self.api_url_v2 = config.get("DEFAULT", "api_url_v2", fallback="")
self.env_name = config.get("DEFAULT", "env_name", fallback="")
for k, v in vars(self).items():
if k in self.required_attrs and v in [None, ""]:
raise ValueError(
f"Attribute '{k}' missing from env file")
# Optional fields that should always be present
self.host = config.get("DEFAULT", "host", fallback="")
self.forced_sampling = config.getboolean(
"DEFAULT", "forced_sampling", fallback=False)
self.local_updater_json = config.get(
"DEFAULT", "local_updater_json", fallback=None)
except ValueError as e:
msg = f"Could not load environment file for {self.addon_name}"
raise RuntimeError(msg) from e
else:
# Assume production environment and set fallback values
self.api_url = "https://api.poliigon.com/api/v1"
self.api_url_v2 = "https://apiv2.poliigon.com/api/v2"
self.env_name = "prod"
self.host = ""
self.forced_sampling = False
def _update_files(self, path):
"""Updates files in the specified path within the addon."""
update_key = "_update"
search_key = "env" + update_key
files_to_update = [f for f in os.listdir(path)
if os.path.isfile(os.path.join(path, f))
and os.path.splitext(f)[0].endswith(search_key)]
for f in files_to_update:
f_split = os.path.splitext(f)
tgt_file = f_split[0][:-len(update_key)] + f_split[1]
try:
os.replace(os.path.join(path, f), os.path.join(path, tgt_file))
print(f"Updated {tgt_file}")
except PermissionError as e:
print(f"Encountered 'file_permission_error': {e}")
except OSError as e:
print(f"Encountered 'os_error': {e}")
@@ -0,0 +1,238 @@
import logging
import os
from io import StringIO
from typing import Optional, List
from .env import PoliigonEnvironment
# Numerical comment -> values for .ini
NOT_SET = logging.NOTSET # 0
DEBUG = logging.DEBUG # 10
INFO = logging.INFO # 20
WARNING = logging.WARNING # 30
ERROR = logging.ERROR # 40
CRITICAL = logging.CRITICAL # 50
# Import and use get_addon_logger() to get hold of the logger manager
addon_logger = None
class MockLogger:
"""Placeholder logger which accepts any arguments and does nothing"""
def __getattr__(self, method_name):
# Names matching the built in logging class methods
log_methods = ['critical', 'error', 'exception', 'fatal',
'debug', 'info', 'log',
'warn', 'warning']
if method_name not in log_methods:
raise RuntimeError("Invalid logger method")
def method(*args, **kwargs):
return None
return method
class AddonLogger:
"""Class to store all the data (created loggers, formatting information
and handlers) and functionality related to the Addon Logs.
loggers: Stores all created loggers;
dcc_handlers: Stores all handlers created on DCC side.
For adding new Handlers, use the method set_dcc_handlers;
write_to_file: Defines if the new loggers write the output into a .log file;
log_file_path: The output path to create the .log file;
"""
loggers = []
dcc_handlers = []
write_to_file: bool = False
str_format = ("%(name)s, %(levelname)s, %(threadName)s, %(asctime)s, "
"%(filename)s/%(funcName)s:%(lineno)d: %(message)s")
date_format = "%I:%M:%S"
def __init__(
self,
env: Optional[PoliigonEnvironment] = None,
file_path: Optional[str] = None):
self.addon_env = env
self.file_handler = None
self.stream_handler = None
addon_core_path = os.path.split(os.path.split(os.path.realpath(__file__))[0])[0]
self.log_file_path = os.path.join(addon_core_path, "logs.log")
if file_path is not None:
self.log_file_path = file_path
self.write_to_file = True
else:
try:
self.write_to_file = self.addon_env.config.getboolean(
"DEFAULT", "log_to_file", fallback=False)
except AttributeError:
pass # no .ini
def _init_filehandler(self, have_filehandler: bool) -> None:
"""Optionally initializes the default file handler."""
if not have_filehandler:
self.file_handler = None
return
log_file_exists = os.path.isfile(self.log_file_path)
if log_file_exists:
log_file_write_access = os.access(self.log_file_path, os.W_OK)
else:
log_file_write_access = True
log_dir = os.path.dirname(self.log_file_path)
log_dir_exists = os.path.isdir(log_dir)
if log_dir_exists:
log_dir_write_access = os.access(log_dir, os.W_OK)
else:
log_dir_write_access = False
log_file_overwrite = log_file_exists and log_file_write_access
log_file_create_allowed = log_dir_exists and log_dir_write_access and not log_file_exists
if log_file_overwrite or log_file_create_allowed:
if self.file_handler is None:
self.file_handler = AddonFileHandler(self.log_file_path, self)
def initialize_logger(self, module_name: Optional[str] = None,
*,
log_lvl: Optional[int] = None,
log_stream: Optional[StringIO] = None,
base_name: str = "Addon",
append_dcc_handlers: bool = True,
have_filehandler: bool = True
) -> logging.Logger:
"""Set format, log level and returns a logger instance
Args:
module_name: The name of the module a required argument
Env log_lvl variable name is derived as follows:
Logger name: Addon => log_lvl
Logger name: Addon.DL => log_lvl_dl
Logger name: Addon.P4C.UI => log_lvl_p4c_ui
But also:
Logger name: bonnie => log_lvl
Logger name: clyde.whatever => log_lvl_whatever
log_lvl: Integer specifying which logs to be printed, one of:
https://docs.python.org/3/library/logging.html#levels
log_stream: Output to StringIO stream instead of the console if not None
base_name: By default all loggers get derived from logger "Addon".
append_dcc_handlers: Defines if the handlers cached in self.dcc_handlers
should be added for the new logger.
have_filehandler: Set to fault too disable logging to a file.
Returns:
Returns a reference to the initialized logger instance
Raises:
AttributeError: If log_lvl and env are both None.
"""
if module_name is None:
logger_name = f"{base_name}"
name_hierarchy = []
else:
logger_name = f"{base_name}.{module_name}"
name_hierarchy = module_name.split(".")
if log_lvl is None:
log_lvl_name = "log_lvl"
for name in name_hierarchy:
log_lvl_name += f"_{name.lower()}"
try:
log_lvl = self.addon_env.config.getint(
"DEFAULT", log_lvl_name, fallback=NOT_SET)
except AttributeError:
log_lvl = NOT_SET # no .ini
logger = logging.getLogger(logger_name)
logger.propagate = False
logger.setLevel(log_lvl)
if self.stream_handler is None:
self.stream_handler = logging.StreamHandler(log_stream)
self.set_dcc_handlers(handlers=[self.stream_handler])
self._init_filehandler(have_filehandler)
if self.file_handler is not None:
self.set_dcc_handlers(handlers=[self.file_handler])
if append_dcc_handlers and len(self.dcc_handlers) > 1:
self.set_dcc_handlers(loggers=[logger])
self.loggers.append(logger)
return logger
def set_write_to_file_handler(
self, enabled: bool, path: Optional[str] = None) -> None:
self.write_to_file = enabled
if path is not None:
self.log_file_path = path
def set_dcc_handlers(
self,
loggers: Optional[List[logging.Logger]] = None,
handlers: Optional[List[logging.Handler]] = None) -> None:
loggers_to_add = loggers if loggers is not None else self.loggers
handlers_to_add = handlers if handlers is not None else self.dcc_handlers
for _logger in loggers_to_add:
for _handler in handlers_to_add:
self.set_handlers_formatter([_handler])
_logger.addHandler(_handler)
if _handler not in self.dcc_handlers:
self.dcc_handlers.append(_handler)
def set_handlers_formatter(
self,
handlers: Optional[List[logging.Handler]] = None) -> None:
handlers_to_format = handlers if handlers is not None else self.dcc_handlers
formatter = logging.Formatter(
fmt=self.str_format, datefmt=self.date_format)
for _handler in handlers_to_format:
_handler.setFormatter(formatter)
class AddonFileHandler(logging.FileHandler):
def __init__(self, filepath: str, log_manager: AddonLogger):
super(AddonFileHandler, self).__init__(filepath)
self.log_manager = log_manager
self.original_emit_event = self.emit
self.emit = self.custom_emit
def change_log_filename(self, filename: str):
if os.path.isfile(filename):
self.baseFilename = filename
def custom_emit(self, record: logging.LogRecord) -> None:
if not self.log_manager.write_to_file:
return
if self.log_manager.log_file_path != self.baseFilename:
self.baseFilename = self.log_manager.log_file_path
self.close()
self.stream = None
# Ensure original emit event runs
self.original_emit_event(record)
def get_addon_logger(env: Optional[PoliigonEnvironment] = None) -> AddonLogger:
global addon_logger
if addon_logger is None:
addon_logger = AddonLogger(env)
return addon_logger
@@ -0,0 +1,219 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
from enum import IntEnum
from dataclasses import dataclass
from typing import Dict, Optional
from .multilingual import _t
# MAPS_TYPE_NAMES defined at the end of the file (needs MapType defined)
@dataclass()
class MapDescription:
description: str
display_name: str
class MapType(IntEnum):
"""Supported texture map types.
NOTE: When extending, existing values MUST NEVER be changed.
NOTE 2: Derived from IntEnum for easier "to JSON serialization"
"""
# Convention 0 values
DEFAULT = 1
UNKNOWN = 1
ALPHA = 2 # Usually associated with a brush
ALPHAMASKED = 3
AO = 4
BUMP = 5
BUMP16 = 6
COL = 7
DIFF = 8
DISP = 9
DISP16 = 10
EMISSIVE = 11
EMISSION = 11
ENV = 12 # Environment for an HDRI, typically a .jpg file
JPG = 12 # Environment for an HDRI, type_code as in ApiResponse
FUZZ = 13
GLOSS = 14
IDMAP = 15
LIGHT = 16 # Lighting for an HDRI, typically a .exr file
HDR = 16 # Lighting for an HDRI, type_code as in ApiResponse
MASK = 17 # Mask here means opacity
METALNESS = 18
NRM = 19
NRM16 = 20
OVERLAY = 21
REFL = 22
ROUGHNESS = 23
SSS = 24
TRANSLUCENCY = 25
TRANSMISSION = 26
OPACITY = 27
UNDEF = 28
# Non convention 0 types (needed for convention conversion)
NA_ORM = 50
NA_VERTEXBLEND = 51
# Convention 1, values 100 to 149 (150 and up for convention 1 only maps)
# NOTE: Value - 100 should match convention 0
AmbientOcclusion = 104
BaseColor = 107
BaseColorOpacity = 103 # realtime only
BaseColorVertexBlend = 151
Displacement = 109
Emission = 111
Environment = 112 # Map for converting HDRI into Env
HDRI = 116
ORM = 150 # packmap where R:AO, G:Roughness, B:Metalness, realtime only
Metallic = 118
Normal = 119
Opacity = 117
Roughness = 123
ScatteringColor = 124
SheenColor = 113
Translucency = 125
Transmission = 126
@classmethod
def from_type_code(cls, map_type_code: str):
if map_type_code in MAPS_TYPE_NAMES:
return cls[map_type_code]
map_type_code = map_type_code.split("_")[0]
if map_type_code in MAPS_TYPE_NAMES:
return cls[map_type_code]
return cls.UNKNOWN
def get_convention(self) -> int:
return self.value // 100
def get_effective(self): # -> MapType
return self.convert_convention(0)
def convert_convention(self, target_convention: int): # -> MapType
convention_in = self.value // 100
convention_diff = target_convention - convention_in
return MapType(self.value + (100 * convention_diff))
def get_description(self) -> Optional[MapDescription]:
return MAP_DESCRIPTIONS.get(self.get_effective(), None)
MAP_DESCRIPTIONS: Dict = {
MapType.ALPHAMASKED: MapDescription(
description=_t(
"This texture map is identical to the Base Color Map, but with "
"an added Alpha channel containing the opacity map. This is "
"included in materials containing empty see-through space such "
"as sheer fabrics and leaves."),
display_name=_t("Base Color Opacity")),
MapType.AO: MapDescription(
description=_t(
"Defines the shadows in the crevices of the material. It's combined "
"with the color map by using a Multiply layer blend operation."),
display_name=_t("Ambient Occlusion")),
MapType.COL: MapDescription(
description=_t(
"Contains the pure color information of the surface, "
"devoid of any shadow or reflection."),
display_name=_t("Base Color")),
MapType.DISP: MapDescription(
description=_t(
"This black and white image defines the height information of the "
"surface. Light values are raised, dark values are reduced, "
"mid-grey (0.5) represents the flat mid-point of the surface."),
display_name=_t("Displacement")),
MapType.FUZZ: MapDescription(
description=_t(
"Defines the fine fuzz of microfibers in cloth-like surfaces. "
"Included with many fabrics textures. "
"The sheen color defines only the color."),
display_name=_t("Sheen Color")),
MapType.METALNESS: MapDescription(
description=_t(
"This black and white image defines which parts are metal "
"(white) and which are non-metal (black)."),
display_name=_t("Metallic")),
MapType.NRM: MapDescription(
description=_t(
"This purple-ish image defines the height information, which is faked "
"by shader (not physically altering the mesh)"),
display_name=_t("Normal")),
MapType.ROUGHNESS: MapDescription(
description=_t(
"This black and white image defines how sharp or diffuse the "
"reflections are. Blacker values are glossy, "
"whiter values are matte."),
display_name=_t("Roughness")),
MapType.SSS: MapDescription(
description=_t(
"Defines the color of light passing through solid closed manifold "
"objects like food or fabric. This is included in fabric "
"and vegetation textures."),
display_name=_t("Scattering Color")),
MapType.TRANSLUCENCY: MapDescription(
description=_t(
"Defines the color of light penetrating and appearing on the "
"backside of a flat thinshell meshes. "
"This is included in fabric and vegetation textures."),
display_name=_t("Translucency")),
MapType.TRANSMISSION: MapDescription(
description=_t(
"Defines which parts of the texture are refracting light, "
"and is included in textures like glass or liquids. "
"The IOR (Index of Refraction) should be set be defined "
"by you depending on the material."),
display_name=_t("Transmission")),
MapType.MASK: MapDescription(
description=_t(
"Defines which parts of the texture are opaque, or transparent "
"(completely invisible, without refraction). "
"This is included in materials containing empty see-through "
"space such as sheer fabrics and leaves."),
display_name=_t("Opacity")),
MapType.NA_ORM: MapDescription(
description=_t(
"This special texture stores the same Ambient Occlusion, Roughness "
"and Metalness information, but each are stored in the separate Red, "
"Green and Blue channels respectively. This special map is "
"typically only used in realtime rendering and game applications."),
display_name=_t("ORM"))
}
MAPS_TYPE_NAMES = MapType.__members__
@@ -0,0 +1,82 @@
from typing import List, Callable
import gettext
def _m(message: str) -> str:
"""Placeholder to mark strings to be translated and stored on .pot file"""
return message
def _o(message: str) -> str:
"""Gets the original string, e.g after set a translated variable"""
return gettext.dgettext("en-US", message)
def _t(message: str) -> str:
"""Gets the translated string with the current domain setup"""
try:
return gettext.gettext(message) # noqa
except NameError:
# If none domain is initialized, _() will not be defined yet;
return message
class MsgFallback(gettext.NullTranslations):
"""Fallback to report if one is trying to translate a message that is not
registered on the stored domains (platforms and languages);
"""
def __init__(self, fallback_method: Callable = None) -> None:
super().__init__()
self.fallback = fallback_method
def gettext(self, msg) -> str:
self.fallback(msg)
return msg
class Multilingual:
"""Class to store and manage all the domains for multilingual translation."""
report_message: Callable # Report function to be set on the addon;
curr_language: str
supported_languages = ["en-US", "test_dummy"]
# All domains already registered.
# NOTE: Do not change this from outside this module;
_domains: List[gettext.GNUTranslations]
def __init__(self):
self._domains = []
self.report_message = None
self.curr_language = None
def install_domain(self,
language: str,
dir_lang: str,
domain: str = "addon-core") -> None:
if language not in self.supported_languages:
return
current_domain = gettext.translation(domain,
localedir=dir_lang,
languages=[language, "en-US"])
current_domain.install()
# If there are already installed domains, they are used as fallback if
# a given message is not found. Each new domain will fall back to the
# previous one - until the first added domain, which will call
# report_message_missing;
if len(self._domains) > 0:
current_domain.add_fallback(self._domains[-1])
else:
current_domain.add_fallback(MsgFallback(self.report_message_missing))
self.curr_language = language
self._domains.append(current_domain)
def report_message_missing(self, msg: str) -> None:
if self.report_message is not None and self.curr_language != "en-US":
error_msg = f"{self.curr_language}:\"{msg}\""
self.report_message("missing_locale_msg", error_msg, "error")
@@ -0,0 +1,678 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
"""Module fo asynchronous user notificatiions."""
from dataclasses import dataclass, field
from functools import wraps
from enum import IntEnum
from queue import Queue
from threading import Lock
from typing import Callable, Dict, List, Optional, Any
from .thread_manager import PoolKeys
from .multilingual import _m
# Predefined priority values (lower numbers -> higher prio)
NOTICE_PRIO_LOWEST = 200
NOTICE_PRIO_LOW = 100
NOTICE_PRIO_MEDIUM = 50
NOTICE_PRIO_HIGH = 20
NOTICE_PRIO_URGENT = 1
NOTICE_PRIO_MAT_TEMPLATE = NOTICE_PRIO_HIGH
NOTICE_PRIO_NO_INET = NOTICE_PRIO_LOW # Show, but other errors have precedent
NOTICE_PRIO_PROXY = NOTICE_PRIO_MEDIUM
NOTICE_PRIO_SETTINGS_WRITE = NOTICE_PRIO_HIGH
NOTICE_PRIO_SURVEY = NOTICE_PRIO_LOWEST
NOTICE_PRIO_UPDATE = NOTICE_PRIO_HIGH + 5 # urgent, but room for "more urgent"
NOTICE_PRIO_RESTART = NOTICE_PRIO_LOW
# Predefined notice IDs
NOTICE_ID_MAT_TEMPLATE = "MATERIAL_TEMPLATE_ERROR"
NOTICE_ID_NO_INET = "NO_INTERNET_CONNECTION"
NOTICE_ID_PROXY = "PROXY_CONNECTION_ERROR"
NOTICE_ID_SETTINGS_WRITE = "SETTINGS_WRITE_ERROR"
NOTICE_ID_SURVEY_FREE = "NPS_INAPP_FREE"
NOTICE_ID_SURVEY_ACTIVE = "NPS_INAPP_ACTIVE"
NOTICE_ID_UPDATE = "UPDATE_READY_MANUAL_INSTALL"
NOTICE_ID_VERSION_ALERT = "ADDON_VERSION_ALERT"
NOTICE_ID_RESTART_ALERT = "NOTICE_ID_RESTART_ALERT"
# Predefined notice titles
# Used a default param in create functions, but should usually be overridden
# by passing in localized titles from DCC.
NOTICE_TITLE_MAT_TEMPLATE = _m("Material template error")
NOTICE_TITLE_NO_INET = _m("No internet access")
NOTICE_TITLE_PROXY = _m("Encountered proxy error")
NOTICE_TITLE_SETTINGS_WRITE = _m("Failed to write settings")
NOTICE_TITLE_SURVEY = _m("How's the addon?")
NOTICE_TITLE_UPDATE = _m("Update ready")
NOTICE_TITLE_DEPRECATED = _m("Deprecated version")
NOTICE_TITLE_RESTART = _m("Restart needed")
# Predefined notice Labels (text to be displayed on the notification Banner)
# Used a default param in create functions, but should usually be overridden
# by passing in localized titles from DCC.
NOTICE_LABEL_NO_INET = _m("Connection Lost")
NOTICE_LABEL_PROXY_ERROR = _m("Proxy Error")
NOTICE_LABEL_RESTART = _m("Restart needed")
# Predefined notice Body (text to be displayed on the notification Popup)
# Used a default param in create functions, but should usually be overridden
# by passing in localized titles from DCC.
NOTICE_BODY_NO_INET = _m("Cannot reach Poliigon, double check your "
"firewall is configured to access Poliigon servers: "
"*poliigon.com / *poliigon.net / *imagedelivery.net. "
"If this persists, please reach out to support.")
NOTICE_BODY_RESTART = _m("Please restart your 3D software")
# Predefined icons, assign DCC specific key/reference via init_icons()
NOTICE_ICON_WARN = "ICON_WARN"
NOTICE_ICON_INFO = "ICON_INFO"
NOTICE_ICON_SURVEY = "ICON_SURVEY"
NOTICE_ICON_NO_CONNECTION = "ICON_NO_CONNECTION"
class ActionType(IntEnum):
# Note: Numerical values are still same as in P4B, but entries got
# sorted alphabetically
OPEN_URL = 1
POPUP_MESSAGE = 3
RUN_OPERATOR = 4
UPDATE_READY = 2
class SignalType(IntEnum):
"""Types of each interaction with the notifications."""
VIEWED = 0
DISMISSED = 1
CLICKED = 2
@dataclass
class Notification():
"""Container object for a user notification.
NOTE: Do not instance Notification directly, but instead either use
NotificationSystem.create_... functions or instance derived
NotificationXYZ classes.
"""
# Unique id for this specific kind of notice, if possible re-use above
# NOTICE_ID_xyz.
id_notice: str
# Main title, should be short
title: str
# Indicator of how to structure and draw notification.
action: ActionType = field(init=False) # does NOT get auto initialized
# Priority is always > 0, lower values = higher priority
priority: int
# Label to be shown in the notification banner - to be defined per addon
label: str = ""
# Allow the user to dismiss the notification.
allow_dismiss: bool = True
# Dismiss after user interacted with the notification
auto_dismiss: bool = False
# Hover-over tooltip, if there is a button
tooltip: str = ""
# In Blender icon's are referenced in strings (icon enum),
# but this may differ per DCC. For prebuilt notices init_icons() has
# to be used to store DCC dependent icons, once.
icon: Optional[any] = None
# Defines if notification has to open a popup with more information and
# options for the user to then address or dismiss the notice
open_popup: bool = False
# Text for the button to execute notify callable when it opens a popup
action_string: Optional[str] = None
# action callable to be attached to the notification, to be executed
# when notification is clicked
action_callable: Optional[Callable] = None
# function to be called when the notification is dismissed (viewed or not)
on_dismiss_callable: Optional[Callable] = None
viewed: bool = False # False until actually drawn
clicked: bool = False # False until user interact with the notice
@dataclass
class AddonNotificationsParameters:
"""Parameters to be parsed from the addon.
parameters:
update_callable: Callable to be set as action_callable to update notifications
update_action_text: Action text for updates - used as popup update button text
update_body: Text with a description for update - used as popup text
NOTE: Feel free to add here any other parameter needed from the addon.
"""
update_callable: Optional[Callable] = None
update_action_text: str = NOTICE_TITLE_UPDATE
update_body: str = ""
@dataclass
class NotificationOpenUrl(Notification):
url: str = ""
def __post_init__(self):
self.action = ActionType.OPEN_URL
def get_key(self) -> str:
return "".join([self.action.name, self.url, self.label])
@dataclass
class NotificationPopup(Notification):
body: str = ""
url: str = ""
alert: bool = True
def __post_init__(self):
self.action = ActionType.POPUP_MESSAGE
def get_key(self) -> str:
return "".join([self.action.name, self.url, self.body])
@dataclass
class NotificationRunOperator(Notification):
# For Blender ops_name will be string, for C4D not so sure, yet...
# I guess, we could even store a callable in here.
ops_name: Optional[any] = None
def __post_init__(self):
self.action = ActionType.RUN_OPERATOR
def get_key(self) -> str:
return "".join([self.action.name, self.ops_name])
@dataclass
class NotificationUpdateReady(Notification):
download_url: str = ""
download_label: str = ""
logs_url: str = ""
logs_label: str = ""
body: str = ""
def __post_init__(self):
self.action = ActionType.UPDATE_READY
def get_key(self) -> str:
return "".join([self.action.name, self.download_url, self.download_label])
class NotificationSystem():
"""Abstraction to handle asynchronous user notification.
Each DCC has to populate icon_dcc_map.
"""
_api = None # PoliigonConector
_tm = None # Thread Manager
_queue_notice: Queue = Queue()
_lock_notice: Lock = Lock()
_notices: Dict = {} # {key: Notification}
# Each DCC use init_icons() to populate these values as is fitting for
# themselves.
icon_dcc_map: Dict[str, Optional[any]] = {
NOTICE_ICON_WARN: None,
NOTICE_ICON_INFO: None,
NOTICE_ICON_SURVEY: None,
NOTICE_ICON_NO_CONNECTION: None
}
addon_params = AddonNotificationsParameters()
def __init__(self, addon):
if addon is None:
return
self._api = addon._api
self._tm = addon._tm
def init_icons(
self,
icon_warn: Optional[Any] = None,
icon_info: Optional[Any] = None,
icon_survey: Optional[Any] = None,
icon_no_connection: Optional[Any] = None
) -> None:
self.icon_dcc_map[NOTICE_ICON_WARN] = icon_warn
self.icon_dcc_map[NOTICE_ICON_INFO] = icon_info
self.icon_dcc_map[NOTICE_ICON_SURVEY] = icon_survey
self.icon_dcc_map[NOTICE_ICON_NO_CONNECTION] = icon_no_connection
def _run_threaded(key_pool: PoolKeys,
max_threads: Optional[int] = None,
foreground: bool = False) -> Callable:
"""Schedules a function to run in a thread of a chosen pool."""
def wrapped_func(func: Callable) -> Callable:
@wraps(func)
def wrapped_func_call(self, *args, **kwargs):
args = (self, ) + args
return self._tm.queue_thread(func,
key_pool,
max_threads,
foreground,
*args,
**kwargs)
return wrapped_func_call
return wrapped_func
def _consume_queued_notices(self) -> None:
"""Empties the notice queue and stores all new notices in _notices.
Note: If an identical notice already exists, it will get skipped.
"""
with self._lock_notice:
while self._queue_notice.qsize() > 0:
notice = self._queue_notice.get(block=False)
key = notice.get_key()
if key in self._notices:
continue
self._notices[key] = notice
def _get_sorted_notices(self) -> List[Notification]:
"""Returns a priority sorted list with all notices."""
with self._lock_notice:
all_notices = list(self._notices.values())
all_notices.sort(key=lambda notice: notice.priority)
return all_notices
@_run_threaded(PoolKeys.INTERACTIVE)
def _thread_signal(
self, notice: Notification, signal_type: SignalType) -> None:
"""Asynchronously signals "notice got viewed" to server"."""
if signal_type == SignalType.VIEWED:
self._api.signal_view_notification(notice.id_notice)
elif signal_type == SignalType.DISMISSED:
self._api.signal_dismiss_notification(notice.id_notice)
elif signal_type == SignalType.CLICKED:
self._api.signal_click_notification(notice.id_notice, notice.action)
def _signal_view(self, notice: Notification) -> None:
"""Internally used to start the signal view thread."""
if self._api is None or not self._api._is_opted_in():
return
self._thread_signal(notice, SignalType.VIEWED)
def _signal_clicked(self, notice: Notification) -> None:
"""Internally used to start the signal click thread."""
if self._api is None or not self._api._is_opted_in():
return
self._thread_signal(notice, SignalType.CLICKED)
def _signal_dismiss(self, notice: Notification) -> None:
"""Internally used to start the signal dismiss thread."""
if self._api is None or not self._api._is_opted_in():
return
self._thread_signal(notice, SignalType.DISMISSED)
def enqueue_notice(self, notice: Notification) -> None:
"""Enqueues a new notification."""
self._queue_notice.put(notice)
def dismiss_notice(
self, notice: Notification, force: bool = False) -> None:
"""Dismisses a notice.
Use force parameter to dismiss 'un-dismissable' notices, e.g.
a 'no internet' notice, when internet is back on.
"""
if not notice.allow_dismiss and not force:
return
if not notice.clicked:
self._signal_dismiss(notice)
if notice.on_dismiss_callable is not None:
notice.on_dismiss_callable()
key = notice.get_key()
with self._lock_notice:
if key in self._notices:
del self._notices[key]
def clicked_notice(self, notice: Notification) -> None:
"""To be called, when a user interacted with the notice."""
notice.clicked = True
self._signal_clicked(notice)
if notice.action_callable is not None:
notice.action_callable()
if not notice.auto_dismiss:
return
self.dismiss_notice(notice)
def get_all_notices(self) -> List[Notification]:
"""Returns a priority sorted list with all notices.
Usually called from draw code.
"""
self._consume_queued_notices()
return self._get_sorted_notices()
def get_top_notice(
self, do_signal_view: bool = False) -> Optional[Notification]:
"""Returns current highest priority notice.
Usually called from draw code.
"""
notices_by_prio = self.get_all_notices()
try:
notice = notices_by_prio[0]
if do_signal_view and not notice.viewed:
self._signal_view(notice)
notice.viewed = True
except (KeyError, IndexError):
notice = None
return notice
def notification_popup(
self, notice: Notification, do_signal_view: bool = False) -> None:
"""Called when a popup notification is drawn"""
if do_signal_view and not notice.viewed:
self._signal_view(notice)
notice.viewed = True
def flush_all(self) -> None:
"""Flushes all existing notices."""
while not self._queue_notice.empty():
self._queue_notice.get(block=False)
with self._lock_notice:
self._notices = {}
def create_restart_needed(self,
title: str = NOTICE_TITLE_RESTART,
*,
label: str = NOTICE_LABEL_RESTART,
tooltip: str = "",
body: str = NOTICE_BODY_RESTART,
action_string: Optional[str] = None,
auto_enqueue: bool = True
) -> Notification:
"""Returns a pre-built 'Restart Needed' notice."""
notice = NotificationPopup(
id_notice=NOTICE_ID_RESTART_ALERT,
title=title,
label=label,
priority=NOTICE_PRIO_RESTART,
allow_dismiss=False,
open_popup=True,
action_string=action_string,
tooltip=tooltip,
icon=self.icon_dcc_map[NOTICE_ICON_WARN],
body=body
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_no_internet(self,
title: str = NOTICE_TITLE_NO_INET,
*,
label: str = NOTICE_LABEL_NO_INET,
tooltip: str = "",
body: str = NOTICE_BODY_NO_INET,
auto_enqueue: bool = True
) -> Notification:
"""Returns a pre-built 'No internet' notice."""
notice = NotificationPopup(
id_notice=NOTICE_ID_NO_INET,
title=title,
label=label,
priority=NOTICE_PRIO_NO_INET,
allow_dismiss=False,
open_popup=True,
action_string=None,
tooltip=tooltip,
icon=self.icon_dcc_map[NOTICE_ICON_NO_CONNECTION],
body=body
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_proxy(self,
title: str = NOTICE_TITLE_PROXY,
*,
label: str = NOTICE_LABEL_PROXY_ERROR,
tooltip: str = "",
body: str = NOTICE_BODY_NO_INET,
auto_enqueue: bool = True
) -> Notification:
"""Returns a pre-built 'Proxy error' notice."""
notice = NotificationPopup(
id_notice=NOTICE_ID_PROXY,
title=title,
label=label,
priority=NOTICE_PRIO_PROXY,
allow_dismiss=False,
open_popup=True,
action_string=None,
tooltip=tooltip,
icon=self.icon_dcc_map[NOTICE_ICON_WARN],
body=body
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_survey(self,
title: str = NOTICE_TITLE_SURVEY,
*,
is_free_user: bool,
tooltip: str,
free_survey_url: str,
active_survey_url: str,
label: str,
auto_enqueue: bool = True,
on_dismiss_callable: Optional[Callable] = None
) -> Notification:
"""Returns a pre-built 'user survey' notice."""
if is_free_user:
id_notice = NOTICE_ID_SURVEY_FREE
url = free_survey_url
else:
id_notice = NOTICE_ID_SURVEY_ACTIVE
url = active_survey_url
notice = NotificationOpenUrl(
id_notice=id_notice,
title=title,
priority=NOTICE_PRIO_SURVEY,
allow_dismiss=True,
auto_dismiss=True,
tooltip=tooltip,
url=url,
label=label,
icon=self.icon_dcc_map[NOTICE_ICON_SURVEY],
on_dismiss_callable=on_dismiss_callable
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_write_mat_template(self,
title: str = NOTICE_TITLE_MAT_TEMPLATE,
*,
tooltip: str,
body: str,
auto_enqueue: bool = True
) -> Notification:
"""Returns a pre-built 'Material template error' notice."""
notice = NotificationPopup(
id_notice=NOTICE_ID_MAT_TEMPLATE,
title=title,
priority=NOTICE_PRIO_MAT_TEMPLATE,
allow_dismiss=True,
tooltip=tooltip,
icon=self.icon_dcc_map[NOTICE_ICON_WARN],
body=body,
alert=True
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_version_alert(self,
title: str = NOTICE_TITLE_DEPRECATED,
*,
priority: int,
label: str,
tooltip: str,
open_popup: bool,
allow_dismiss: bool = True,
auto_dismiss: bool = True,
body: Optional[str] = None,
action_string: Optional[str] = None,
url: Optional[str] = None,
auto_enqueue: bool = True
) -> Notification:
"""Returns a pre-built 'Version Alert' notice.
Note: An Alert Notification can be a NotificationPopup or a
NotificationOpenUrl, depending on the given AlertData
"""
if open_popup:
notice = NotificationPopup(
id_notice=NOTICE_ID_VERSION_ALERT,
title=title,
priority=priority,
allow_dismiss=allow_dismiss,
auto_dismiss=auto_dismiss,
tooltip=tooltip,
label=label,
icon=self.icon_dcc_map[NOTICE_ICON_WARN],
open_popup=open_popup,
body=body,
action_string=action_string
)
else:
notice = NotificationOpenUrl(
id_notice=NOTICE_ID_VERSION_ALERT,
url=url,
title=title,
priority=priority,
allow_dismiss=allow_dismiss,
auto_dismiss=auto_dismiss,
tooltip=tooltip,
label=label,
icon=self.icon_dcc_map[NOTICE_ICON_WARN]
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_write_settings_error(self,
title: str = NOTICE_TITLE_SETTINGS_WRITE,
*,
tooltip: str,
body: str,
auto_enqueue: bool = True
) -> Notification:
"""Returns a pre-built 'write settings error' notice."""
notice = NotificationPopup(
id_notice=NOTICE_ID_SETTINGS_WRITE,
title=title,
priority=NOTICE_PRIO_SETTINGS_WRITE,
allow_dismiss=True,
tooltip=tooltip,
icon=self.icon_dcc_map[NOTICE_ICON_WARN],
body=body,
alert=True
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
def create_update(self,
title: str = NOTICE_TITLE_UPDATE,
*,
tooltip: str,
label: str,
download_url: str,
download_label: str = "",
logs_url: str = "",
logs_label: str = "",
auto_enqueue: bool = True,
open_popup: bool = True,
auto_dismiss: bool = True,
action_string: Optional[str] = None,
body: Optional[str] = None,
action_callable: Optional[Callable] = None
) -> Notification:
"""Returns a pre-built 'Update available' notice."""
if action_string is None:
action_string = self.addon_params.update_action_text
if body is None:
body = self.addon_params.update_body
if action_callable is None:
action_callable = self.addon_params.update_callable
notice = NotificationUpdateReady(
id_notice=NOTICE_ID_UPDATE,
title=title,
priority=NOTICE_PRIO_UPDATE,
allow_dismiss=True,
auto_dismiss=auto_dismiss,
tooltip=tooltip,
download_url=download_url,
download_label=download_label,
label=label,
logs_url=logs_url,
logs_label=logs_label,
icon=self.icon_dcc_map[NOTICE_ICON_INFO],
open_popup=open_popup,
action_string=action_string,
body=body,
action_callable=action_callable
)
if auto_enqueue:
self.enqueue_notice(notice)
return notice
@@ -0,0 +1,555 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
import html
from typing import Any, Dict, List, Optional
from .notifications import Notification, ActionType
from .api import STR_NO_PLAN
SIGNAL_PLAN_SUBSCRIBE_EDU = "PLAN_SUBSCRIBE_EDU"
SIGNAL_PLAN_SUBSCRIBE = "PLAN_SUBSCRIBE"
SIGNAL_PLAN_UPGRADE_NO_DLS = "PLAN_UPGRADE_NO_DLS"
SIGNAL_PLAN_PROMPT_UNLIMITED = "PLAN_PROMPT_UNLIMITED"
SIGNAL_PLAN_RESUME_PAUSED = "PLAN_RESUME_PAUSED"
SIGNAL_PLAN_RESUME_CANCELLATION = "PLAN_RESUME_SCHEDULED_CANCEL"
SIGNAL_PLAN_RESUME_SCHEDULED_PAUSE = "PLAN_RESUME_SCHEDULED_PAUSE"
class SubscriptionState(Enum):
"""Values for allowed user subscription states."""
NOT_POPULATED = 0
FREE = 1,
ACTIVE = 2,
PAUSED = 3,
PAUSE_SCHEDULED = 4,
CANCELLED = 4
class PlanUpgradeStatus(Enum):
NOT_POPULATED = 0
STUDENT_DISCOUNT = 1
TEACHER_DISCOUNT = 2
BECOME_PRO = 3
UPGRADE_PLAN_BALANCE = 4
RESUME_PLAN = 5
REMOVE_SCHEDULED_PAUSE = 6
REMOVE_CANCELLATION = 7
UPGRADE_PLAN_UNLIMITED = 8
NO_UPGRADE_AVAILABLE = 9
def get_signal_string(self) -> Optional[str]:
if self == self.NOT_POPULATED:
return None
elif self in [self.STUDENT_DISCOUNT, self.TEACHER_DISCOUNT]:
return SIGNAL_PLAN_SUBSCRIBE_EDU
elif self == self.BECOME_PRO:
return SIGNAL_PLAN_SUBSCRIBE
elif self == self.UPGRADE_PLAN_BALANCE:
return SIGNAL_PLAN_UPGRADE_NO_DLS
elif self == self.RESUME_PLAN:
return SIGNAL_PLAN_RESUME_PAUSED
elif self == self.REMOVE_SCHEDULED_PAUSE:
return SIGNAL_PLAN_RESUME_SCHEDULED_PAUSE
elif self == self.REMOVE_CANCELLATION:
return SIGNAL_PLAN_RESUME_CANCELLATION
elif self == self.UPGRADE_PLAN_UNLIMITED:
return SIGNAL_PLAN_PROMPT_UNLIMITED
else:
return None
def _decode_currency_symbol(currency_str: str) -> str:
decoded_str = ""
chars = currency_str.split(";")
for _char in chars:
# Processing chrs in html format (e.g "82;&#36" => R$)
try:
int_char = int(_char)
_char = chr(int_char)
except ValueError:
_char = html.unescape(_char)
if len(_char) != 1:
_char = ""
decoded_str += _char
return decoded_str
@dataclass
class PoliigonPlanUpgradeInfo:
ok: bool
error: Optional[str] = None
action: Optional[str] = None
amount_due: Optional[str] = None
amount_due_renewal: Optional[str] = None
renewal_date: Optional[str] = None
tax_rate: Optional[int] = None
currency_code: Optional[str] = None
currency_symbol: Optional[str] = None
previous_assets: Optional[int] = None
new_assets: Optional[int] = None
previous_users: Optional[int] = None
new_users: Optional[int] = None
@classmethod
def from_dict(cls, dictionary: Dict):
new = cls(ok=True)
new.action = dictionary.get("action")
new.amount_due = dictionary.get("amount_due")
new.amount_due_renewal = dictionary.get("amount_due_renewal")
new.renewal_date = dictionary.get("renewal_date")
new.tax_rate = dictionary.get("tax_rate")
new.currency_code = dictionary.get("currency_code")
new.currency_symbol = _decode_currency_symbol(
dictionary.get("currency_symbol"))
new.previous_assets = dictionary.get("previous_assets")
if isinstance(new.previous_assets, str):
new.previous_assets = new.previous_assets.title()
new.new_assets = dictionary.get("new_assets")
if isinstance(new.new_assets, str):
new.new_assets = new.new_assets.title()
new.previous_users = dictionary.get("previous_users")
new.new_users = dictionary.get("new_users")
return new
@dataclass
class PoliigonSubscription:
"""Container object for a subscription."""
plan_name: Optional[str] = None
plan_credit: Optional[int] = None
next_credit_renewal_date: Optional[datetime] = None
current_term_end: Optional[datetime] = None
next_subscription_renewal_date: Optional[datetime] = None
plan_paused_at: Optional[datetime] = None
plan_paused_until: Optional[datetime] = None
subscription_state: Optional[SubscriptionState] = SubscriptionState.NOT_POPULATED
period_unit: Optional[str] = None # e.g. per "month" or "year" for renewing
plan_price_id: Optional[str] = None
plan_price: Optional[str] = None # e.g. "123"
currency_code: Optional[str] = None # e.g. "USD"
base_price: Optional[float] = None # e.g. 123.45
currency_symbol: Optional[str] = None # e.g. "$" (special character)
is_unlimited: Optional[bool] = None
has_team: Optional[bool] = None
@staticmethod
def _to_float(value: Optional[str]) -> Optional[float]:
if value is None:
return None
# Replacing commas from the string, so they can be formatted as floats.
# For some users, this value can be formatted as "1,000"
if isinstance(value, str) and "," in value:
value.replace(",", "")
try:
return float(value)
except (ValueError, TypeError):
return None
def update_from_upgrade_dict(self, plan_dictionary: Dict) -> Any: # Returns an instance of the class
"""Creates a class instance from an API V2 available plans
response plan dictionary."""
if plan_dictionary.get("name") and plan_dictionary["name"] != STR_NO_PLAN:
self.plan_name = plan_dictionary["name"]
self.plan_credit = plan_dictionary.get("meta", {}).get("credits")
self.period_unit = plan_dictionary.get("periodUnit", None)
self.plan_price_id = plan_dictionary.get("id", None)
self.plan_price = plan_dictionary.get("price", None)
self.currency_code = plan_dictionary.get("currencyCode", None)
self.base_price = self._to_float(plan_dictionary.get("basePrice", None))
self.currency_symbol = _decode_currency_symbol(
plan_dictionary.get("currency_symbol", ""))
self.is_unlimited = bool(plan_dictionary.get("meta", {}).get("unlimited"))
self.has_team = bool(plan_dictionary.get("meta", {}).get("hasTeams"))
else:
self.plan_name = None
self.plan_credit = None
self.next_subscription_renewal_date = None
self.next_credit_renewal_date = None
self.subscription_state = SubscriptionState.FREE
self.period_unit = None
self.plan_price_id = None
self.plan_price = None
self.currency_code = None
self.base_price = None
self.currency_symbol = None
def update_from_dict(self, plan_dictionary: Dict) -> Any: # Returns an instance of the class
"""TCreates a class instance from an API V1 Subscription Data (Some API
V2 still returning with same structure as API V1 - e.g. put_upgrade_plan).
"""
if plan_dictionary.get("plan_name") and plan_dictionary["plan_name"] != STR_NO_PLAN:
self.plan_name = plan_dictionary["plan_name"]
self.plan_credit = plan_dictionary.get("plan_credit", None)
# Extract "2022-08-19" from "2022-08-19 23:58:37"
renew = plan_dictionary.get("next_subscription_renewal_date", None)
try:
renew = datetime.strptime(renew, "%Y-%m-%d %H:%M:%S")
self.next_subscription_renewal_date = renew
except (ValueError, TypeError):
self.next_subscription_renewal_date = None
end_plan = plan_dictionary.get("current_term_end", None)
try:
end_plan = datetime.strptime(end_plan, "%Y-%m-%d %H:%M:%S")
self.current_term_end = end_plan
except (ValueError, TypeError):
self.current_term_end = None
next_credits = plan_dictionary.get("next_credit_renewal_date", None)
try:
next_credits = datetime.strptime(
next_credits, "%Y-%m-%d %H:%M:%S")
self.next_credit_renewal_date = next_credits
except (ValueError, TypeError):
self.next_credit_renewal_date = None
paused_plan_info = plan_dictionary.get("paused_info", None)
not_renewing = self.next_subscription_renewal_date is None
if paused_plan_info is not None:
self.subscription_state = SubscriptionState.PAUSED
paused_date = paused_plan_info.get("pause_date", None)
resume_date = paused_plan_info.get("resume_date", None)
try:
self.plan_paused_at = datetime.strptime(
paused_date, "%Y-%m-%d %H:%M:%S")
self.plan_paused_until = datetime.strptime(
resume_date, "%Y-%m-%d %H:%M:%S")
now = datetime.now()
if now < self.plan_paused_at or now > self.plan_paused_until:
self.subscription_state = SubscriptionState.PAUSE_SCHEDULED
except (ValueError, TypeError):
self.plan_paused_until = None
self.plan_paused_at = None
elif not_renewing:
self.subscription_state = SubscriptionState.CANCELLED
else:
self.plan_paused_until = None
self.plan_paused_at = None
self.subscription_state = SubscriptionState.ACTIVE
self.period_unit = plan_dictionary.get("period_unit", None)
self.plan_price_id = plan_dictionary.get("plan_price_id", None)
self.plan_price = plan_dictionary.get("plan_price", None)
self.currency_code = plan_dictionary.get("currency_code", None)
self.base_price = self._to_float(plan_dictionary.get("base_price", None))
self.currency_symbol = _decode_currency_symbol(
plan_dictionary.get("currency_symbol", ""))
unlimited = plan_dictionary.get("unlimited", None)
if unlimited is not None:
self.is_unlimited = bool(unlimited)
has_team = bool(plan_dictionary.get("team_id", None))
if has_team is not None:
self.has_team = bool(has_team)
else:
self.plan_name = None
self.plan_credit = None
self.next_subscription_renewal_date = None
self.next_credit_renewal_date = None
self.subscription_state = SubscriptionState.FREE
self.period_unit = None
self.plan_price_id = None
self.plan_price = None
self.currency_code = None
self.base_price = None
self.currency_symbol = None
class PoliigonPlanUpgradeManager:
available_plans: List[Any] # List[PoliigonSubscription]
upgrade_plan: Optional[Any] = None # Optional[PoliigonSubscription]
status: Optional[PlanUpgradeStatus] = PlanUpgradeStatus.NOT_POPULATED
upgrade_info: Optional[PoliigonPlanUpgradeInfo] = None
show_banner: bool = False
upgrade_dismissed: bool = False
banner_status_emitted: Optional[PlanUpgradeStatus] = None
# Upgrade banner and popups content to be used on DCC UI;
content: Optional[Any] = None # UpgradeContent (Circular Import)
def __init__(self,
addon: Any # PoliigonAddon
):
self.addon = addon
self.user = self.addon.user
self.available_plans = []
self.set_upgrade_status()
def refresh(self,
plans_info: Optional[Dict] = None,
only_resume_popup: bool = False,
clean_plans: bool = False
) -> None:
self.user = self.addon.user
if clean_plans:
self.available_plans = []
if plans_info is not None:
self.set_available_plans(plans_info)
self.set_upgrade_plan()
self.set_upgrade_status()
self.set_show_banner()
if self.content is not None:
self.content.refresh(self, only_resume_popup)
def get_last_dismiss(self) -> Optional[datetime]:
last_dismiss = self.addon.settings_config.get(
"upgrade", "last_dismiss", fallback=None)
if last_dismiss is None:
return None
return datetime.strptime(last_dismiss, "%Y-%m-%d %H:%M:%S")
def check_last_dismiss_interval(self, day_interval: int = 7) -> bool:
last_dismiss = self.get_last_dismiss()
if last_dismiss is None:
return True
diff = datetime.now() - last_dismiss
if diff.days >= day_interval:
return True
return False
def set_show_banner(self) -> None:
if self.user is None:
return
if self.addon.user.credits is None:
return
upgrade_available = None not in [self.upgrade_info, self.upgrade_plan]
if self.status in [PlanUpgradeStatus.NOT_POPULATED,
PlanUpgradeStatus.NO_UPGRADE_AVAILABLE]:
self.show_banner = False
elif self.status in [PlanUpgradeStatus.STUDENT_DISCOUNT,
PlanUpgradeStatus.TEACHER_DISCOUNT,
PlanUpgradeStatus.BECOME_PRO,
PlanUpgradeStatus.RESUME_PLAN,
PlanUpgradeStatus.REMOVE_SCHEDULED_PAUSE,
PlanUpgradeStatus.REMOVE_CANCELLATION]:
self.show_banner = True
elif self.status == PlanUpgradeStatus.UPGRADE_PLAN_BALANCE:
self.show_banner = upgrade_available
elif self.status == PlanUpgradeStatus.UPGRADE_PLAN_UNLIMITED:
if self.addon.user.plan.plan_credit is None or not self.check_last_dismiss_interval():
self.show_banner = False
else:
self.show_banner = upgrade_available
else:
self.show_banner = False
def emit_signal(self,
view: bool = False,
dismiss: bool = False,
clicked: bool = False) -> None:
if self.status is None:
return
signal_str = self.status.get_signal_string()
if signal_str is None:
return
action_type = ActionType.OPEN_URL
if self.content is not None and self.content.open_popup:
action_type = ActionType.POPUP_MESSAGE
# Mocked Notification to be used for signals
signal_notice = Notification(id_notice=signal_str,
title=self.status.name,
priority=0,
label=self.status.name)
signal_notice.action = action_type
if dismiss and not self.upgrade_dismissed:
self.addon.notify._signal_dismiss(signal_notice)
elif view:
self.addon.notify._signal_view(signal_notice)
elif clicked:
self.addon.notify._signal_clicked(signal_notice)
def check_show_banner(self) -> bool:
if self.user is None:
return False
do_show_banner = self.show_banner
# Checks if the status changed since the last view signal
different_emit_signal_status = self.banner_status_emitted != self.status
if do_show_banner and different_emit_signal_status:
self.emit_signal(view=True)
self.banner_status_emitted = self.status
return do_show_banner
def dismiss_upgrade(self) -> None:
date_now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.addon.settings_config.set(
"upgrade", "last_dismiss", date_now)
self.set_show_banner()
self.addon._settings.save_settings()
self.emit_signal(dismiss=True)
self.upgrade_dismissed = True
def set_upgrade_status(self) -> None:
if self.user is None:
return
if self.user.credits is None:
return
subscription_state = self.user.plan.subscription_state
if subscription_state == SubscriptionState.FREE:
if self.user.is_student:
self.status = PlanUpgradeStatus.STUDENT_DISCOUNT
elif self.user.is_teacher:
self.status = PlanUpgradeStatus.TEACHER_DISCOUNT
else:
self.status = PlanUpgradeStatus.BECOME_PRO
elif subscription_state == SubscriptionState.PAUSED:
self.status = PlanUpgradeStatus.RESUME_PLAN
elif subscription_state == SubscriptionState.PAUSE_SCHEDULED:
self.status = PlanUpgradeStatus.REMOVE_SCHEDULED_PAUSE
elif subscription_state == SubscriptionState.CANCELLED:
self.status = PlanUpgradeStatus.REMOVE_CANCELLATION
elif subscription_state == SubscriptionState.ACTIVE:
if self.upgrade_plan is None:
self.status = PlanUpgradeStatus.NO_UPGRADE_AVAILABLE
elif self.user.credits == 0:
self.status = PlanUpgradeStatus.UPGRADE_PLAN_BALANCE
elif self.upgrade_plan.is_unlimited:
self.status = PlanUpgradeStatus.UPGRADE_PLAN_UNLIMITED
else:
self.status = PlanUpgradeStatus.NO_UPGRADE_AVAILABLE
else:
self.status = PlanUpgradeStatus.NO_UPGRADE_AVAILABLE
def set_available_plans(self, plans_dict: Dict) -> None:
yearly_plans = plans_dict.get("plan_year", [])
monthly_plans = plans_dict.get("plan_month", [])
# Clean available plans before populating again
self.available_plans = []
for _plan in (yearly_plans + monthly_plans):
plan_data = PoliigonSubscription()
plan_data.update_from_upgrade_dict(_plan)
self.available_plans.append(plan_data)
def set_upgrade_plan(self) -> None:
"""Method to define what is the next plan to offer to the user.
We have two main scenarios for upgrading:
Upgrade to Pro Plan: If a given user has a next Pro plan available and
(only if) their credits are empty, we offer the next pro plan
(not dismissible);
Upgrade to Unlimited:
For any Pro plan user that credits are more than zero, we should show
the Upgrade to Unlimited banner (this one is dismissible);"""
if self.user is None or len(self.available_plans) == 0:
return
if self.user.plan.is_unlimited:
# The only benefit to upgrade is to get more downloads, if you're
# already unlimited, there's nothing to upgrade to
return
if self.user.plan.has_team:
# Let's not offer updates to team members, since these contracts
# are handled separately
return
upgrade_pro_plan = None
upgrade_unlimited_plan = None
filter_period_unit = [_plan for _plan in self.available_plans
if _plan.period_unit == self.user.plan.period_unit]
filter_has_team = [_plan for _plan in filter_period_unit
if _plan.has_team == self.user.plan.has_team]
sorted_price_plans = sorted(filter_has_team, key=lambda plan: plan.plan_credit)
for _plan in sorted_price_plans:
if self.user.plan.is_unlimited and not _plan.is_unlimited:
# Don't offer credit-based plans if they are already unlimited
continue
if self.user.plan.plan_credit >= _plan.plan_credit and not _plan.is_unlimited:
# Don't offer plans which have the same or fewer credits
continue
if _plan.is_unlimited and upgrade_unlimited_plan is None:
upgrade_unlimited_plan = _plan
if upgrade_pro_plan is None:
upgrade_pro_plan = _plan
if None not in [upgrade_unlimited_plan, upgrade_pro_plan]:
break
if upgrade_pro_plan is None and upgrade_unlimited_plan is None:
self.upgrade_plan = None
elif self.user.credits == 0 and upgrade_pro_plan is not None:
self.upgrade_plan = upgrade_pro_plan
else:
self.upgrade_plan = upgrade_unlimited_plan
def finish_upgrade_plan(self) -> None:
"""This method should be called to confirm an update, resume or
choosing plan in the addon dcc side"""
if self.addon.api_rc is None:
self.addon.logger.error("API RC not defined")
return
choose_plans_status = [PlanUpgradeStatus.STUDENT_DISCOUNT,
PlanUpgradeStatus.TEACHER_DISCOUNT,
PlanUpgradeStatus.BECOME_PRO]
resume_status = [PlanUpgradeStatus.RESUME_PLAN,
PlanUpgradeStatus.REMOVE_SCHEDULED_PAUSE,
PlanUpgradeStatus.REMOVE_CANCELLATION]
upgrade_status = [PlanUpgradeStatus.UPGRADE_PLAN_BALANCE,
PlanUpgradeStatus.UPGRADE_PLAN_UNLIMITED]
if self.status in choose_plans_status:
self.addon._api.open_poliigon_link("subscribe")
elif self.status in resume_status:
callback = self.addon.api_rc._addon_params.callback_resume_plan
self.addon.api_rc.add_job_resume_plan(callback_done=callback)
elif self.status in upgrade_status:
callback = self.addon.api_rc._addon_params.callback_put_upgrade_plan
self.addon.api_rc.add_job_put_upgrade_plan(callback_done=callback)
else:
self.addon.logger.error(
f"Current user not available for upgrade: {self.user}")
@@ -0,0 +1,133 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
import os
try:
import ConfigParser
except Exception:
import configparser as ConfigParser
class PoliigonSettings():
"""Settings used for the addon."""
addon_name: str # e.g. poliigon-addon-3dsmax
base: str # Path to base directory of addon or package
software_source: str # e.g. blender
settings_filename: str
config: ConfigParser.ConfigParser = None
def __init__(self,
addon_name: str,
software_source: str,
base: str = os.path.join(os.path.expanduser("~"), "Poliigon"),
settings_filename: str = "settings.ini"):
self.addon_name = addon_name
self.base = os.path.join(base, software_source)
self.settings_filename = settings_filename
self.get_settings()
def _ensure_sections_exist(self):
sections = [
"download",
"library",
"update",
"logging",
"purchase",
"import",
"onboarding",
"ui",
"user",
"upgrade"
]
_ = [
self.config.add_section(sec)
for sec in sections
if not self.config.has_section(sec)
]
def _populate_default_settings(self):
self._ensure_sections_exist()
defaults = {
"download": {
"brush": "2K",
"download_lods": "true",
"hdri_bg": "8K",
"hdri_light": "1K",
"lod": "NONE",
"model_res": "NONE",
"tex_res": "2K"
},
"map_preferences": {},
"library": {
"primary": ""
},
"directories": {},
"logging": {
"reporting_opt_in": "true",
"verbose_logs": "true"
},
"purchase": {
"auto_download": "true"
},
"user": {
"token": "",
"first_local_asset": ""
}
}
for section_name, section_defaults in defaults.items():
if len(section_defaults.items()) == 0:
self.config.add_section(section_name)
continue
for option, value in section_defaults.items():
self.config.set(section_name, option, value)
def get_settings(self):
# https://docs.python.org/3/library/configparser.html#configparser.ConfigParser.optionxform
self.config = ConfigParser.ConfigParser()
self.config.optionxform = str
self._populate_default_settings()
settings_file = os.path.join(self.base, self.settings_filename)
if os.path.exists(settings_file):
try:
self.config.read(settings_file)
except ValueError as e:
print(f"Could not load settings for {self.addon_name}!")
print(e)
def save_settings(self):
if self.config is None:
print(f"No settings found for {self.addon_name}! Initializing...")
self.get_settings()
if not os.path.exists(self.base):
try:
os.makedirs(self.base)
except Exception as e:
print("Failed to create directory: ", e)
raise
settings_file = os.path.join(self.base, self.settings_filename)
with open(settings_file, "w+") as f:
self.config.write(f)
@@ -0,0 +1,213 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
"""Module for thread management and thread queues for Poliigon software."""
from typing import Dict, List, Optional, Union, Callable
from concurrent.futures import (CancelledError,
Future,
ThreadPoolExecutor)
from enum import Enum
import functools
import sys
import traceback
class PoolKeys(Enum):
""" Enum for the different ways to label a thread."""
INTERACTIVE = 0 # Should be the default and highest prempetive order
PREVIEW_DL = 1 # Preview thumbnails should be second order
ASSET_DL = 2 # Asset downloads lowest, don't occupy the 'last thread'
MP = 3 # Mixpannel signaling
def print_exc(fut: Future, key_pool: PoolKeys):
"""Default function to print exceptions from pool thread's done handler."""
try:
exc = fut.exception()
except CancelledError:
exc = None
if exc is None:
return
print((f"=== ThreadManager[{key_pool.name}]: Thread Exception "
f"({exc.__class__.__name__}): {exc}"))
traceback.print_tb(exc.__traceback__)
class ThreadManager:
"""The class which manages state of the threads.
ThreadPools are created upon first use.
Number of threads per pool can be set "globally" upon creation
of the ThreadPoolManager or per pool, when a pool is used the first time.
Decorator to be implemented in a class using the ThreadManager.
Parameters pool and foreground are explained in detail for queue_thread().
The code expects the ThreadManager instance in a member variable tm.
Adapt as needed:
def run_threaded(key_pool: PoolKeys,
max_threads: Optional[int] = None,
foreground: bool = False) -> callable:
# Schedule a function to run in a thread of a chosen pool
def wrapped_func(func: callable) -> callable:
@functools.wraps(func)
def wrapped_func_call(self, *args, **kwargs):
args = (self, ) + args
return self.tm.queue_thread(func, key_pool, max_threads,
foreground, *args, **kwargs)
return wrapped_func_call
return wrapped_func
"""
max_threads: int # "global" max_threads, used if not overriden
thread_pools: Dict[PoolKeys, ThreadPoolExecutor] = {}
# function from the reporting addon side to report Sentry messages from
# threaded functions. Expected to receive as parameter the function name
# and a partial of the function to be threaded
reporting_callable: Optional[Callable] = None
def __init__(self,
max_threads: int = 10,
callback_print_exc: Optional[Callable] = None,
):
"""Arguments:
print_exc: Callable to be used instead of the default print_exc function.
The callable needs to have the following interface:
print_exc(fut: Future, key_pool: PoolKeys)
Partial wrap if more parameters needed.
"""
self.thread_pools = {}
self.max_threads = max_threads
if callback_print_exc is None:
self.print_exc = print_exc
else:
self.print_exc = callback_print_exc
def get_pool(self,
key_pool: PoolKeys,
max_threads: Optional[int] = None,
no_create: bool = False
) -> Optional[ThreadPoolExecutor]:
"""Returns the thread pool for a given key.
If the pool does not exist, yet, it will be created unless
no_create is set to True, in which case None gets returned.
No need to call exernally.
"""
if key_pool in self.thread_pools:
return self.thread_pools[key_pool]
if no_create:
return None
if max_threads is None:
max_threads = self.max_threads
tpe = ThreadPoolExecutor(max_workers=max_threads)
self.thread_pools[key_pool] = tpe
return tpe
def queue_thread(self,
func: callable,
key_pool: Optional[PoolKeys] = None,
max_threads: Optional[int] = None,
foreground: bool = False,
*args, **kwargs) -> Union[Future, any]:
"""Enqueue a function for threaded execution via a thread pool.
Parameters:
key_pool: Selects the pool to be used, see PoolKeys enum.
max_threads: The maximum number of threads can only be set once upon
pool's first usage. It can not be changed later on.
foreground: Set to True to have the function directly executed
instead of being submitted to a thread pool.
Return value:
Usually the Future belonging to a scheduled thread.
If foreground option is used, it may actually be anything,
as the return value of the function gets returned directly.
"""
if max_threads is None or max_threads <= 0:
max_threads = self.max_threads
if key_pool is None:
key_pool = PoolKeys.INTERACTIVE
report_func = None
if self.reporting_callable is not None:
partial_func = functools.partial(func, *args, **kwargs)
report_func = self.reporting_callable(func.__name__, partial_func)
if foreground:
# With foreground option the function gets called directly
# NOTE: When using foreground option, the function returns
# the return value of the called function instead of a Future
if report_func is not None:
fut = report_func()
else:
fut = func(*args, **kwargs)
else:
# Create ThreadPoolExecutor, if not already in thread_pools dict
thread_pool = self.get_pool(key_pool, max_threads)
# Finally, kick the can
# Schedule the function for threaded execution
if report_func is not None:
fut = thread_pool.submit(report_func)
else:
fut = thread_pool.submit(func, *args, **kwargs)
func_print = functools.partial(self.print_exc,
key_pool=key_pool)
fut.add_done_callback(func_print)
return fut
def shutdown(self,
key_pool: Optional[PoolKeys] = None,
wait: bool = True) -> None:
"""Shutdown one or all (key_pool=None) ThreadPoolExecutors."""
if key_pool is None:
for tpe in self.thread_pools.values():
if sys.version_info >= (3, 8, 0):
tpe.shutdown(wait=wait, cancel_futures=True)
else:
tpe.shutdown(wait=wait)
self.thread_pools = {}
elif key_pool in self.thread_pools:
self.thread_pools[key_pool].shutdown(wait=wait)
del self.thread_pools[key_pool]
def pool_keys(self) -> List[PoolKeys]:
"""Returns a list containing the pool keys of current pools."""
return list(self.thread_pools.keys())
def number_of_pools(self) -> int:
"""Returns the number of currently active ThreadPoolExecutors.
This does NOT mean, these ThreadPoolExecutors are currently
actively executing threads.
"""
return len(self.thread_pools)
@@ -0,0 +1,491 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
"""Module for general purpose updating for Poliigon software."""
from typing import Dict, Optional, Sequence, Tuple, Callable, Any
from dataclasses import dataclass
import datetime
import json
import os
import threading
import requests
from .multilingual import _m
from .notifications import (Notification,
NotificationSystem,
NOTICE_TITLE_UPDATE)
BASE_URL = "https://software.poliigon.com"
TIMEOUT = 20.0
# Status texts
FAIL_GET_VERSIONS = _m("Failed to get versions")
def v2t(value: str) -> tuple:
"""Take a version string like v1.2.3 and convert it to a tuple."""
if not value or "." not in value:
return None
if value.lower().startswith("v"):
value = value[1:]
return tuple([int(ind) for ind in value.split(".")])
def t2v(ver: tuple) -> str:
"""Take a tuple like (2, 80) and construct a string like v2.80."""
return "v" + ".".join(list(ver))
@dataclass
class AlertData:
title: Optional[str] = None
label: Optional[str] = None
body: Optional[str] = None
url: Optional[str] = None
priority: Optional[int] = None
action_string: Optional[str] = None
open_popup: Optional[bool] = None
allow_dismiss: bool = True
auto_dismiss: bool = True
valid: bool = True
def validate_field(
self,
value: Any,
field: str,
field_type: type,
mandatory: bool,
report_callable: Optional[Callable] = None):
exists = value is not None
ok_exists = (exists or not mandatory)
ok_type = True
if exists:
ok_type = isinstance(value, field_type)
if not ok_exists or not ok_type:
self.valid = False
if report_callable is not None:
report_callable(
"invalid_alert_information",
f"Invalid {field} {value}",
"error")
def validate_data(self, report_callable: Optional[Callable] = None) -> bool:
rc = report_callable
self.validate_field(self.title, "Title", str, True, rc)
self.validate_field(self.label, "Label", str, True, rc)
self.validate_field(self.priority, "Priority", int, True, rc)
self.validate_field(self.url, "Url", str, False, rc)
self.validate_field(self.action_string, "Action String", str, False, rc)
self.validate_field(self.auto_dismiss, "Auto Dismiss", bool, False, rc)
self.validate_field(self.allow_dismiss, "Allow Dismiss", bool, False, rc)
if self.url is None:
# one of url or body have to be available with right format
self.validate_field(self.body, "Body", str, True, rc)
def update_from_dict(
self, data: Dict, report_callable: Optional[Callable] = None) -> None:
self.title = data.get("title")
self.label = data.get("label")
self.body = data.get("body")
self.url = data.get("url")
self.action_string = data.get("action_string")
self.priority = data.get("priority")
self.allow_dismiss = data.get("allow_dismiss", True)
self.auto_dismiss = data.get("auto_dismiss", True)
if self.url is None or self.url == "":
self.open_popup = True
self.validate_data(report_callable)
def create_notification(
self, notification_system: NotificationSystem) -> Optional[Notification]:
if not self.valid:
return None
notice = notification_system.create_version_alert(
title=self.title,
priority=self.priority,
label=self.label,
tooltip=self.label,
body=self.body,
action_string=self.action_string,
url=self.url,
open_popup=self.open_popup,
allow_dismiss=self.allow_dismiss,
auto_dismiss=self.auto_dismiss
)
return notice
@dataclass
class VersionData:
"""Container for a single version of the software."""
version: Optional[tuple] = None
url: Optional[str] = None
min_software_version: Optional[tuple] = None # Inclusive.
max_software_version: Optional[tuple] = None # Not inclusive.
required: Optional[bool] = None
release_timestamp: Optional[datetime.datetime] = None
alert: Optional[AlertData] = None
# Internal, huamn readable current status.
status_title: str = ""
status_details: str = ""
status_ok: bool = True
# Reporting rate for the version
error_sample_rate: Optional[float] = None
traces_sample_rate: Optional[float] = None
def update_from_dict(
self, data: Dict, report_callable: Optional[Callable] = None) -> None:
self.version = v2t(data.get("version"))
self.url = data.get("url", "")
# List format like [2, 80]
self.min_software_version = tuple(data.get("min_software_version"))
self.max_software_version = tuple(data.get("max_software_version"))
self.required = data.get("required")
self.release_timestamp = data.get("release_timestamp")
alert_data = data.get("alert", None)
if alert_data is not None:
self.alert = AlertData()
self.alert.update_from_dict(alert_data, report_callable)
self.error_sample_rate = (data.get("error_sample_rate"))
self.traces_sample_rate = (data.get("traces_sample_rate"))
def create_alert_notification(
self, notification_system: NotificationSystem) -> Optional[Notification]:
if self.alert is None or notification_system is None:
return
return self.alert.create_notification(notification_system)
def create_update_notification(
self, notification_system: NotificationSystem) -> Optional[Notification]:
if self.url is None or notification_system is None:
return
version = str(self.version)
version = version.replace(", ", ".")
label = f"{NOTICE_TITLE_UPDATE} {version}"
notice = notification_system.create_update(
tooltip=NOTICE_TITLE_UPDATE,
label=label,
download_url=self.url
)
return notice
class SoftwareUpdater():
"""Primary class which implements checks for updates and installs."""
# Versions of software available.
stable: Optional[VersionData]
latest: Optional[VersionData]
all_versions: Sequence
# Always initialized
addon_name: str # e.g. poliigon-addon-blender.
addon_version: tuple # Current addon version.
software_version: tuple # DCC software version, e.g. (3, 0).
base_url: str # Primary url where updates and version data is hosted.
# State properties.
update_ready: Optional[bool] = None # None until proven true or false.
update_data: Optional[VersionData] = None
_last_check: Optional[datetime.datetime] = None
last_check_callback: Optional[Callable] = None # When last_check changes.
check_interval: Optional[int] = None # interval in seconds between auto check.
verbose: bool = True
# Classes to be imported from the addon
notification_system: Optional[NotificationSystem] = None
reporting_callable: Optional[Callable] = None
# Notifications
alert_notice: Optional[Notification] = None
update_notice: Optional[Notification] = None
# Bool value to be set by addon to take the update
# data from the latest version instead of the stable one
update_from_latest: bool = False
_check_thread: Optional[threading.Thread] = None
def __init__(self,
addon_name: str,
addon_version: tuple,
software_version: tuple,
base_url: Optional[str] = None,
notification_system: Optional[NotificationSystem] = None,
local_json: Optional[str] = None):
self.addon_name = addon_name
self.addon_version = addon_version
self.notification_system = notification_system
self.software_version = software_version
self.base_url = base_url if base_url is not None else BASE_URL
self.local_json = local_json
self.current_version = VersionData()
self._clear_versions()
@property
def is_checking(self) -> bool:
"""Interface for other modules to see if a check for update running."""
return self._check_thread and self._check_thread.is_alive()
@property
def last_check(self) -> str:
if not self._last_check:
return ""
try:
return self._last_check.strftime("%Y-%m-%d %H:%M")
except ValueError as err:
print("Get last update check error:", err)
return ""
@last_check.setter
def last_check(self, value: str) -> None:
try:
self._last_check = datetime.datetime.strptime(
value, "%Y-%m-%d %H:%M")
except ValueError as err:
print("Assign last update check error:", value, err)
print(err)
self._last_check = None
if self.last_check_callback:
self.last_check_callback(self.last_check) # The string version.
def _clear_versions(self) -> None:
self.stable = None
self.latest = None
self.all_versions = []
def _clear_update(self) -> None:
self.update_ready = None # Set to None until proven true or false.
self.update_data = None
self.status_ok = True
def has_time_elapsed(self, hours: int = 24) -> bool:
"""Checks if a given number of hours have passed since last check."""
now = datetime.datetime.now()
if not self._last_check:
return True # No check on record.
diff = now - self._last_check
return diff.total_seconds() / 3600.0 > hours
def print_debug(self, *args):
if self.verbose:
print(*args)
def update_versions(self) -> None:
"""Fetch the latest versions available from the server."""
self.status_ok = True # True until proven false.
self._clear_versions()
url = f"{self.base_url}/{self.addon_name}-versions.json"
try:
res = requests.get(url, timeout=TIMEOUT)
except requests.exceptions.ConnectionError:
self.status_title = FAIL_GET_VERSIONS
self.status_ok = False
self.status_details = "Updater ConnectionError"
return
except requests.exceptions.Timeout:
self.status_title = FAIL_GET_VERSIONS
self.status_ok = False
self.status_details = "Updater Timeout"
return
except requests.exceptions.ProxyError:
self.status_title = FAIL_GET_VERSIONS
self.status_ok = False
self.status_details = "Updater ProxyError"
return
if not res.ok:
self.status_title = FAIL_GET_VERSIONS
self.status_details = (
"Did not get OK response while fetching available versions "
f"from {url}")
self.status_ok = False
print(self.status_details)
return
if res.status_code != 200:
self.status_title = FAIL_GET_VERSIONS
self.status_details = (
"Did not get OK code while fetching available versions")
self.status_ok = False
print(self.status_details)
return
try:
resp = json.loads(res.text)
if self.local_json is not None and os.path.isfile(self.local_json):
with open(self.local_json) as f:
resp = json.load(f)
except json.decoder.JSONDecodeError as e:
self.status_title = FAIL_GET_VERSIONS
self.status_details = "Could not parse json response for versions"
self.status_ok = False
self.status_is_error = True
print(self.status_details)
print(e)
return
if resp.get("stable"):
self.stable = VersionData()
self.stable.update_from_dict(resp["stable"])
if resp.get("latest"):
self.latest = VersionData()
self.latest.update_from_dict(resp["latest"])
if resp.get("versions"):
for itm in resp["versions"]:
ver = VersionData()
ver.update_from_dict(itm, self.reporting_callable)
self.all_versions.append(ver)
if ver.version == self.addon_version:
self.current_version = ver
self._last_check = datetime.datetime.now()
self.last_check = self.last_check # Trigger callback.
def _update_notification_msg(self) -> None:
if self.notification_system is None:
return
body = self.notification_system.addon_params.update_body
if self.update_data is not None and body is not None:
version = self.update_data.version
version = ".".join(map(str, version))
self.notification_system.addon_params.update_body = body.format(
version)
def _create_notifications(self) -> Tuple[Notification, Notification]:
alert_notif = None
update_notif = None
self._update_notification_msg()
if self.current_version.alert is not None:
alert_notif = self.current_version.create_alert_notification(
self.notification_system)
if self.update_data is not None:
update_notif = self.update_data.create_update_notification(
self.notification_system)
return update_notif, alert_notif
def check_for_update(self,
callback: Optional[callable] = None,
create_notifications: bool = False) -> bool:
"""Fetch and check versions to see if a new update is available."""
self._clear_update()
self.update_versions()
if not self.status_ok:
if callback:
callback()
return False
# First compare against latest
if self.stable and self._check_eligible(self.stable):
update_version = self.stable
if self.update_from_latest:
update_version = self.latest
self.print_debug(
"Using latest stable:",
update_version.version,
"vs current addon: ",
self.addon_version)
if update_version.version > self.addon_version:
self.update_data = update_version
self.update_ready = True
else:
self.update_ready = False
if create_notifications and self.notification_system is not None:
self.update_notice, self.alert_notice = self._create_notifications()
if callback:
callback()
return True
# Eligible wasn't present or more eligible, find next best.
self.print_debug("Unable to use current stable release")
max_version = self.get_max_eligible()
if max_version:
if max_version.version > self.addon_version:
self.update_data = max_version
self.update_ready = True
else:
self.update_ready = False
else:
self.print_debug("No eligible releases found")
self.update_ready = False
if create_notifications and self.notification_system is not None:
self.update_notice, self.alert_notice = self._create_notifications()
if callback is not None:
callback()
return True
def _check_eligible(self, version: VersionData) -> bool:
"""Verify if input version is compatible with the current software."""
eligible = True
if version.min_software_version:
if self.software_version < version.min_software_version:
eligible = False
elif version.max_software_version:
# Inclusive so that if max is 3.0, must be 2.99 or lower.
if self.software_version >= version.max_software_version:
eligible = False
return eligible
def get_max_eligible(self) -> Optional[VersionData]:
"""Find the eligible version with the highest version number."""
max_eligible = None
for ver in self.all_versions:
if not self._check_eligible(ver):
continue
elif max_eligible is None:
max_eligible = ver
elif ver.version > max_eligible.version:
max_eligible = ver
return max_eligible
def async_check_for_update(
self, callback: Callable = None, create_notifications: bool = False):
"""Start a background thread which will check for updates."""
if self.is_checking:
return
self._check_thread = threading.Thread(
target=self.check_for_update,
args=(callback, create_notifications))
self._check_thread.daemon = True
self._check_thread.start()
@@ -0,0 +1,498 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, Optional, Tuple
from .plan_manager import (PoliigonPlanUpgradeManager,
PlanUpgradeStatus,
PoliigonSubscription)
from .multilingual import _t
@dataclass
class UpgradeIcons:
check: str
info: str
unlimited: str
class UpgradeContent:
"""Class to be created in the DCC side which will define and store all
UI content information."""
upgrade_manager: PoliigonPlanUpgradeManager
current_plan: PoliigonSubscription
banner_primary_text: str = ""
banner_secondary_text: str = ""
banner_button_text: str = ""
allow_dismiss: bool = False
open_popup: bool = False
icon_path: Optional[str] = None
# For upgrade popup
upgrade_popup_title: Optional[str] = None
upgrade_popup_table: Optional[Dict[str, str]] = None
upgrade_popup_key_value: Optional[Dict[str, str]] = None
upgrade_popup_text: Optional[str] = None
upgrade_popup_confirm_button: Optional[str] = None
upgrade_popup_pricing_button: Optional[str] = None
upgrade_popup_terms_button: Optional[str] = None
# Upgrading process messages
upgrading_primary_text: Optional[str] = None
upgrading_secondary_text: Optional[str] = None
# For success popups
success_popup_title: Optional[str] = None
success_popup_text: Optional[str] = None
# For error popups
error_popup_title: Optional[str] = None
error_popup_text: Optional[str] = None
# Flags to be used for P4B UI
as_single_paragraph: bool = False
# Flag to use in P4B & P4C UI
use_single_policy_link: bool = True
# Stores the icon paths for each dcc
icons: Optional[UpgradeIcons] = None
def __init__(self,
upgrade_manager: PoliigonPlanUpgradeManager,
as_single_paragraph: bool = False,
use_single_policy_link: bool = True,
icons: Optional[Tuple[str, str, str]] = None):
"""Class to handle all the Content for Upgrade UI in each DCC.
Parameters:
upgrade_manager: addon.upgrade_manager instance of PoliigonPlanUpgradeManager;
as_single_paragraph: If True, the banner_secondary_text will be set as
None, and all the text will be represented as
a single paragraph in banner_primary_text;
icons: The icon paths to be used in upgrade manager (order: check, info, unlimited)
NOTE: This class instance should be created in the addon side and
stored in addon.upgrade_manager.content;
"""
self.as_single_paragraph = as_single_paragraph
self.use_single_policy_link = use_single_policy_link
if icons is not None:
# Icon paths order: check, info, unlimited
self.icons = UpgradeIcons(icons[0], icons[1], icons[2])
self.refresh(upgrade_manager)
def refresh(self,
upgrade_manager: PoliigonPlanUpgradeManager,
only_resume_popup: bool = False
) -> None:
if upgrade_manager is not None:
self.upgrade_manager = upgrade_manager
if self.upgrade_manager is None:
return
if self.upgrade_manager.addon.user is None:
return
self.current_plan = self.upgrade_manager.addon.user.plan
if only_resume_popup:
# When resuming, the renewal date is updated, so in this
# scenario we use this flag to only update the popup. If populate()
# was called, now the status would be NO_UPGRADE_AVAILABLE and
# the popup message will be another
self.set_resume_success_popup()
return
self.populate()
def student_discount(self, is_teacher: bool = False) -> Any:
primary = _t("Access the entire library by joining Pro")
secondary = _t("{0} can claim a 50% discount".format(
_t("Students") if not is_teacher else _t("Teachers")))
if self.as_single_paragraph:
self.banner_primary_text = f"{primary}. {secondary}"
self.banner_secondary_text = None
else:
self.banner_primary_text = primary
self.banner_secondary_text = secondary
self.banner_button_text = _t("Choose Your Plan")
self.allow_dismiss = False
self.open_popup = False
if self.icons is not None:
self.icon_path = self.icons.check
def become_pro(self) -> Any:
primary = _t("Access the entire library by joining Pro")
secondary = _t("Download and import from the entire Poliigon library")
if self.as_single_paragraph:
# To keep it slimmer, do just the primary text.
self.banner_primary_text = f"{primary}"
self.banner_secondary_text = None
else:
self.banner_primary_text = primary
self.banner_secondary_text = secondary
self.banner_button_text = _t("Choose Your Plan")
self.allow_dismiss = False
self.open_popup = False
if self.icons is not None:
self.icon_path = self.icons.check
def upgrade_balance(self) -> Any:
primary_text, secondary_text = self._get_upgrade_text_upgrade_balance()
self.banner_primary_text = primary_text
self.banner_secondary_text = secondary_text
self.banner_button_text = _t("Get More Downloads")
self.allow_dismiss = False
self.open_popup = True
if self.icons is not None:
self.icon_path = self.icons.info
def resume_plan(self) -> Any:
primary_text, secondary_text = self._get_upgrade_text_paused_until()
self.banner_primary_text = primary_text
self.banner_secondary_text = secondary_text
self.banner_button_text = _t("Resume Plan")
self.allow_dismiss = False
self.open_popup = True
if self.icons is not None:
self.icon_path = self.icons.info
def remove_cancel(self) -> Any:
primary_text, secondary_text = self._get_upgrade_text_term_end()
self.banner_primary_text = primary_text
self.banner_secondary_text = secondary_text
self.banner_button_text = _t("Resume Plan")
self.allow_dismiss = False
self.open_popup = True
if self.icons is not None:
self.icon_path = self.icons.info
def remove_pause(self) -> Any:
primary_text, secondary_text = self._get_upgrade_text_paused_at()
self.banner_primary_text = primary_text
self.banner_secondary_text = secondary_text
self.banner_button_text = _t("Cancel Pause")
self.allow_dismiss = False
self.open_popup = True
if self.icons is not None:
self.icon_path = self.icons.info
def upgrade_unlimited(self) -> Any:
primary_text, secondary_text = self._get_upgrade_text_unlimited()
self.banner_primary_text = primary_text
self.banner_secondary_text = secondary_text
self.banner_button_text = _t("Upgrade to Unlimited")
self.allow_dismiss = True
self.open_popup = True
if self.icons is not None:
self.icon_path = self.icons.unlimited
return self
def _get_upgrade_text_upgrade_balance(self) -> Tuple[str, Optional[str]]:
"""Returns text to display in case of a scheduled pause subscription."""
next_renewal_date = self.current_plan.next_subscription_renewal_date
diff = next_renewal_date - datetime.now()
head = _t("Youre out of downloads")
text = _t("Youll get more in {0} days or upgrade "
"to download now").format(diff.days)
if self.as_single_paragraph:
return f"{head}. {text}", None
else:
return head, text
def _get_upgrade_text_paused_at(self) -> Tuple[str, Optional[str]]:
"""Returns text to display in case of a scheduled pause subscription."""
pause_date = self.current_plan.plan_paused_at
date_paused_until = None
if pause_date is not None:
date_paused_until = self.current_plan.plan_paused_at.strftime("%d %b %Y")
head = _t("Your plan will pause on {0}").format(date_paused_until)
text = _t("Cancel pause to keep downloading")
if self.as_single_paragraph:
return f"{head}. {text}", None
else:
return head, text
def _get_upgrade_text_paused_until(self) -> Tuple[str, Optional[str]]:
"""Returns text to display in case of a paused subscription."""
date_paused_until = self.current_plan.plan_paused_until.strftime("%d %b %Y")
head = _t("Your plan is paused until {0}").format(date_paused_until)
text = _t("Resume your plan to download new assets")
if self.as_single_paragraph:
return f"{head}. {text}", None
else:
return head, text
def _get_upgrade_text_term_end(self) -> Tuple[str, Optional[str]]:
"""Returns text to display in case of a cancelled subscription."""
date_term_end = self.current_plan.current_term_end.strftime("%d %b %Y")
head = _t("Your plan will end on {0}").format(date_term_end)
text = _t("Resume your plan to keep downloading")
if self.as_single_paragraph:
return f"{head}. {text}", None
else:
return head, text
def _get_upgrade_text_unlimited(self) -> Tuple[str, Optional[str]]:
"""Returns text to display in case of non-unlimited subscription."""
head = _t("Need more downloads?")
text = _t("Upgrade to unlimited and never worry about limits again")
if self.as_single_paragraph:
# Per Blender design, don't include the second bit of text
return f"{head}", None
else:
return head, text
def _get_text_price_change(self) -> str:
price_old = self.current_plan.base_price
price_new = self.upgrade_manager.upgrade_plan.base_price
currency_code = self.upgrade_manager.upgrade_info.currency_code
currency_symbol = self.upgrade_manager.upgrade_info.currency_symbol
if price_old is not None:
price_old = f"{currency_symbol}{price_old:.2f} {currency_code}"
if price_new is not None:
price_new = f"{currency_symbol}{price_new:.2f} {currency_code}"
return f"{price_old} \u2192 {price_new}"
def _get_text_licence(self) -> str:
"""Decodes boolean has_team into 'Team' or 'Individual'."""
has_team = self.upgrade_manager.upgrade_plan.has_team
if has_team:
text_licence = _t("Team")
else:
text_licence = _t("Individual")
return text_licence
def _get_text_billing_period(self) -> str:
"""Decodes period_unit into 'Yearly' or 'Monthly'."""
period_unit = self.upgrade_manager.upgrade_plan.period_unit
if period_unit == "year":
text_billing = _t("Yearly")
elif period_unit == "month":
text_billing = _t("Monthly")
else:
text_billing = period_unit
return text_billing
def _get_text_assets_change(self) -> str:
"""Returns the change in assets count as string ('previous -> new')."""
new_assets = self.upgrade_manager.upgrade_info.new_assets
prev_assets = self.upgrade_manager.upgrade_info.previous_assets
text_assets = f"{prev_assets} \u2192 {new_assets}"
return text_assets
def _get_text_users_change(self) -> str:
"""Returns the change in user count as string ('previous -> new')."""
new_users = self.upgrade_manager.upgrade_info.new_users
previous_users = self.upgrade_manager.upgrade_info.previous_users
text_users = f"{previous_users} \u2192 {new_users}"
return text_users
def _get_text_amount_due(self) -> str:
"""Returns amount due as string with currency code and symbol."""
amount_due = self.upgrade_manager.upgrade_info.amount_due
currency_code = self.upgrade_manager.upgrade_info.currency_code
currency_symbol = self.upgrade_manager.upgrade_info.currency_symbol
text_amount_due = f"{currency_symbol}{amount_due} {currency_code}"
return text_amount_due
def _get_text_amount_due_renewal(self) -> str:
"""Returns amount due on renewal as string with
currency code and symbol.
"""
amount_due_renewal = self.upgrade_manager.upgrade_info.amount_due_renewal
currency_code = self.upgrade_manager.upgrade_info.currency_code
currency_symbol = self.upgrade_manager.upgrade_info.currency_symbol
text_amount_due_renewal = (f"{currency_symbol}{amount_due_renewal} "
f"{currency_code}")
return text_amount_due_renewal
def set_resume_popup_information(self):
# TODO: Check string phrasing here (Maybe different texts for each one of the scenarios of resuming)
self.upgrade_popup_text = _t("Would you like to resume your plan? "
"You will be charged for renewal and can "
"start downloading straight away.")
self.upgrade_popup_title = _t("Resume Plan")
self.upgrade_popup_confirm_button = _t("Resume Plan")
self.upgrading_primary_text = _t("Resuming plan...")
self.upgrading_secondary_text = _t("This may take a few seconds.")
self.upgrade_popup_table = None
self.upgrade_popup_key_value = None
self.upgrade_popup_pricing_button = None
self.upgrade_popup_terms_button = None
def set_remove_scheduled_pause_popup_information(self):
# TODO: Check string phrasing here (Maybe different texts for each one of the scenarios of resuming)
self.upgrade_popup_text = _t("Would you like to remove the scheduled pause? ")
self.upgrade_popup_title = _t("Cancel Pause")
self.upgrade_popup_confirm_button = _t("Cancel Pause")
self.upgrading_primary_text = _t("Cancelling Pause...")
self.upgrading_secondary_text = _t("This may take a few seconds.")
self.upgrade_popup_table = None
self.upgrade_popup_key_value = None
self.upgrade_popup_pricing_button = None
self.upgrade_popup_terms_button = None
def set_remove_scheduled_cancel_popup_information(self):
# TODO: Check string phrasing here (Maybe different texts for each one of the scenarios of resuming)
self.upgrade_popup_text = _t("Would you like to remove the scheduled cancellation?")
self.upgrade_popup_title = _t("Remove Cancellation")
self.upgrade_popup_confirm_button = _t("Remove Cancellation")
self.upgrading_primary_text = _t("Removing Cancellation...")
self.upgrading_secondary_text = _t("This may take a few seconds.")
self.upgrade_popup_table = None
self.upgrade_popup_key_value = None
self.upgrade_popup_pricing_button = None
self.upgrade_popup_terms_button = None
def set_upgrade_popup_information(self):
self.upgrade_popup_table = {
_t("Assets per month:"): self._get_text_assets_change(),
# The following line are commented due to a decision of not showing
# team related information on confirmation popup;
# _t("Users:"): self._get_text_users_change(),
# _t("License:"): self._get_text_licence(),
_t(f"{self._get_text_billing_period()} price:"): self._get_text_price_change(),
_t("Starts:"): _t("Today"),
_t("Billing frequency:"): self._get_text_billing_period(),
_t("Renewal:"): self.upgrade_manager.upgrade_info.renewal_date
}
self.upgrade_popup_key_value = {
_t("Due today:"): self._get_text_amount_due(),
_t("Due on renewal:"): self._get_text_amount_due_renewal()
}
if self.use_single_policy_link:
confirm_text = _t("By confirming you agree to our Unlimited Fair Use "
"Policy, Terms & Conditions, Privacy & Refund Policy below.")
else:
confirm_text = _t("By confirming you agree to our Unlimited Fair Use "
"Policy, Terms & Conditions, Privacy & Refund Policy.")
tax_text = ""
if self.upgrade_manager.upgrade_info.tax_rate not in [None, 0]:
tax_text = _t("Due today and renewal prices include {0}% tax. ").format(
self.upgrade_manager.upgrade_info.tax_rate)
self.upgrade_popup_text = f"{tax_text}{confirm_text}"
self.upgrade_popup_title = _t("Change Plan")
self.upgrade_popup_confirm_button = _t("Confirm Plan Change")
self.upgrade_popup_pricing_button = _t("View All Pricing")
self.upgrade_popup_terms_button = _t("Terms & Policy Documents")
def set_upgrade_success_popup(self):
self.success_popup_title = _t("Plan Change Successful")
self.success_popup_text = _t("You have successfully updated your plan.")
# TODO(Joao): Implement different hard coded error messages for each
# scenario - look for strings in the api error (maybe error codes)
self.error_popup_title = _t("Error Upgrading Plan")
self.error_popup_text = _t("Upgrade Plan Failed. \n\n{0}\n\n"
"Try again later or reach out to support.")
self.upgrading_primary_text = _t("Upgrading Plan...")
self.upgrading_secondary_text = _t("This may take a few seconds.")
def set_resume_success_popup(self):
# TODO: Check string phrasing here (Just a mock str for now)
renewal_date = self.current_plan.next_subscription_renewal_date
text = _t("Your plan has successfully resumed")
if renewal_date is not None:
renewal_text = self.current_plan.next_subscription_renewal_date.strftime("%d %b %Y")
renewal_date_text = _t("and will renew on {0}").format(renewal_text)
text = f"{text} {renewal_date_text}"
else:
text = f"{text}."
self.success_popup_title = _t("Plan Resumed")
self.success_popup_text = text
self.error_popup_title = _t("Error Resuming Plan")
self.error_popup_text = _t("Resume Plan Failed. \n\n{0}\n\n"
"Try again later or reach out to support.")
def populate(self) -> None:
upgrade_satus = self.upgrade_manager.status
if upgrade_satus == PlanUpgradeStatus.STUDENT_DISCOUNT:
self.student_discount()
if upgrade_satus == PlanUpgradeStatus.TEACHER_DISCOUNT:
self.student_discount(is_teacher=True)
elif upgrade_satus == PlanUpgradeStatus.BECOME_PRO:
self.become_pro()
elif upgrade_satus == PlanUpgradeStatus.UPGRADE_PLAN_BALANCE:
self.upgrade_balance()
if self.upgrade_manager.upgrade_info is None:
return
self.set_upgrade_popup_information()
self.set_upgrade_success_popup()
elif upgrade_satus == PlanUpgradeStatus.RESUME_PLAN:
self.resume_plan()
self.set_resume_popup_information()
self.set_resume_success_popup()
elif upgrade_satus == PlanUpgradeStatus.REMOVE_SCHEDULED_PAUSE:
self.remove_pause()
self.set_remove_scheduled_pause_popup_information()
self.set_resume_success_popup()
elif upgrade_satus == PlanUpgradeStatus.REMOVE_CANCELLATION:
self.remove_cancel()
self.set_remove_scheduled_cancel_popup_information()
self.set_resume_success_popup()
elif upgrade_satus == PlanUpgradeStatus.UPGRADE_PLAN_UNLIMITED:
self.upgrade_unlimited()
if self.upgrade_manager.upgrade_info is None:
return
self.set_upgrade_popup_information()
self.set_upgrade_success_popup()
@@ -0,0 +1,129 @@
# #### BEGIN GPL LICENSE BLOCK #####
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
from dataclasses import dataclass
from typing import Dict, List, Optional
from .assets import (MapType)
from .plan_manager import PoliigonSubscription
from .logger import (DEBUG, # noqa F401, allowing downstream const usage
ERROR,
INFO,
get_addon_logger,
NOT_SET,
WARNING)
@dataclass
class MapFormats:
map_type: MapType
default: str
required: bool
extensions: Dict[str, bool]
enabled: bool
selected: Optional[str] = None
class UserDownloadPreferences:
resolution_options: List[str]
default_resolution: str
selected_resolution: Optional[str] = None
texture_maps: List[MapFormats]
software_selected: Optional[str] = None
render_engine_selected: Optional[str] = None
lod_options: List[str]
lod_selected: Optional[str] = None
def __init__(self, res: Dict):
self.res = res
self.set_resolution()
self.set_lods()
self.set_software()
self.set_texture_maps()
def set_resolution(self) -> None:
resolution_info = self.res.get("default_resolution", {})
self.resolution_options = resolution_info.get("resolution_options")
self.default_resolution = resolution_info.get("default")
self.selected_resolution = resolution_info.get("selected")
def set_lods(self) -> None:
lods_info = self.res.get("lods", {})
self.lod_options = lods_info.get("lod_options")
self.lod_selected = lods_info.get("selected")
def set_software(self) -> None:
software_info = self.res.get("softwares", {})
for _soft, soft_inf in software_info.items():
soft_selected = soft_inf.get("selected", None)
renderer_selected = soft_inf.get("selected_render_engine", None)
if soft_selected is not None and renderer_selected is not None:
self.software_selected = soft_selected
self.render_engine_selected = renderer_selected
break
def set_texture_maps(self) -> None:
self.texture_maps = []
texture_maps_info = self.res.get("texture_maps", {})
for _map, map_info in texture_maps_info.items():
map_type = MapType.from_type_code(_map)
enabled = map_info.get("selected") is not None
map_format = MapFormats(map_type=map_type,
default=map_info.get("default"),
enabled=enabled,
selected=map_info.get("selected"),
required=map_info.get("required"),
extensions=map_info.get("formats"))
self.texture_maps.append(map_format)
def string_stamp(self) -> str:
string_stamp = ""
for _map in self.texture_maps:
string_stamp += f"{_map.map_type.name}:{str(_map.selected)};"
return string_stamp
def get_map_preferences(self, map_type: MapType) -> Optional[MapFormats]:
for _map in self.texture_maps:
if _map.map_type.get_effective() == map_type.get_effective():
return _map
return None
def get_all_maps_enabled(self):
return [_map for _map in self.texture_maps if _map.enabled]
@dataclass
class PoliigonUser:
"""Container object for a user."""
user_name: str
user_id: int
is_student: Optional[bool] = False
is_teacher: Optional[bool] = False
credits: Optional[int] = None
credits_od: Optional[int] = None
plan: Optional[PoliigonSubscription] = None
map_preferences: Optional[UserDownloadPreferences] = None
# Todo(Joao): remove this flag when all addons are using map prefs
use_preferences_on_download: bool = False
@@ -0,0 +1,57 @@
from sentry_sdk.scope import Scope
from sentry_sdk.transport import Transport, HttpTransport
from sentry_sdk.client import Client
from sentry_sdk.api import * # noqa
from sentry_sdk.consts import VERSION # noqa
__all__ = [ # noqa
"Hub",
"Scope",
"Client",
"Transport",
"HttpTransport",
"integrations",
# From sentry_sdk.api
"init",
"add_breadcrumb",
"capture_event",
"capture_exception",
"capture_message",
"configure_scope",
"continue_trace",
"flush",
"get_baggage",
"get_client",
"get_global_scope",
"get_isolation_scope",
"get_current_scope",
"get_current_span",
"get_traceparent",
"is_initialized",
"isolation_scope",
"last_event_id",
"new_scope",
"push_scope",
"set_context",
"set_extra",
"set_level",
"set_measurement",
"set_tag",
"set_tags",
"set_user",
"start_span",
"start_transaction",
"trace",
"monitor",
]
# Initialize the debug support after everything is loaded
from sentry_sdk.debug import init_debug_support
init_debug_support()
del init_debug_support
# circular imports
from sentry_sdk.hub import Hub
@@ -0,0 +1,98 @@
import sys
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import TypeVar
T = TypeVar("T")
PY37 = sys.version_info[0] == 3 and sys.version_info[1] >= 7
PY38 = sys.version_info[0] == 3 and sys.version_info[1] >= 8
PY310 = sys.version_info[0] == 3 and sys.version_info[1] >= 10
PY311 = sys.version_info[0] == 3 and sys.version_info[1] >= 11
def with_metaclass(meta, *bases):
# type: (Any, *Any) -> Any
class MetaClass(type):
def __new__(metacls, name, this_bases, d):
# type: (Any, Any, Any, Any) -> Any
return meta(name, bases, d)
return type.__new__(MetaClass, "temporary_class", (), {})
def check_uwsgi_thread_support():
# type: () -> bool
# We check two things here:
#
# 1. uWSGI doesn't run in threaded mode by default -- issue a warning if
# that's the case.
#
# 2. Additionally, if uWSGI is running in preforking mode (default), it needs
# the --py-call-uwsgi-fork-hooks option for the SDK to work properly. This
# is because any background threads spawned before the main process is
# forked are NOT CLEANED UP IN THE CHILDREN BY DEFAULT even if
# --enable-threads is on. One has to explicitly provide
# --py-call-uwsgi-fork-hooks to force uWSGI to run regular cpython
# after-fork hooks that take care of cleaning up stale thread data.
try:
from uwsgi import opt # type: ignore
except ImportError:
return True
from sentry_sdk.consts import FALSE_VALUES
def enabled(option):
# type: (str) -> bool
value = opt.get(option, False)
if isinstance(value, bool):
return value
if isinstance(value, bytes):
try:
value = value.decode()
except Exception:
pass
return value and str(value).lower() not in FALSE_VALUES
# When `threads` is passed in as a uwsgi option,
# `enable-threads` is implied on.
threads_enabled = "threads" in opt or enabled("enable-threads")
fork_hooks_on = enabled("py-call-uwsgi-fork-hooks")
lazy_mode = enabled("lazy-apps") or enabled("lazy")
if lazy_mode and not threads_enabled:
from warnings import warn
warn(
Warning(
"IMPORTANT: "
"We detected the use of uWSGI without thread support. "
"This might lead to unexpected issues. "
'Please run uWSGI with "--enable-threads" for full support.'
)
)
return False
elif not lazy_mode and (not threads_enabled or not fork_hooks_on):
from warnings import warn
warn(
Warning(
"IMPORTANT: "
"We detected the use of uWSGI in preforking mode without "
"thread support. This might lead to crashing workers. "
'Please run uWSGI with both "--enable-threads" and '
'"--py-call-uwsgi-fork-hooks" for full support.'
)
)
return False
return True
@@ -0,0 +1,84 @@
import warnings
from typing import TYPE_CHECKING
import sentry_sdk
if TYPE_CHECKING:
from typing import Any, ContextManager, Optional
import sentry_sdk.consts
class _InitGuard:
_CONTEXT_MANAGER_DEPRECATION_WARNING_MESSAGE = (
"Using the return value of sentry_sdk.init as a context manager "
"and manually calling the __enter__ and __exit__ methods on the "
"return value are deprecated. We are no longer maintaining this "
"functionality, and we will remove it in the next major release."
)
def __init__(self, client):
# type: (sentry_sdk.Client) -> None
self._client = client
def __enter__(self):
# type: () -> _InitGuard
warnings.warn(
self._CONTEXT_MANAGER_DEPRECATION_WARNING_MESSAGE,
stacklevel=2,
category=DeprecationWarning,
)
return self
def __exit__(self, exc_type, exc_value, tb):
# type: (Any, Any, Any) -> None
warnings.warn(
self._CONTEXT_MANAGER_DEPRECATION_WARNING_MESSAGE,
stacklevel=2,
category=DeprecationWarning,
)
c = self._client
if c is not None:
c.close()
def _check_python_deprecations():
# type: () -> None
# Since we're likely to deprecate Python versions in the future, I'm keeping
# this handy function around. Use this to detect the Python version used and
# to output logger.warning()s if it's deprecated.
pass
def _init(*args, **kwargs):
# type: (*Optional[str], **Any) -> ContextManager[Any]
"""Initializes the SDK and optionally integrations.
This takes the same arguments as the client constructor.
"""
client = sentry_sdk.Client(*args, **kwargs)
sentry_sdk.get_global_scope().set_client(client)
_check_python_deprecations()
rv = _InitGuard(client)
return rv
if TYPE_CHECKING:
# Make mypy, PyCharm and other static analyzers think `init` is a type to
# have nicer autocompletion for params.
#
# Use `ClientConstructor` to define the argument types of `init` and
# `ContextManager[Any]` to tell static analyzers about the return type.
class init(sentry_sdk.consts.ClientConstructor, _InitGuard): # noqa: N801
pass
else:
# Alias `init` for actual usage. Go through the lambda indirection to throw
# PyCharm off of the weakly typed signature (it would otherwise discover
# both the weakly typed signature of `_init` and our faked `init` type).
init = (lambda: _init)()
@@ -0,0 +1,47 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
_SENTINEL = object()
class LRUCache:
def __init__(self, max_size):
# type: (int) -> None
if max_size <= 0:
raise AssertionError(f"invalid max_size: {max_size}")
self.max_size = max_size
self._data = {} # type: dict[Any, Any]
self.hits = self.misses = 0
self.full = False
def set(self, key, value):
# type: (Any, Any) -> None
current = self._data.pop(key, _SENTINEL)
if current is not _SENTINEL:
self._data[key] = value
elif self.full:
self._data.pop(next(iter(self._data)))
self._data[key] = value
else:
self._data[key] = value
self.full = len(self._data) >= self.max_size
def get(self, key, default=None):
# type: (Any, Any) -> Any
try:
ret = self._data.pop(key)
except KeyError:
self.misses += 1
ret = default
else:
self.hits += 1
self._data[key] = ret
return ret
def get_all(self):
# type: () -> list[tuple[Any, Any]]
return list(self._data.items())
@@ -0,0 +1,289 @@
"""
A fork of Python 3.6's stdlib queue (found in Pythons 'cpython/Lib/queue.py')
with Lock swapped out for RLock to avoid a deadlock while garbage collecting.
https://github.com/python/cpython/blob/v3.6.12/Lib/queue.py
See also
https://codewithoutrules.com/2017/08/16/concurrency-python/
https://bugs.python.org/issue14976
https://github.com/sqlalchemy/sqlalchemy/blob/4eb747b61f0c1b1c25bdee3856d7195d10a0c227/lib/sqlalchemy/queue.py#L1
We also vendor the code to evade eventlet's broken monkeypatching, see
https://github.com/getsentry/sentry-python/pull/484
Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010,
2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software Foundation;
All Rights Reserved
PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2
--------------------------------------------
1. This LICENSE AGREEMENT is between the Python Software Foundation
("PSF"), and the Individual or Organization ("Licensee") accessing and
otherwise using this software ("Python") in source or binary form and
its associated documentation.
2. Subject to the terms and conditions of this License Agreement, PSF hereby
grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce,
analyze, test, perform and/or display publicly, prepare derivative works,
distribute, and otherwise use Python alone or in any derivative version,
provided, however, that PSF's License Agreement and PSF's notice of copyright,
i.e., "Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010,
2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020 Python Software Foundation;
All Rights Reserved" are retained in Python alone or in any derivative version
prepared by Licensee.
3. In the event Licensee prepares a derivative work that is based on
or incorporates Python or any part thereof, and wants to make
the derivative work available to others as provided herein, then
Licensee hereby agrees to include in any such work a brief summary of
the changes made to Python.
4. PSF is making Python available to Licensee on an "AS IS"
basis. PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR
IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND
DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS
FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT
INFRINGE ANY THIRD PARTY RIGHTS.
5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON
FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS
A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON,
OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF.
6. This License Agreement will automatically terminate upon a material
breach of its terms and conditions.
7. Nothing in this License Agreement shall be deemed to create any
relationship of agency, partnership, or joint venture between PSF and
Licensee. This License Agreement does not grant permission to use PSF
trademarks or trade name in a trademark sense to endorse or promote
products or services of Licensee, or any third party.
8. By copying, installing or otherwise using Python, Licensee
agrees to be bound by the terms and conditions of this License
Agreement.
"""
import threading
from collections import deque
from time import time
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
__all__ = ["EmptyError", "FullError", "Queue"]
class EmptyError(Exception):
"Exception raised by Queue.get(block=0)/get_nowait()."
pass
class FullError(Exception):
"Exception raised by Queue.put(block=0)/put_nowait()."
pass
class Queue:
"""Create a queue object with a given maximum size.
If maxsize is <= 0, the queue size is infinite.
"""
def __init__(self, maxsize=0):
self.maxsize = maxsize
self._init(maxsize)
# mutex must be held whenever the queue is mutating. All methods
# that acquire mutex must release it before returning. mutex
# is shared between the three conditions, so acquiring and
# releasing the conditions also acquires and releases mutex.
self.mutex = threading.RLock()
# Notify not_empty whenever an item is added to the queue; a
# thread waiting to get is notified then.
self.not_empty = threading.Condition(self.mutex)
# Notify not_full whenever an item is removed from the queue;
# a thread waiting to put is notified then.
self.not_full = threading.Condition(self.mutex)
# Notify all_tasks_done whenever the number of unfinished tasks
# drops to zero; thread waiting to join() is notified to resume
self.all_tasks_done = threading.Condition(self.mutex)
self.unfinished_tasks = 0
def task_done(self):
"""Indicate that a formerly enqueued task is complete.
Used by Queue consumer threads. For each get() used to fetch a task,
a subsequent call to task_done() tells the queue that the processing
on the task is complete.
If a join() is currently blocking, it will resume when all items
have been processed (meaning that a task_done() call was received
for every item that had been put() into the queue).
Raises a ValueError if called more times than there were items
placed in the queue.
"""
with self.all_tasks_done:
unfinished = self.unfinished_tasks - 1
if unfinished <= 0:
if unfinished < 0:
raise ValueError("task_done() called too many times")
self.all_tasks_done.notify_all()
self.unfinished_tasks = unfinished
def join(self):
"""Blocks until all items in the Queue have been gotten and processed.
The count of unfinished tasks goes up whenever an item is added to the
queue. The count goes down whenever a consumer thread calls task_done()
to indicate the item was retrieved and all work on it is complete.
When the count of unfinished tasks drops to zero, join() unblocks.
"""
with self.all_tasks_done:
while self.unfinished_tasks:
self.all_tasks_done.wait()
def qsize(self):
"""Return the approximate size of the queue (not reliable!)."""
with self.mutex:
return self._qsize()
def empty(self):
"""Return True if the queue is empty, False otherwise (not reliable!).
This method is likely to be removed at some point. Use qsize() == 0
as a direct substitute, but be aware that either approach risks a race
condition where a queue can grow before the result of empty() or
qsize() can be used.
To create code that needs to wait for all queued tasks to be
completed, the preferred technique is to use the join() method.
"""
with self.mutex:
return not self._qsize()
def full(self):
"""Return True if the queue is full, False otherwise (not reliable!).
This method is likely to be removed at some point. Use qsize() >= n
as a direct substitute, but be aware that either approach risks a race
condition where a queue can shrink before the result of full() or
qsize() can be used.
"""
with self.mutex:
return 0 < self.maxsize <= self._qsize()
def put(self, item, block=True, timeout=None):
"""Put an item into the queue.
If optional args 'block' is true and 'timeout' is None (the default),
block if necessary until a free slot is available. If 'timeout' is
a non-negative number, it blocks at most 'timeout' seconds and raises
the FullError exception if no free slot was available within that time.
Otherwise ('block' is false), put an item on the queue if a free slot
is immediately available, else raise the FullError exception ('timeout'
is ignored in that case).
"""
with self.not_full:
if self.maxsize > 0:
if not block:
if self._qsize() >= self.maxsize:
raise FullError()
elif timeout is None:
while self._qsize() >= self.maxsize:
self.not_full.wait()
elif timeout < 0:
raise ValueError("'timeout' must be a non-negative number")
else:
endtime = time() + timeout
while self._qsize() >= self.maxsize:
remaining = endtime - time()
if remaining <= 0.0:
raise FullError()
self.not_full.wait(remaining)
self._put(item)
self.unfinished_tasks += 1
self.not_empty.notify()
def get(self, block=True, timeout=None):
"""Remove and return an item from the queue.
If optional args 'block' is true and 'timeout' is None (the default),
block if necessary until an item is available. If 'timeout' is
a non-negative number, it blocks at most 'timeout' seconds and raises
the EmptyError exception if no item was available within that time.
Otherwise ('block' is false), return an item if one is immediately
available, else raise the EmptyError exception ('timeout' is ignored
in that case).
"""
with self.not_empty:
if not block:
if not self._qsize():
raise EmptyError()
elif timeout is None:
while not self._qsize():
self.not_empty.wait()
elif timeout < 0:
raise ValueError("'timeout' must be a non-negative number")
else:
endtime = time() + timeout
while not self._qsize():
remaining = endtime - time()
if remaining <= 0.0:
raise EmptyError()
self.not_empty.wait(remaining)
item = self._get()
self.not_full.notify()
return item
def put_nowait(self, item):
"""Put an item into the queue without blocking.
Only enqueue the item if a free slot is immediately available.
Otherwise raise the FullError exception.
"""
return self.put(item, block=False)
def get_nowait(self):
"""Remove and return an item from the queue without blocking.
Only get an item if one is immediately available. Otherwise
raise the EmptyError exception.
"""
return self.get(block=False)
# Override these methods to implement other queue organizations
# (e.g. stack or priority queue).
# These will only be called with appropriate locks held
# Initialize the queue representation
def _init(self, maxsize):
self.queue = deque() # type: Any
def _qsize(self):
return len(self.queue)
# Put a new item in the queue
def _put(self, item):
self.queue.append(item)
# Get an item from the queue
def _get(self):
return self.queue.popleft()
@@ -0,0 +1,300 @@
from typing import TYPE_CHECKING, TypeVar, Union
# Re-exported for compat, since code out there in the wild might use this variable.
MYPY = TYPE_CHECKING
SENSITIVE_DATA_SUBSTITUTE = "[Filtered]"
class AnnotatedValue:
"""
Meta information for a data field in the event payload.
This is to tell Relay that we have tampered with the fields value.
See:
https://github.com/getsentry/relay/blob/be12cd49a0f06ea932ed9b9f93a655de5d6ad6d1/relay-general/src/types/meta.rs#L407-L423
"""
__slots__ = ("value", "metadata")
def __init__(self, value, metadata):
# type: (Optional[Any], Dict[str, Any]) -> None
self.value = value
self.metadata = metadata
def __eq__(self, other):
# type: (Any) -> bool
if not isinstance(other, AnnotatedValue):
return False
return self.value == other.value and self.metadata == other.metadata
@classmethod
def removed_because_raw_data(cls):
# type: () -> AnnotatedValue
"""The value was removed because it could not be parsed. This is done for request body values that are not json nor a form."""
return AnnotatedValue(
value="",
metadata={
"rem": [ # Remark
[
"!raw", # Unparsable raw data
"x", # The fields original value was removed
]
]
},
)
@classmethod
def removed_because_over_size_limit(cls):
# type: () -> AnnotatedValue
"""The actual value was removed because the size of the field exceeded the configured maximum size (specified with the max_request_body_size sdk option)"""
return AnnotatedValue(
value="",
metadata={
"rem": [ # Remark
[
"!config", # Because of configured maximum size
"x", # The fields original value was removed
]
]
},
)
@classmethod
def substituted_because_contains_sensitive_data(cls):
# type: () -> AnnotatedValue
"""The actual value was removed because it contained sensitive information."""
return AnnotatedValue(
value=SENSITIVE_DATA_SUBSTITUTE,
metadata={
"rem": [ # Remark
[
"!config", # Because of SDK configuration (in this case the config is the hard coded removal of certain django cookies)
"s", # The fields original value was substituted
]
]
},
)
T = TypeVar("T")
Annotated = Union[AnnotatedValue, T]
if TYPE_CHECKING:
from collections.abc import Container, MutableMapping, Sequence
from datetime import datetime
from types import TracebackType
from typing import Any
from typing import Callable
from typing import Dict
from typing import Mapping
from typing import NotRequired
from typing import Optional
from typing import Tuple
from typing import Type
from typing_extensions import Literal, TypedDict
class SDKInfo(TypedDict):
name: str
version: str
packages: Sequence[Mapping[str, str]]
# "critical" is an alias of "fatal" recognized by Relay
LogLevelStr = Literal["fatal", "critical", "error", "warning", "info", "debug"]
DurationUnit = Literal[
"nanosecond",
"microsecond",
"millisecond",
"second",
"minute",
"hour",
"day",
"week",
]
InformationUnit = Literal[
"bit",
"byte",
"kilobyte",
"kibibyte",
"megabyte",
"mebibyte",
"gigabyte",
"gibibyte",
"terabyte",
"tebibyte",
"petabyte",
"pebibyte",
"exabyte",
"exbibyte",
]
FractionUnit = Literal["ratio", "percent"]
MeasurementUnit = Union[DurationUnit, InformationUnit, FractionUnit, str]
MeasurementValue = TypedDict(
"MeasurementValue",
{
"value": float,
"unit": NotRequired[Optional[MeasurementUnit]],
},
)
Event = TypedDict(
"Event",
{
"breadcrumbs": dict[
Literal["values"], list[dict[str, Any]]
], # TODO: We can expand on this type
"check_in_id": str,
"contexts": dict[str, dict[str, object]],
"dist": str,
"duration": Optional[float],
"environment": str,
"errors": list[dict[str, Any]], # TODO: We can expand on this type
"event_id": str,
"exception": dict[
Literal["values"], list[dict[str, Any]]
], # TODO: We can expand on this type
"extra": MutableMapping[str, object],
"fingerprint": list[str],
"level": LogLevelStr,
"logentry": Mapping[str, object],
"logger": str,
"measurements": dict[str, MeasurementValue],
"message": str,
"modules": dict[str, str],
"monitor_config": Mapping[str, object],
"monitor_slug": Optional[str],
"platform": Literal["python"],
"profile": object, # Should be sentry_sdk.profiler.Profile, but we can't import that here due to circular imports
"release": str,
"request": dict[str, object],
"sdk": Mapping[str, object],
"server_name": str,
"spans": Annotated[list[dict[str, object]]],
"stacktrace": dict[
str, object
], # We access this key in the code, but I am unsure whether we ever set it
"start_timestamp": datetime,
"status": Optional[str],
"tags": MutableMapping[
str, str
], # Tags must be less than 200 characters each
"threads": dict[
Literal["values"], list[dict[str, Any]]
], # TODO: We can expand on this type
"timestamp": Optional[datetime], # Must be set before sending the event
"transaction": str,
"transaction_info": Mapping[str, Any], # TODO: We can expand on this type
"type": Literal["check_in", "transaction"],
"user": dict[str, object],
"_dropped_spans": int,
"_metrics_summary": dict[str, object],
},
total=False,
)
ExcInfo = Union[
tuple[Type[BaseException], BaseException, Optional[TracebackType]],
tuple[None, None, None],
]
Hint = Dict[str, Any]
Breadcrumb = Dict[str, Any]
BreadcrumbHint = Dict[str, Any]
SamplingContext = Dict[str, Any]
EventProcessor = Callable[[Event, Hint], Optional[Event]]
ErrorProcessor = Callable[[Event, ExcInfo], Optional[Event]]
BreadcrumbProcessor = Callable[[Breadcrumb, BreadcrumbHint], Optional[Breadcrumb]]
TransactionProcessor = Callable[[Event, Hint], Optional[Event]]
TracesSampler = Callable[[SamplingContext], Union[float, int, bool]]
# https://github.com/python/mypy/issues/5710
NotImplementedType = Any
EventDataCategory = Literal[
"default",
"error",
"crash",
"transaction",
"security",
"attachment",
"session",
"internal",
"profile",
"profile_chunk",
"metric_bucket",
"monitor",
"span",
]
SessionStatus = Literal["ok", "exited", "crashed", "abnormal"]
ContinuousProfilerMode = Literal["thread", "gevent", "unknown"]
ProfilerMode = Union[ContinuousProfilerMode, Literal["sleep"]]
# Type of the metric.
MetricType = Literal["d", "s", "g", "c"]
# Value of the metric.
MetricValue = Union[int, float, str]
# Internal representation of tags as a tuple of tuples (this is done in order to allow for the same key to exist
# multiple times).
MetricTagsInternal = Tuple[Tuple[str, str], ...]
# External representation of tags as a dictionary.
MetricTagValue = Union[str, int, float, None]
MetricTags = Mapping[str, MetricTagValue]
# Value inside the generator for the metric value.
FlushedMetricValue = Union[int, float]
BucketKey = Tuple[MetricType, str, MeasurementUnit, MetricTagsInternal]
MetricMetaKey = Tuple[MetricType, str, MeasurementUnit]
MonitorConfigScheduleType = Literal["crontab", "interval"]
MonitorConfigScheduleUnit = Literal[
"year",
"month",
"week",
"day",
"hour",
"minute",
"second", # not supported in Sentry and will result in a warning
]
MonitorConfigSchedule = TypedDict(
"MonitorConfigSchedule",
{
"type": MonitorConfigScheduleType,
"value": Union[int, str],
"unit": MonitorConfigScheduleUnit,
},
total=False,
)
MonitorConfig = TypedDict(
"MonitorConfig",
{
"schedule": MonitorConfigSchedule,
"timezone": str,
"checkin_margin": int,
"max_runtime": int,
"failure_issue_threshold": int,
"recovery_threshold": int,
},
total=False,
)
HttpStatusCodeRange = Union[int, Container[int]]
@@ -0,0 +1,98 @@
"""
Copyright (c) 2007 by the Pallets team.
Some rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:
* Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE AND DOCUMENTATION IS PROVIDED BY THE COPYRIGHT HOLDERS AND
CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING,
BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
THIS SOFTWARE AND DOCUMENTATION, EVEN IF ADVISED OF THE POSSIBILITY OF
SUCH DAMAGE.
"""
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Dict
from typing import Iterator
from typing import Tuple
#
# `get_headers` comes from `werkzeug.datastructures.EnvironHeaders`
# https://github.com/pallets/werkzeug/blob/0.14.1/werkzeug/datastructures.py#L1361
#
# We need this function because Django does not give us a "pure" http header
# dict. So we might as well use it for all WSGI integrations.
#
def _get_headers(environ):
# type: (Dict[str, str]) -> Iterator[Tuple[str, str]]
"""
Returns only proper HTTP headers.
"""
for key, value in environ.items():
key = str(key)
if key.startswith("HTTP_") and key not in (
"HTTP_CONTENT_TYPE",
"HTTP_CONTENT_LENGTH",
):
yield key[5:].replace("_", "-").title(), value
elif key in ("CONTENT_TYPE", "CONTENT_LENGTH"):
yield key.replace("_", "-").title(), value
#
# `get_host` comes from `werkzeug.wsgi.get_host`
# https://github.com/pallets/werkzeug/blob/1.0.1/src/werkzeug/wsgi.py#L145
#
def get_host(environ, use_x_forwarded_for=False):
# type: (Dict[str, str], bool) -> str
"""
Return the host for the given WSGI environment.
"""
if use_x_forwarded_for and "HTTP_X_FORWARDED_HOST" in environ:
rv = environ["HTTP_X_FORWARDED_HOST"]
if environ["wsgi.url_scheme"] == "http" and rv.endswith(":80"):
rv = rv[:-3]
elif environ["wsgi.url_scheme"] == "https" and rv.endswith(":443"):
rv = rv[:-4]
elif environ.get("HTTP_HOST"):
rv = environ["HTTP_HOST"]
if environ["wsgi.url_scheme"] == "http" and rv.endswith(":80"):
rv = rv[:-3]
elif environ["wsgi.url_scheme"] == "https" and rv.endswith(":443"):
rv = rv[:-4]
elif environ.get("SERVER_NAME"):
rv = environ["SERVER_NAME"]
if (environ["wsgi.url_scheme"], environ["SERVER_PORT"]) not in (
("https", "443"),
("http", "80"),
):
rv += ":" + environ["SERVER_PORT"]
else:
# In spite of the WSGI spec, SERVER_NAME might not be present.
rv = "unknown"
return rv
@@ -0,0 +1,115 @@
import inspect
from functools import wraps
import sentry_sdk.utils
from sentry_sdk import start_span
from sentry_sdk.tracing import Span
from sentry_sdk.utils import ContextVar
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional, Callable, Any
_ai_pipeline_name = ContextVar("ai_pipeline_name", default=None)
def set_ai_pipeline_name(name):
# type: (Optional[str]) -> None
_ai_pipeline_name.set(name)
def get_ai_pipeline_name():
# type: () -> Optional[str]
return _ai_pipeline_name.get()
def ai_track(description, **span_kwargs):
# type: (str, Any) -> Callable[..., Any]
def decorator(f):
# type: (Callable[..., Any]) -> Callable[..., Any]
def sync_wrapped(*args, **kwargs):
# type: (Any, Any) -> Any
curr_pipeline = _ai_pipeline_name.get()
op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
with start_span(name=description, op=op, **span_kwargs) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_tag(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_data(k, v)
if curr_pipeline:
span.set_data("ai.pipeline.name", curr_pipeline)
return f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = f(*args, **kwargs)
except Exception as e:
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "ai_monitoring", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
raise e from None
finally:
_ai_pipeline_name.set(None)
return res
async def async_wrapped(*args, **kwargs):
# type: (Any, Any) -> Any
curr_pipeline = _ai_pipeline_name.get()
op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
with start_span(name=description, op=op, **span_kwargs) as span:
for k, v in kwargs.pop("sentry_tags", {}).items():
span.set_tag(k, v)
for k, v in kwargs.pop("sentry_data", {}).items():
span.set_data(k, v)
if curr_pipeline:
span.set_data("ai.pipeline.name", curr_pipeline)
return await f(*args, **kwargs)
else:
_ai_pipeline_name.set(description)
try:
res = await f(*args, **kwargs)
except Exception as e:
event, hint = sentry_sdk.utils.event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "ai_monitoring", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
raise e from None
finally:
_ai_pipeline_name.set(None)
return res
if inspect.iscoroutinefunction(f):
return wraps(f)(async_wrapped)
else:
return wraps(f)(sync_wrapped)
return decorator
def record_token_usage(
span, prompt_tokens=None, completion_tokens=None, total_tokens=None
):
# type: (Span, Optional[int], Optional[int], Optional[int]) -> None
ai_pipeline_name = get_ai_pipeline_name()
if ai_pipeline_name:
span.set_data("ai.pipeline.name", ai_pipeline_name)
if prompt_tokens is not None:
span.set_measurement("ai_prompt_tokens_used", value=prompt_tokens)
if completion_tokens is not None:
span.set_measurement("ai_completion_tokens_used", value=completion_tokens)
if (
total_tokens is None
and prompt_tokens is not None
and completion_tokens is not None
):
total_tokens = prompt_tokens + completion_tokens
if total_tokens is not None:
span.set_measurement("ai_total_tokens_used", total_tokens)
@@ -0,0 +1,32 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from sentry_sdk.tracing import Span
from sentry_sdk.utils import logger
def _normalize_data(data):
# type: (Any) -> Any
# convert pydantic data (e.g. OpenAI v1+) to json compatible format
if hasattr(data, "model_dump"):
try:
return data.model_dump()
except Exception as e:
logger.warning("Could not convert pydantic data to JSON: %s", e)
return data
if isinstance(data, list):
if len(data) == 1:
return _normalize_data(data[0]) # remove empty dimensions
return list(_normalize_data(x) for x in data)
if isinstance(data, dict):
return {k: _normalize_data(v) for (k, v) in data.items()}
return data
def set_data_normalized(span, key, value):
# type: (Span, str, Any) -> None
normalized = _normalize_data(value)
span.set_data(key, normalized)
@@ -0,0 +1,433 @@
import inspect
import warnings
from contextlib import contextmanager
from sentry_sdk import tracing_utils, Client
from sentry_sdk._init_implementation import init
from sentry_sdk.consts import INSTRUMENTER
from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope
from sentry_sdk.tracing import NoOpSpan, Transaction, trace
from sentry_sdk.crons import monitor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Mapping
from typing import Any
from typing import Dict
from typing import Generator
from typing import Optional
from typing import overload
from typing import Callable
from typing import TypeVar
from typing import ContextManager
from typing import Union
from typing_extensions import Unpack
from sentry_sdk.client import BaseClient
from sentry_sdk._types import (
Event,
Hint,
Breadcrumb,
BreadcrumbHint,
ExcInfo,
MeasurementUnit,
LogLevelStr,
SamplingContext,
)
from sentry_sdk.tracing import Span, TransactionKwargs
T = TypeVar("T")
F = TypeVar("F", bound=Callable[..., Any])
else:
def overload(x):
# type: (T) -> T
return x
# When changing this, update __all__ in __init__.py too
__all__ = [
"init",
"add_breadcrumb",
"capture_event",
"capture_exception",
"capture_message",
"configure_scope",
"continue_trace",
"flush",
"get_baggage",
"get_client",
"get_global_scope",
"get_isolation_scope",
"get_current_scope",
"get_current_span",
"get_traceparent",
"is_initialized",
"isolation_scope",
"last_event_id",
"new_scope",
"push_scope",
"set_context",
"set_extra",
"set_level",
"set_measurement",
"set_tag",
"set_tags",
"set_user",
"start_span",
"start_transaction",
"trace",
"monitor",
]
def scopemethod(f):
# type: (F) -> F
f.__doc__ = "%s\n\n%s" % (
"Alias for :py:meth:`sentry_sdk.Scope.%s`" % f.__name__,
inspect.getdoc(getattr(Scope, f.__name__)),
)
return f
def clientmethod(f):
# type: (F) -> F
f.__doc__ = "%s\n\n%s" % (
"Alias for :py:meth:`sentry_sdk.Client.%s`" % f.__name__,
inspect.getdoc(getattr(Client, f.__name__)),
)
return f
@scopemethod
def get_client():
# type: () -> BaseClient
return Scope.get_client()
def is_initialized():
# type: () -> bool
"""
.. versionadded:: 2.0.0
Returns whether Sentry has been initialized or not.
If a client is available and the client is active
(meaning it is configured to send data) then
Sentry is initialized.
"""
return get_client().is_active()
@scopemethod
def get_global_scope():
# type: () -> Scope
return Scope.get_global_scope()
@scopemethod
def get_isolation_scope():
# type: () -> Scope
return Scope.get_isolation_scope()
@scopemethod
def get_current_scope():
# type: () -> Scope
return Scope.get_current_scope()
@scopemethod
def last_event_id():
# type: () -> Optional[str]
"""
See :py:meth:`sentry_sdk.Scope.last_event_id` documentation regarding
this method's limitations.
"""
return Scope.last_event_id()
@scopemethod
def capture_event(
event, # type: Event
hint=None, # type: Optional[Hint]
scope=None, # type: Optional[Any]
**scope_kwargs, # type: Any
):
# type: (...) -> Optional[str]
return get_current_scope().capture_event(event, hint, scope=scope, **scope_kwargs)
@scopemethod
def capture_message(
message, # type: str
level=None, # type: Optional[LogLevelStr]
scope=None, # type: Optional[Any]
**scope_kwargs, # type: Any
):
# type: (...) -> Optional[str]
return get_current_scope().capture_message(
message, level, scope=scope, **scope_kwargs
)
@scopemethod
def capture_exception(
error=None, # type: Optional[Union[BaseException, ExcInfo]]
scope=None, # type: Optional[Any]
**scope_kwargs, # type: Any
):
# type: (...) -> Optional[str]
return get_current_scope().capture_exception(error, scope=scope, **scope_kwargs)
@scopemethod
def add_breadcrumb(
crumb=None, # type: Optional[Breadcrumb]
hint=None, # type: Optional[BreadcrumbHint]
**kwargs, # type: Any
):
# type: (...) -> None
return get_isolation_scope().add_breadcrumb(crumb, hint, **kwargs)
@overload
def configure_scope():
# type: () -> ContextManager[Scope]
pass
@overload
def configure_scope( # noqa: F811
callback, # type: Callable[[Scope], None]
):
# type: (...) -> None
pass
def configure_scope( # noqa: F811
callback=None, # type: Optional[Callable[[Scope], None]]
):
# type: (...) -> Optional[ContextManager[Scope]]
"""
Reconfigures the scope.
:param callback: If provided, call the callback with the current scope.
:returns: If no callback is provided, returns a context manager that returns the scope.
"""
warnings.warn(
"sentry_sdk.configure_scope is deprecated and will be removed in the next major version. "
"Please consult our migration guide to learn how to migrate to the new API: "
"https://docs.sentry.io/platforms/python/migration/1.x-to-2.x#scope-configuring",
DeprecationWarning,
stacklevel=2,
)
scope = get_isolation_scope()
scope.generate_propagation_context()
if callback is not None:
# TODO: used to return None when client is None. Check if this changes behavior.
callback(scope)
return None
@contextmanager
def inner():
# type: () -> Generator[Scope, None, None]
yield scope
return inner()
@overload
def push_scope():
# type: () -> ContextManager[Scope]
pass
@overload
def push_scope( # noqa: F811
callback, # type: Callable[[Scope], None]
):
# type: (...) -> None
pass
def push_scope( # noqa: F811
callback=None, # type: Optional[Callable[[Scope], None]]
):
# type: (...) -> Optional[ContextManager[Scope]]
"""
Pushes a new layer on the scope stack.
:param callback: If provided, this method pushes a scope, calls
`callback`, and pops the scope again.
:returns: If no `callback` is provided, a context manager that should
be used to pop the scope again.
"""
warnings.warn(
"sentry_sdk.push_scope is deprecated and will be removed in the next major version. "
"Please consult our migration guide to learn how to migrate to the new API: "
"https://docs.sentry.io/platforms/python/migration/1.x-to-2.x#scope-pushing",
DeprecationWarning,
stacklevel=2,
)
if callback is not None:
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
with push_scope() as scope:
callback(scope)
return None
return _ScopeManager()
@scopemethod
def set_tag(key, value):
# type: (str, Any) -> None
return get_isolation_scope().set_tag(key, value)
@scopemethod
def set_tags(tags):
# type: (Mapping[str, object]) -> None
return get_isolation_scope().set_tags(tags)
@scopemethod
def set_context(key, value):
# type: (str, Dict[str, Any]) -> None
return get_isolation_scope().set_context(key, value)
@scopemethod
def set_extra(key, value):
# type: (str, Any) -> None
return get_isolation_scope().set_extra(key, value)
@scopemethod
def set_user(value):
# type: (Optional[Dict[str, Any]]) -> None
return get_isolation_scope().set_user(value)
@scopemethod
def set_level(value):
# type: (LogLevelStr) -> None
return get_isolation_scope().set_level(value)
@clientmethod
def flush(
timeout=None, # type: Optional[float]
callback=None, # type: Optional[Callable[[int, float], None]]
):
# type: (...) -> None
return get_client().flush(timeout=timeout, callback=callback)
@scopemethod
def start_span(
**kwargs, # type: Any
):
# type: (...) -> Span
return get_current_scope().start_span(**kwargs)
@scopemethod
def start_transaction(
transaction=None, # type: Optional[Transaction]
instrumenter=INSTRUMENTER.SENTRY, # type: str
custom_sampling_context=None, # type: Optional[SamplingContext]
**kwargs, # type: Unpack[TransactionKwargs]
):
# type: (...) -> Union[Transaction, NoOpSpan]
"""
Start and return a transaction on the current scope.
Start an existing transaction if given, otherwise create and start a new
transaction with kwargs.
This is the entry point to manual tracing instrumentation.
A tree structure can be built by adding child spans to the transaction,
and child spans to other spans. To start a new child span within the
transaction or any span, call the respective `.start_child()` method.
Every child span must be finished before the transaction is finished,
otherwise the unfinished spans are discarded.
When used as context managers, spans and transactions are automatically
finished at the end of the `with` block. If not using context managers,
call the `.finish()` method.
When the transaction is finished, it will be sent to Sentry with all its
finished child spans.
:param transaction: The transaction to start. If omitted, we create and
start a new transaction.
:param instrumenter: This parameter is meant for internal use only. It
will be removed in the next major version.
:param custom_sampling_context: The transaction's custom sampling context.
:param kwargs: Optional keyword arguments to be passed to the Transaction
constructor. See :py:class:`sentry_sdk.tracing.Transaction` for
available arguments.
"""
return get_current_scope().start_transaction(
transaction, instrumenter, custom_sampling_context, **kwargs
)
def set_measurement(name, value, unit=""):
# type: (str, float, MeasurementUnit) -> None
transaction = get_current_scope().transaction
if transaction is not None:
transaction.set_measurement(name, value, unit)
def get_current_span(scope=None):
# type: (Optional[Scope]) -> Optional[Span]
"""
Returns the currently active span if there is one running, otherwise `None`
"""
return tracing_utils.get_current_span(scope)
def get_traceparent():
# type: () -> Optional[str]
"""
Returns the traceparent either from the active span or from the scope.
"""
return get_current_scope().get_traceparent()
def get_baggage():
# type: () -> Optional[str]
"""
Returns Baggage either from the active span or from the scope.
"""
baggage = get_current_scope().get_baggage()
if baggage is not None:
return baggage.serialize()
return None
def continue_trace(
environ_or_headers, op=None, name=None, source=None, origin="manual"
):
# type: (Dict[str, Any], Optional[str], Optional[str], Optional[str], str) -> Transaction
"""
Sets the propagation context from environment or headers and returns a transaction.
"""
return get_isolation_scope().continue_trace(
environ_or_headers, op, name, source, origin
)
@@ -0,0 +1,75 @@
import os
import mimetypes
from sentry_sdk.envelope import Item, PayloadRef
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional, Union, Callable
class Attachment:
"""Additional files/data to send along with an event.
This class stores attachments that can be sent along with an event. Attachments are files or other data, e.g.
config or log files, that are relevant to an event. Attachments are set on the ``Scope``, and are sent along with
all non-transaction events (or all events including transactions if ``add_to_transactions`` is ``True``) that are
captured within the ``Scope``.
To add an attachment to a ``Scope``, use :py:meth:`sentry_sdk.Scope.add_attachment`. The parameters for
``add_attachment`` are the same as the parameters for this class's constructor.
:param bytes: Raw bytes of the attachment, or a function that returns the raw bytes. Must be provided unless
``path`` is provided.
:param filename: The filename of the attachment. Must be provided unless ``path`` is provided.
:param path: Path to a file to attach. Must be provided unless ``bytes`` is provided.
:param content_type: The content type of the attachment. If not provided, it will be guessed from the ``filename``
parameter, if available, or the ``path`` parameter if ``filename`` is ``None``.
:param add_to_transactions: Whether to add this attachment to transactions. Defaults to ``False``.
"""
def __init__(
self,
bytes=None, # type: Union[None, bytes, Callable[[], bytes]]
filename=None, # type: Optional[str]
path=None, # type: Optional[str]
content_type=None, # type: Optional[str]
add_to_transactions=False, # type: bool
):
# type: (...) -> None
if bytes is None and path is None:
raise TypeError("path or raw bytes required for attachment")
if filename is None and path is not None:
filename = os.path.basename(path)
if filename is None:
raise TypeError("filename is required for attachment")
if content_type is None:
content_type = mimetypes.guess_type(filename)[0]
self.bytes = bytes
self.filename = filename
self.path = path
self.content_type = content_type
self.add_to_transactions = add_to_transactions
def to_envelope_item(self):
# type: () -> Item
"""Returns an envelope item for this attachment."""
payload = None # type: Union[None, PayloadRef, bytes]
if self.bytes is not None:
if callable(self.bytes):
payload = self.bytes()
else:
payload = self.bytes
else:
payload = PayloadRef(path=self.path)
return Item(
payload=payload,
type="attachment",
content_type=self.content_type,
filename=self.filename,
)
def __repr__(self):
# type: () -> str
return "<Attachment %r>" % (self.filename,)
@@ -0,0 +1,959 @@
import os
import uuid
import random
import socket
from collections.abc import Mapping
from datetime import datetime, timezone
from importlib import import_module
from typing import TYPE_CHECKING, List, Dict, cast, overload
import warnings
from sentry_sdk._compat import PY37, check_uwsgi_thread_support
from sentry_sdk.utils import (
AnnotatedValue,
ContextVar,
capture_internal_exceptions,
current_stacktrace,
env_to_bool,
format_timestamp,
get_sdk_name,
get_type_name,
get_default_release,
handle_in_app,
is_gevent,
logger,
)
from sentry_sdk.serializer import serialize
from sentry_sdk.tracing import trace
from sentry_sdk.transport import BaseHttpTransport, make_transport
from sentry_sdk.consts import (
DEFAULT_MAX_VALUE_LENGTH,
DEFAULT_OPTIONS,
INSTRUMENTER,
VERSION,
ClientConstructor,
)
from sentry_sdk.integrations import _DEFAULT_INTEGRATIONS, setup_integrations
from sentry_sdk.sessions import SessionFlusher
from sentry_sdk.envelope import Envelope
from sentry_sdk.profiler.continuous_profiler import setup_continuous_profiler
from sentry_sdk.profiler.transaction_profiler import (
has_profiling_enabled,
Profile,
setup_profiler,
)
from sentry_sdk.scrubber import EventScrubber
from sentry_sdk.monitor import Monitor
from sentry_sdk.spotlight import setup_spotlight
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import Optional
from typing import Sequence
from typing import Type
from typing import Union
from typing import TypeVar
from sentry_sdk._types import Event, Hint, SDKInfo
from sentry_sdk.integrations import Integration
from sentry_sdk.metrics import MetricsAggregator
from sentry_sdk.scope import Scope
from sentry_sdk.session import Session
from sentry_sdk.spotlight import SpotlightClient
from sentry_sdk.transport import Transport
I = TypeVar("I", bound=Integration) # noqa: E741
_client_init_debug = ContextVar("client_init_debug")
SDK_INFO = {
"name": "sentry.python", # SDK name will be overridden after integrations have been loaded with sentry_sdk.integrations.setup_integrations()
"version": VERSION,
"packages": [{"name": "pypi:sentry-sdk", "version": VERSION}],
} # type: SDKInfo
def _get_options(*args, **kwargs):
# type: (*Optional[str], **Any) -> Dict[str, Any]
if args and (isinstance(args[0], (bytes, str)) or args[0] is None):
dsn = args[0] # type: Optional[str]
args = args[1:]
else:
dsn = None
if len(args) > 1:
raise TypeError("Only single positional argument is expected")
rv = dict(DEFAULT_OPTIONS)
options = dict(*args, **kwargs)
if dsn is not None and options.get("dsn") is None:
options["dsn"] = dsn
for key, value in options.items():
if key not in rv:
raise TypeError("Unknown option %r" % (key,))
rv[key] = value
if rv["dsn"] is None:
rv["dsn"] = os.environ.get("SENTRY_DSN")
if rv["release"] is None:
rv["release"] = get_default_release()
if rv["environment"] is None:
rv["environment"] = os.environ.get("SENTRY_ENVIRONMENT") or "production"
if rv["debug"] is None:
rv["debug"] = env_to_bool(os.environ.get("SENTRY_DEBUG", "False"), strict=True)
if rv["server_name"] is None and hasattr(socket, "gethostname"):
rv["server_name"] = socket.gethostname()
if rv["instrumenter"] is None:
rv["instrumenter"] = INSTRUMENTER.SENTRY
if rv["project_root"] is None:
try:
project_root = os.getcwd()
except Exception:
project_root = None
rv["project_root"] = project_root
if rv["enable_tracing"] is True and rv["traces_sample_rate"] is None:
rv["traces_sample_rate"] = 1.0
if rv["event_scrubber"] is None:
rv["event_scrubber"] = EventScrubber(
send_default_pii=(
False if rv["send_default_pii"] is None else rv["send_default_pii"]
)
)
if rv["socket_options"] and not isinstance(rv["socket_options"], list):
logger.warning(
"Ignoring socket_options because of unexpected format. See urllib3.HTTPConnection.socket_options for the expected format."
)
rv["socket_options"] = None
if rv["enable_tracing"] is not None:
warnings.warn(
"The `enable_tracing` parameter is deprecated. Please use `traces_sample_rate` instead.",
DeprecationWarning,
stacklevel=2,
)
return rv
try:
# Python 3.6+
module_not_found_error = ModuleNotFoundError
except Exception:
# Older Python versions
module_not_found_error = ImportError # type: ignore
class BaseClient:
"""
.. versionadded:: 2.0.0
The basic definition of a client that is used for sending data to Sentry.
"""
spotlight = None # type: Optional[SpotlightClient]
def __init__(self, options=None):
# type: (Optional[Dict[str, Any]]) -> None
self.options = (
options if options is not None else DEFAULT_OPTIONS
) # type: Dict[str, Any]
self.transport = None # type: Optional[Transport]
self.monitor = None # type: Optional[Monitor]
self.metrics_aggregator = None # type: Optional[MetricsAggregator]
def __getstate__(self, *args, **kwargs):
# type: (*Any, **Any) -> Any
return {"options": {}}
def __setstate__(self, *args, **kwargs):
# type: (*Any, **Any) -> None
pass
@property
def dsn(self):
# type: () -> Optional[str]
return None
def should_send_default_pii(self):
# type: () -> bool
return False
def is_active(self):
# type: () -> bool
"""
.. versionadded:: 2.0.0
Returns whether the client is active (able to send data to Sentry)
"""
return False
def capture_event(self, *args, **kwargs):
# type: (*Any, **Any) -> Optional[str]
return None
def capture_session(self, *args, **kwargs):
# type: (*Any, **Any) -> None
return None
if TYPE_CHECKING:
@overload
def get_integration(self, name_or_class):
# type: (str) -> Optional[Integration]
...
@overload
def get_integration(self, name_or_class):
# type: (type[I]) -> Optional[I]
...
def get_integration(self, name_or_class):
# type: (Union[str, type[Integration]]) -> Optional[Integration]
return None
def close(self, *args, **kwargs):
# type: (*Any, **Any) -> None
return None
def flush(self, *args, **kwargs):
# type: (*Any, **Any) -> None
return None
def __enter__(self):
# type: () -> BaseClient
return self
def __exit__(self, exc_type, exc_value, tb):
# type: (Any, Any, Any) -> None
return None
class NonRecordingClient(BaseClient):
"""
.. versionadded:: 2.0.0
A client that does not send any events to Sentry. This is used as a fallback when the Sentry SDK is not yet initialized.
"""
pass
class _Client(BaseClient):
"""
The client is internally responsible for capturing the events and
forwarding them to sentry through the configured transport. It takes
the client options as keyword arguments and optionally the DSN as first
argument.
Alias of :py:class:`sentry_sdk.Client`. (Was created for better intelisense support)
"""
def __init__(self, *args, **kwargs):
# type: (*Any, **Any) -> None
super(_Client, self).__init__(options=get_options(*args, **kwargs))
self._init_impl()
def __getstate__(self):
# type: () -> Any
return {"options": self.options}
def __setstate__(self, state):
# type: (Any) -> None
self.options = state["options"]
self._init_impl()
def _setup_instrumentation(self, functions_to_trace):
# type: (Sequence[Dict[str, str]]) -> None
"""
Instruments the functions given in the list `functions_to_trace` with the `@sentry_sdk.tracing.trace` decorator.
"""
for function in functions_to_trace:
class_name = None
function_qualname = function["qualified_name"]
module_name, function_name = function_qualname.rsplit(".", 1)
try:
# Try to import module and function
# ex: "mymodule.submodule.funcname"
module_obj = import_module(module_name)
function_obj = getattr(module_obj, function_name)
setattr(module_obj, function_name, trace(function_obj))
logger.debug("Enabled tracing for %s", function_qualname)
except module_not_found_error:
try:
# Try to import a class
# ex: "mymodule.submodule.MyClassName.member_function"
module_name, class_name = module_name.rsplit(".", 1)
module_obj = import_module(module_name)
class_obj = getattr(module_obj, class_name)
function_obj = getattr(class_obj, function_name)
function_type = type(class_obj.__dict__[function_name])
traced_function = trace(function_obj)
if function_type in (staticmethod, classmethod):
traced_function = staticmethod(traced_function)
setattr(class_obj, function_name, traced_function)
setattr(module_obj, class_name, class_obj)
logger.debug("Enabled tracing for %s", function_qualname)
except Exception as e:
logger.warning(
"Can not enable tracing for '%s'. (%s) Please check your `functions_to_trace` parameter.",
function_qualname,
e,
)
except Exception as e:
logger.warning(
"Can not enable tracing for '%s'. (%s) Please check your `functions_to_trace` parameter.",
function_qualname,
e,
)
def _init_impl(self):
# type: () -> None
old_debug = _client_init_debug.get(False)
def _capture_envelope(envelope):
# type: (Envelope) -> None
if self.transport is not None:
self.transport.capture_envelope(envelope)
try:
_client_init_debug.set(self.options["debug"])
self.transport = make_transport(self.options)
self.monitor = None
if self.transport:
if self.options["enable_backpressure_handling"]:
self.monitor = Monitor(self.transport)
self.session_flusher = SessionFlusher(capture_func=_capture_envelope)
self.metrics_aggregator = None # type: Optional[MetricsAggregator]
experiments = self.options.get("_experiments", {})
if experiments.get("enable_metrics", True):
# Context vars are not working correctly on Python <=3.6
# with gevent.
metrics_supported = not is_gevent() or PY37
if metrics_supported:
from sentry_sdk.metrics import MetricsAggregator
self.metrics_aggregator = MetricsAggregator(
capture_func=_capture_envelope,
enable_code_locations=bool(
experiments.get("metric_code_locations", True)
),
)
else:
logger.info(
"Metrics not supported on Python 3.6 and lower with gevent."
)
max_request_body_size = ("always", "never", "small", "medium")
if self.options["max_request_body_size"] not in max_request_body_size:
raise ValueError(
"Invalid value for max_request_body_size. Must be one of {}".format(
max_request_body_size
)
)
if self.options["_experiments"].get("otel_powered_performance", False):
logger.debug(
"[OTel] Enabling experimental OTel-powered performance monitoring."
)
self.options["instrumenter"] = INSTRUMENTER.OTEL
if (
"sentry_sdk.integrations.opentelemetry.integration.OpenTelemetryIntegration"
not in _DEFAULT_INTEGRATIONS
):
_DEFAULT_INTEGRATIONS.append(
"sentry_sdk.integrations.opentelemetry.integration.OpenTelemetryIntegration",
)
self.integrations = setup_integrations(
self.options["integrations"],
with_defaults=self.options["default_integrations"],
with_auto_enabling_integrations=self.options[
"auto_enabling_integrations"
],
disabled_integrations=self.options["disabled_integrations"],
)
spotlight_config = self.options.get("spotlight")
if spotlight_config is None and "SENTRY_SPOTLIGHT" in os.environ:
spotlight_env_value = os.environ["SENTRY_SPOTLIGHT"]
spotlight_config = env_to_bool(spotlight_env_value, strict=True)
self.options["spotlight"] = (
spotlight_config
if spotlight_config is not None
else spotlight_env_value
)
if self.options.get("spotlight"):
self.spotlight = setup_spotlight(self.options)
sdk_name = get_sdk_name(list(self.integrations.keys()))
SDK_INFO["name"] = sdk_name
logger.debug("Setting SDK name to '%s'", sdk_name)
if has_profiling_enabled(self.options):
try:
setup_profiler(self.options)
except Exception as e:
logger.debug("Can not set up profiler. (%s)", e)
else:
try:
setup_continuous_profiler(
self.options,
sdk_info=SDK_INFO,
capture_func=_capture_envelope,
)
except Exception as e:
logger.debug("Can not set up continuous profiler. (%s)", e)
finally:
_client_init_debug.set(old_debug)
self._setup_instrumentation(self.options.get("functions_to_trace", []))
if (
self.monitor
or self.metrics_aggregator
or has_profiling_enabled(self.options)
or isinstance(self.transport, BaseHttpTransport)
):
# If we have anything on that could spawn a background thread, we
# need to check if it's safe to use them.
check_uwsgi_thread_support()
def is_active(self):
# type: () -> bool
"""
.. versionadded:: 2.0.0
Returns whether the client is active (able to send data to Sentry)
"""
return True
def should_send_default_pii(self):
# type: () -> bool
"""
.. versionadded:: 2.0.0
Returns whether the client should send default PII (Personally Identifiable Information) data to Sentry.
"""
result = self.options.get("send_default_pii")
if result is None:
result = not self.options["dsn"] and self.spotlight is not None
return result
@property
def dsn(self):
# type: () -> Optional[str]
"""Returns the configured DSN as string."""
return self.options["dsn"]
def _prepare_event(
self,
event, # type: Event
hint, # type: Hint
scope, # type: Optional[Scope]
):
# type: (...) -> Optional[Event]
previous_total_spans = None # type: Optional[int]
if event.get("timestamp") is None:
event["timestamp"] = datetime.now(timezone.utc)
if scope is not None:
is_transaction = event.get("type") == "transaction"
spans_before = len(cast(List[Dict[str, object]], event.get("spans", [])))
event_ = scope.apply_to_event(event, hint, self.options)
# one of the event/error processors returned None
if event_ is None:
if self.transport:
self.transport.record_lost_event(
"event_processor",
data_category=("transaction" if is_transaction else "error"),
)
if is_transaction:
self.transport.record_lost_event(
"event_processor",
data_category="span",
quantity=spans_before + 1, # +1 for the transaction itself
)
return None
event = event_
spans_delta = spans_before - len(
cast(List[Dict[str, object]], event.get("spans", []))
)
if is_transaction and spans_delta > 0 and self.transport is not None:
self.transport.record_lost_event(
"event_processor", data_category="span", quantity=spans_delta
)
dropped_spans = event.pop("_dropped_spans", 0) + spans_delta # type: int
if dropped_spans > 0:
previous_total_spans = spans_before + dropped_spans
if (
self.options["attach_stacktrace"]
and "exception" not in event
and "stacktrace" not in event
and "threads" not in event
):
with capture_internal_exceptions():
event["threads"] = {
"values": [
{
"stacktrace": current_stacktrace(
include_local_variables=self.options.get(
"include_local_variables", True
),
max_value_length=self.options.get(
"max_value_length", DEFAULT_MAX_VALUE_LENGTH
),
),
"crashed": False,
"current": True,
}
]
}
for key in "release", "environment", "server_name", "dist":
if event.get(key) is None and self.options[key] is not None:
event[key] = str(self.options[key]).strip()
if event.get("sdk") is None:
sdk_info = dict(SDK_INFO)
sdk_info["integrations"] = sorted(self.integrations.keys())
event["sdk"] = sdk_info
if event.get("platform") is None:
event["platform"] = "python"
event = handle_in_app(
event,
self.options["in_app_exclude"],
self.options["in_app_include"],
self.options["project_root"],
)
if event is not None:
event_scrubber = self.options["event_scrubber"]
if event_scrubber:
event_scrubber.scrub_event(event)
if previous_total_spans is not None:
event["spans"] = AnnotatedValue(
event.get("spans", []), {"len": previous_total_spans}
)
# Postprocess the event here so that annotated types do
# generally not surface in before_send
if event is not None:
event = cast(
"Event",
serialize(
cast("Dict[str, Any]", event),
max_request_body_size=self.options.get("max_request_body_size"),
max_value_length=self.options.get("max_value_length"),
custom_repr=self.options.get("custom_repr"),
),
)
before_send = self.options["before_send"]
if (
before_send is not None
and event is not None
and event.get("type") != "transaction"
):
new_event = None
with capture_internal_exceptions():
new_event = before_send(event, hint or {})
if new_event is None:
logger.info("before send dropped event")
if self.transport:
self.transport.record_lost_event(
"before_send", data_category="error"
)
event = new_event
before_send_transaction = self.options["before_send_transaction"]
if (
before_send_transaction is not None
and event is not None
and event.get("type") == "transaction"
):
new_event = None
spans_before = len(cast(List[Dict[str, object]], event.get("spans", [])))
with capture_internal_exceptions():
new_event = before_send_transaction(event, hint or {})
if new_event is None:
logger.info("before send transaction dropped event")
if self.transport:
self.transport.record_lost_event(
reason="before_send", data_category="transaction"
)
self.transport.record_lost_event(
reason="before_send",
data_category="span",
quantity=spans_before + 1, # +1 for the transaction itself
)
else:
spans_delta = spans_before - len(new_event.get("spans", []))
if spans_delta > 0 and self.transport is not None:
self.transport.record_lost_event(
reason="before_send", data_category="span", quantity=spans_delta
)
event = new_event
return event
def _is_ignored_error(self, event, hint):
# type: (Event, Hint) -> bool
exc_info = hint.get("exc_info")
if exc_info is None:
return False
error = exc_info[0]
error_type_name = get_type_name(exc_info[0])
error_full_name = "%s.%s" % (exc_info[0].__module__, error_type_name)
for ignored_error in self.options["ignore_errors"]:
# String types are matched against the type name in the
# exception only
if isinstance(ignored_error, str):
if ignored_error == error_full_name or ignored_error == error_type_name:
return True
else:
if issubclass(error, ignored_error):
return True
return False
def _should_capture(
self,
event, # type: Event
hint, # type: Hint
scope=None, # type: Optional[Scope]
):
# type: (...) -> bool
# Transactions are sampled independent of error events.
is_transaction = event.get("type") == "transaction"
if is_transaction:
return True
ignoring_prevents_recursion = scope is not None and not scope._should_capture
if ignoring_prevents_recursion:
return False
ignored_by_config_option = self._is_ignored_error(event, hint)
if ignored_by_config_option:
return False
return True
def _should_sample_error(
self,
event, # type: Event
hint, # type: Hint
):
# type: (...) -> bool
error_sampler = self.options.get("error_sampler", None)
if callable(error_sampler):
with capture_internal_exceptions():
sample_rate = error_sampler(event, hint)
else:
sample_rate = self.options["sample_rate"]
try:
not_in_sample_rate = sample_rate < 1.0 and random.random() >= sample_rate
except NameError:
logger.warning(
"The provided error_sampler raised an error. Defaulting to sampling the event."
)
# If the error_sampler raised an error, we should sample the event, since the default behavior
# (when no sample_rate or error_sampler is provided) is to sample all events.
not_in_sample_rate = False
except TypeError:
parameter, verb = (
("error_sampler", "returned")
if callable(error_sampler)
else ("sample_rate", "contains")
)
logger.warning(
"The provided %s %s an invalid value of %s. The value should be a float or a bool. Defaulting to sampling the event."
% (parameter, verb, repr(sample_rate))
)
# If the sample_rate has an invalid value, we should sample the event, since the default behavior
# (when no sample_rate or error_sampler is provided) is to sample all events.
not_in_sample_rate = False
if not_in_sample_rate:
# because we will not sample this event, record a "lost event".
if self.transport:
self.transport.record_lost_event("sample_rate", data_category="error")
return False
return True
def _update_session_from_event(
self,
session, # type: Session
event, # type: Event
):
# type: (...) -> None
crashed = False
errored = False
user_agent = None
exceptions = (event.get("exception") or {}).get("values")
if exceptions:
errored = True
for error in exceptions:
mechanism = error.get("mechanism")
if isinstance(mechanism, Mapping) and mechanism.get("handled") is False:
crashed = True
break
user = event.get("user")
if session.user_agent is None:
headers = (event.get("request") or {}).get("headers")
headers_dict = headers if isinstance(headers, dict) else {}
for k, v in headers_dict.items():
if k.lower() == "user-agent":
user_agent = v
break
session.update(
status="crashed" if crashed else None,
user=user,
user_agent=user_agent,
errors=session.errors + (errored or crashed),
)
def capture_event(
self,
event, # type: Event
hint=None, # type: Optional[Hint]
scope=None, # type: Optional[Scope]
):
# type: (...) -> Optional[str]
"""Captures an event.
:param event: A ready-made event that can be directly sent to Sentry.
:param hint: Contains metadata about the event that can be read from `before_send`, such as the original exception object or a HTTP request object.
:param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events.
:returns: An event ID. May be `None` if there is no DSN set or of if the SDK decided to discard the event for other reasons. In such situations setting `debug=True` on `init()` may help.
"""
hint = dict(hint or ()) # type: Hint
if not self._should_capture(event, hint, scope):
return None
profile = event.pop("profile", None)
event_id = event.get("event_id")
if event_id is None:
event["event_id"] = event_id = uuid.uuid4().hex
event_opt = self._prepare_event(event, hint, scope)
if event_opt is None:
return None
# 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
if session:
self._update_session_from_event(session, event)
is_transaction = event_opt.get("type") == "transaction"
is_checkin = event_opt.get("type") == "check_in"
if (
not is_transaction
and not is_checkin
and not self._should_sample_error(event, hint)
):
return None
attachments = hint.get("attachments")
trace_context = event_opt.get("contexts", {}).get("trace") or {}
dynamic_sampling_context = trace_context.pop("dynamic_sampling_context", {})
headers = {
"event_id": event_opt["event_id"],
"sent_at": format_timestamp(datetime.now(timezone.utc)),
} # type: dict[str, object]
if dynamic_sampling_context:
headers["trace"] = dynamic_sampling_context
envelope = Envelope(headers=headers)
if is_transaction:
if isinstance(profile, Profile):
envelope.add_profile(profile.to_json(event_opt, self.options))
envelope.add_transaction(event_opt)
elif is_checkin:
envelope.add_checkin(event_opt)
else:
envelope.add_event(event_opt)
for attachment in attachments or ():
envelope.add_item(attachment.to_envelope_item())
return_value = None
if self.spotlight:
self.spotlight.capture_envelope(envelope)
return_value = event_id
if self.transport is not None:
self.transport.capture_envelope(envelope)
return_value = event_id
return return_value
def capture_session(
self, session # type: Session
):
# type: (...) -> None
if not session.release:
logger.info("Discarded session update because of missing release")
else:
self.session_flusher.add_session(session)
if TYPE_CHECKING:
@overload
def get_integration(self, name_or_class):
# type: (str) -> Optional[Integration]
...
@overload
def get_integration(self, name_or_class):
# type: (type[I]) -> Optional[I]
...
def get_integration(
self, name_or_class # type: Union[str, Type[Integration]]
):
# type: (...) -> Optional[Integration]
"""Returns the integration for this client by name or class.
If the client does not have that integration then `None` is returned.
"""
if isinstance(name_or_class, str):
integration_name = name_or_class
elif name_or_class.identifier is not None:
integration_name = name_or_class.identifier
else:
raise ValueError("Integration has no name")
return self.integrations.get(integration_name)
def close(
self,
timeout=None, # type: Optional[float]
callback=None, # type: Optional[Callable[[int, float], None]]
):
# type: (...) -> None
"""
Close the client and shut down the transport. Arguments have the same
semantics as :py:meth:`Client.flush`.
"""
if self.transport is not None:
self.flush(timeout=timeout, callback=callback)
self.session_flusher.kill()
if self.metrics_aggregator is not None:
self.metrics_aggregator.kill()
if self.monitor:
self.monitor.kill()
self.transport.kill()
self.transport = None
def flush(
self,
timeout=None, # type: Optional[float]
callback=None, # type: Optional[Callable[[int, float], None]]
):
# type: (...) -> None
"""
Wait for the current events to be sent.
:param timeout: Wait for at most `timeout` seconds. If no `timeout` is provided, the `shutdown_timeout` option value is used.
:param callback: Is invoked with the number of pending events and the configured timeout.
"""
if self.transport is not None:
if timeout is None:
timeout = self.options["shutdown_timeout"]
self.session_flusher.flush()
if self.metrics_aggregator is not None:
self.metrics_aggregator.flush()
self.transport.flush(timeout=timeout, callback=callback)
def __enter__(self):
# type: () -> _Client
return self
def __exit__(self, exc_type, exc_value, tb):
# type: (Any, Any, Any) -> None
self.close()
from typing import TYPE_CHECKING
if TYPE_CHECKING:
# Make mypy, PyCharm and other static analyzers think `get_options` is a
# type to have nicer autocompletion for params.
#
# Use `ClientConstructor` to define the argument types of `init` and
# `Dict[str, Any]` to tell static analyzers about the return type.
class get_options(ClientConstructor, Dict[str, Any]): # noqa: N801
pass
class Client(ClientConstructor, _Client):
pass
else:
# Alias `get_options` for actual usage. Go through the lambda indirection
# to throw PyCharm off of the weakly typed signature (it would otherwise
# discover both the weakly typed signature of `_init` and our faked `init`
# type).
get_options = (lambda: _get_options)()
Client = (lambda: _Client)()
@@ -0,0 +1,587 @@
import itertools
from enum import Enum
from typing import TYPE_CHECKING
# up top to prevent circular import due to integration import
DEFAULT_MAX_VALUE_LENGTH = 1024
DEFAULT_MAX_STACK_FRAMES = 100
DEFAULT_ADD_FULL_STACK = False
# Also needs to be at the top to prevent circular import
class EndpointType(Enum):
"""
The type of an endpoint. This is an enum, rather than a constant, for historical reasons
(the old /store endpoint). The enum also preserve future compatibility, in case we ever
have a new endpoint.
"""
ENVELOPE = "envelope"
class CompressionAlgo(Enum):
GZIP = "gzip"
BROTLI = "br"
if TYPE_CHECKING:
import sentry_sdk
from typing import Optional
from typing import Callable
from typing import Union
from typing import List
from typing import Type
from typing import Dict
from typing import Any
from typing import Sequence
from typing import Tuple
from typing_extensions import Literal
from typing_extensions import TypedDict
from sentry_sdk._types import (
BreadcrumbProcessor,
ContinuousProfilerMode,
Event,
EventProcessor,
Hint,
MeasurementUnit,
ProfilerMode,
TracesSampler,
TransactionProcessor,
MetricTags,
MetricValue,
)
# Experiments are feature flags to enable and disable certain unstable SDK
# functionality. Changing them from the defaults (`None`) in production
# code is highly discouraged. They are not subject to any stability
# guarantees such as the ones from semantic versioning.
Experiments = TypedDict(
"Experiments",
{
"max_spans": Optional[int],
"max_flags": Optional[int],
"record_sql_params": Optional[bool],
"continuous_profiling_auto_start": Optional[bool],
"continuous_profiling_mode": Optional[ContinuousProfilerMode],
"otel_powered_performance": Optional[bool],
"transport_zlib_compression_level": Optional[int],
"transport_compression_level": Optional[int],
"transport_compression_algo": Optional[CompressionAlgo],
"transport_num_pools": Optional[int],
"transport_http2": Optional[bool],
"enable_metrics": Optional[bool],
"before_emit_metric": Optional[
Callable[[str, MetricValue, MeasurementUnit, MetricTags], bool]
],
"metric_code_locations": Optional[bool],
},
total=False,
)
DEFAULT_QUEUE_SIZE = 100
DEFAULT_MAX_BREADCRUMBS = 100
MATCH_ALL = r".*"
FALSE_VALUES = [
"false",
"no",
"off",
"n",
"0",
]
class INSTRUMENTER:
SENTRY = "sentry"
OTEL = "otel"
class SPANDATA:
"""
Additional information describing the type of the span.
See: https://develop.sentry.dev/sdk/performance/span-data-conventions/
"""
AI_FREQUENCY_PENALTY = "ai.frequency_penalty"
"""
Used to reduce repetitiveness of generated tokens.
Example: 0.5
"""
AI_PRESENCE_PENALTY = "ai.presence_penalty"
"""
Used to reduce repetitiveness of generated tokens.
Example: 0.5
"""
AI_INPUT_MESSAGES = "ai.input_messages"
"""
The input messages to an LLM call.
Example: [{"role": "user", "message": "hello"}]
"""
AI_MODEL_ID = "ai.model_id"
"""
The unique descriptor of the model being execugted
Example: gpt-4
"""
AI_METADATA = "ai.metadata"
"""
Extra metadata passed to an AI pipeline step.
Example: {"executed_function": "add_integers"}
"""
AI_TAGS = "ai.tags"
"""
Tags that describe an AI pipeline step.
Example: {"executed_function": "add_integers"}
"""
AI_STREAMING = "ai.streaming"
"""
Whether or not the AI model call's repsonse was streamed back asynchronously
Example: true
"""
AI_TEMPERATURE = "ai.temperature"
"""
For an AI model call, the temperature parameter. Temperature essentially means how random the output will be.
Example: 0.5
"""
AI_TOP_P = "ai.top_p"
"""
For an AI model call, the top_p parameter. Top_p essentially controls how random the output will be.
Example: 0.5
"""
AI_TOP_K = "ai.top_k"
"""
For an AI model call, the top_k parameter. Top_k essentially controls how random the output will be.
Example: 35
"""
AI_FUNCTION_CALL = "ai.function_call"
"""
For an AI model call, the function that was called. This is deprecated for OpenAI, and replaced by tool_calls
"""
AI_TOOL_CALLS = "ai.tool_calls"
"""
For an AI model call, the function that was called. This is deprecated for OpenAI, and replaced by tool_calls
"""
AI_TOOLS = "ai.tools"
"""
For an AI model call, the functions that are available
"""
AI_RESPONSE_FORMAT = "ai.response_format"
"""
For an AI model call, the format of the response
"""
AI_LOGIT_BIAS = "ai.response_format"
"""
For an AI model call, the logit bias
"""
AI_PREAMBLE = "ai.preamble"
"""
For an AI model call, the preamble parameter.
Preambles are a part of the prompt used to adjust the model's overall behavior and conversation style.
Example: "You are now a clown."
"""
AI_RAW_PROMPTING = "ai.raw_prompting"
"""
Minimize pre-processing done to the prompt sent to the LLM.
Example: true
"""
AI_RESPONSES = "ai.responses"
"""
The responses to an AI model call. Always as a list.
Example: ["hello", "world"]
"""
AI_SEED = "ai.seed"
"""
The seed, ideally models given the same seed and same other parameters will produce the exact same output.
Example: 123.45
"""
DB_NAME = "db.name"
"""
The name of the database being accessed. For commands that switch the database, this should be set to the target database (even if the command fails).
Example: myDatabase
"""
DB_USER = "db.user"
"""
The name of the database user used for connecting to the database.
See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/database.md
Example: my_user
"""
DB_OPERATION = "db.operation"
"""
The name of the operation being executed, e.g. the MongoDB command name such as findAndModify, or the SQL keyword.
See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/database.md
Example: findAndModify, HMSET, SELECT
"""
DB_SYSTEM = "db.system"
"""
An identifier for the database management system (DBMS) product being used.
See: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/database.md
Example: postgresql
"""
DB_MONGODB_COLLECTION = "db.mongodb.collection"
"""
The MongoDB collection being accessed within the database.
See: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/database/mongodb.md#attributes
Example: public.users; customers
"""
CACHE_HIT = "cache.hit"
"""
A boolean indicating whether the requested data was found in the cache.
Example: true
"""
CACHE_ITEM_SIZE = "cache.item_size"
"""
The size of the requested data in bytes.
Example: 58
"""
CACHE_KEY = "cache.key"
"""
The key of the requested data.
Example: template.cache.some_item.867da7e2af8e6b2f3aa7213a4080edb3
"""
NETWORK_PEER_ADDRESS = "network.peer.address"
"""
Peer address of the network connection - IP address or Unix domain socket name.
Example: 10.1.2.80, /tmp/my.sock, localhost
"""
NETWORK_PEER_PORT = "network.peer.port"
"""
Peer port number of the network connection.
Example: 6379
"""
HTTP_QUERY = "http.query"
"""
The Query string present in the URL.
Example: ?foo=bar&bar=baz
"""
HTTP_FRAGMENT = "http.fragment"
"""
The Fragments present in the URL.
Example: #foo=bar
"""
HTTP_METHOD = "http.method"
"""
The HTTP method used.
Example: GET
"""
HTTP_STATUS_CODE = "http.response.status_code"
"""
The HTTP status code as an integer.
Example: 418
"""
MESSAGING_DESTINATION_NAME = "messaging.destination.name"
"""
The destination name where the message is being consumed from,
e.g. the queue name or topic.
"""
MESSAGING_MESSAGE_ID = "messaging.message.id"
"""
The message's identifier.
"""
MESSAGING_MESSAGE_RETRY_COUNT = "messaging.message.retry.count"
"""
Number of retries/attempts to process a message.
"""
MESSAGING_MESSAGE_RECEIVE_LATENCY = "messaging.message.receive.latency"
"""
The latency between when the task was enqueued and when it was started to be processed.
"""
MESSAGING_SYSTEM = "messaging.system"
"""
The messaging system's name, e.g. `kafka`, `aws_sqs`
"""
SERVER_ADDRESS = "server.address"
"""
Name of the database host.
Example: example.com
"""
SERVER_PORT = "server.port"
"""
Logical server port number
Example: 80; 8080; 443
"""
SERVER_SOCKET_ADDRESS = "server.socket.address"
"""
Physical server IP address or Unix socket address.
Example: 10.5.3.2
"""
SERVER_SOCKET_PORT = "server.socket.port"
"""
Physical server port.
Recommended: If different than server.port.
Example: 16456
"""
CODE_FILEPATH = "code.filepath"
"""
The source code file name that identifies the code unit as uniquely as possible (preferably an absolute file path).
Example: "/app/myapplication/http/handler/server.py"
"""
CODE_LINENO = "code.lineno"
"""
The line number in `code.filepath` best representing the operation. It SHOULD point within the code unit named in `code.function`.
Example: 42
"""
CODE_FUNCTION = "code.function"
"""
The method or function name, or equivalent (usually rightmost part of the code unit's name).
Example: "server_request"
"""
CODE_NAMESPACE = "code.namespace"
"""
The "namespace" within which `code.function` is defined. Usually the qualified class or module name, such that `code.namespace` + some separator + `code.function` form a unique identifier for the code unit.
Example: "http.handler"
"""
THREAD_ID = "thread.id"
"""
Identifier of a thread from where the span originated. This should be a string.
Example: "7972576320"
"""
THREAD_NAME = "thread.name"
"""
Label identifying a thread from where the span originated. This should be a string.
Example: "MainThread"
"""
PROFILER_ID = "profiler_id"
"""
Label identifying the profiler id that the span occurred in. This should be a string.
Example: "5249fbada8d5416482c2f6e47e337372"
"""
class SPANSTATUS:
"""
The status of a Sentry span.
See: https://develop.sentry.dev/sdk/event-payloads/contexts/#trace-context
"""
ABORTED = "aborted"
ALREADY_EXISTS = "already_exists"
CANCELLED = "cancelled"
DATA_LOSS = "data_loss"
DEADLINE_EXCEEDED = "deadline_exceeded"
FAILED_PRECONDITION = "failed_precondition"
INTERNAL_ERROR = "internal_error"
INVALID_ARGUMENT = "invalid_argument"
NOT_FOUND = "not_found"
OK = "ok"
OUT_OF_RANGE = "out_of_range"
PERMISSION_DENIED = "permission_denied"
RESOURCE_EXHAUSTED = "resource_exhausted"
UNAUTHENTICATED = "unauthenticated"
UNAVAILABLE = "unavailable"
UNIMPLEMENTED = "unimplemented"
UNKNOWN_ERROR = "unknown_error"
class OP:
ANTHROPIC_MESSAGES_CREATE = "ai.messages.create.anthropic"
CACHE_GET = "cache.get"
CACHE_PUT = "cache.put"
COHERE_CHAT_COMPLETIONS_CREATE = "ai.chat_completions.create.cohere"
COHERE_EMBEDDINGS_CREATE = "ai.embeddings.create.cohere"
DB = "db"
DB_REDIS = "db.redis"
EVENT_DJANGO = "event.django"
FUNCTION = "function"
FUNCTION_AWS = "function.aws"
FUNCTION_GCP = "function.gcp"
GRAPHQL_EXECUTE = "graphql.execute"
GRAPHQL_MUTATION = "graphql.mutation"
GRAPHQL_PARSE = "graphql.parse"
GRAPHQL_RESOLVE = "graphql.resolve"
GRAPHQL_SUBSCRIPTION = "graphql.subscription"
GRAPHQL_QUERY = "graphql.query"
GRAPHQL_VALIDATE = "graphql.validate"
GRPC_CLIENT = "grpc.client"
GRPC_SERVER = "grpc.server"
HTTP_CLIENT = "http.client"
HTTP_CLIENT_STREAM = "http.client.stream"
HTTP_SERVER = "http.server"
MIDDLEWARE_DJANGO = "middleware.django"
MIDDLEWARE_LITESTAR = "middleware.litestar"
MIDDLEWARE_LITESTAR_RECEIVE = "middleware.litestar.receive"
MIDDLEWARE_LITESTAR_SEND = "middleware.litestar.send"
MIDDLEWARE_STARLETTE = "middleware.starlette"
MIDDLEWARE_STARLETTE_RECEIVE = "middleware.starlette.receive"
MIDDLEWARE_STARLETTE_SEND = "middleware.starlette.send"
MIDDLEWARE_STARLITE = "middleware.starlite"
MIDDLEWARE_STARLITE_RECEIVE = "middleware.starlite.receive"
MIDDLEWARE_STARLITE_SEND = "middleware.starlite.send"
OPENAI_CHAT_COMPLETIONS_CREATE = "ai.chat_completions.create.openai"
OPENAI_EMBEDDINGS_CREATE = "ai.embeddings.create.openai"
HUGGINGFACE_HUB_CHAT_COMPLETIONS_CREATE = (
"ai.chat_completions.create.huggingface_hub"
)
LANGCHAIN_PIPELINE = "ai.pipeline.langchain"
LANGCHAIN_RUN = "ai.run.langchain"
LANGCHAIN_TOOL = "ai.tool.langchain"
LANGCHAIN_AGENT = "ai.agent.langchain"
LANGCHAIN_CHAT_COMPLETIONS_CREATE = "ai.chat_completions.create.langchain"
QUEUE_PROCESS = "queue.process"
QUEUE_PUBLISH = "queue.publish"
QUEUE_SUBMIT_ARQ = "queue.submit.arq"
QUEUE_TASK_ARQ = "queue.task.arq"
QUEUE_SUBMIT_CELERY = "queue.submit.celery"
QUEUE_TASK_CELERY = "queue.task.celery"
QUEUE_TASK_RQ = "queue.task.rq"
QUEUE_SUBMIT_HUEY = "queue.submit.huey"
QUEUE_TASK_HUEY = "queue.task.huey"
QUEUE_SUBMIT_RAY = "queue.submit.ray"
QUEUE_TASK_RAY = "queue.task.ray"
SUBPROCESS = "subprocess"
SUBPROCESS_WAIT = "subprocess.wait"
SUBPROCESS_COMMUNICATE = "subprocess.communicate"
TEMPLATE_RENDER = "template.render"
VIEW_RENDER = "view.render"
VIEW_RESPONSE_RENDER = "view.response.render"
WEBSOCKET_SERVER = "websocket.server"
SOCKET_CONNECTION = "socket.connection"
SOCKET_DNS = "socket.dns"
# This type exists to trick mypy and PyCharm into thinking `init` and `Client`
# take these arguments (even though they take opaque **kwargs)
class ClientConstructor:
def __init__(
self,
dsn=None, # type: Optional[str]
*,
max_breadcrumbs=DEFAULT_MAX_BREADCRUMBS, # type: int
release=None, # type: Optional[str]
environment=None, # type: Optional[str]
server_name=None, # type: Optional[str]
shutdown_timeout=2, # type: float
integrations=[], # type: Sequence[sentry_sdk.integrations.Integration] # noqa: B006
in_app_include=[], # type: List[str] # noqa: B006
in_app_exclude=[], # type: List[str] # noqa: B006
default_integrations=True, # type: bool
dist=None, # type: Optional[str]
transport=None, # type: Optional[Union[sentry_sdk.transport.Transport, Type[sentry_sdk.transport.Transport], Callable[[Event], None]]]
transport_queue_size=DEFAULT_QUEUE_SIZE, # type: int
sample_rate=1.0, # type: float
send_default_pii=None, # type: Optional[bool]
http_proxy=None, # type: Optional[str]
https_proxy=None, # type: Optional[str]
ignore_errors=[], # type: Sequence[Union[type, str]] # noqa: B006
max_request_body_size="medium", # type: str
socket_options=None, # type: Optional[List[Tuple[int, int, int | bytes]]]
keep_alive=False, # type: bool
before_send=None, # type: Optional[EventProcessor]
before_breadcrumb=None, # type: Optional[BreadcrumbProcessor]
debug=None, # type: Optional[bool]
attach_stacktrace=False, # type: bool
ca_certs=None, # type: Optional[str]
propagate_traces=True, # type: bool
traces_sample_rate=None, # type: Optional[float]
traces_sampler=None, # type: Optional[TracesSampler]
profiles_sample_rate=None, # type: Optional[float]
profiles_sampler=None, # type: Optional[TracesSampler]
profiler_mode=None, # type: Optional[ProfilerMode]
profile_lifecycle="manual", # type: Literal["manual", "trace"]
profile_session_sample_rate=None, # type: Optional[float]
auto_enabling_integrations=True, # type: bool
disabled_integrations=None, # type: Optional[Sequence[sentry_sdk.integrations.Integration]]
auto_session_tracking=True, # type: bool
send_client_reports=True, # type: bool
_experiments={}, # type: Experiments # noqa: B006
proxy_headers=None, # type: Optional[Dict[str, str]]
instrumenter=INSTRUMENTER.SENTRY, # type: Optional[str]
before_send_transaction=None, # type: Optional[TransactionProcessor]
project_root=None, # type: Optional[str]
enable_tracing=None, # type: Optional[bool]
include_local_variables=True, # type: Optional[bool]
include_source_context=True, # type: Optional[bool]
trace_propagation_targets=[ # noqa: B006
MATCH_ALL
], # type: Optional[Sequence[str]]
functions_to_trace=[], # type: Sequence[Dict[str, str]] # noqa: B006
event_scrubber=None, # type: Optional[sentry_sdk.scrubber.EventScrubber]
max_value_length=DEFAULT_MAX_VALUE_LENGTH, # type: int
enable_backpressure_handling=True, # type: bool
error_sampler=None, # type: Optional[Callable[[Event, Hint], Union[float, bool]]]
enable_db_query_source=True, # type: bool
db_query_source_threshold_ms=100, # type: int
spotlight=None, # type: Optional[Union[bool, str]]
cert_file=None, # type: Optional[str]
key_file=None, # type: Optional[str]
custom_repr=None, # type: Optional[Callable[..., Optional[str]]]
add_full_stack=DEFAULT_ADD_FULL_STACK, # type: bool
max_stack_frames=DEFAULT_MAX_STACK_FRAMES, # type: Optional[int]
):
# type: (...) -> None
pass
def _get_default_options():
# type: () -> dict[str, Any]
import inspect
a = inspect.getfullargspec(ClientConstructor.__init__)
defaults = a.defaults or ()
kwonlydefaults = a.kwonlydefaults or {}
return dict(
itertools.chain(
zip(a.args[-len(defaults) :], defaults),
kwonlydefaults.items(),
)
)
DEFAULT_OPTIONS = _get_default_options()
del _get_default_options
VERSION = "2.22.0"
@@ -0,0 +1,10 @@
from sentry_sdk.crons.api import capture_checkin
from sentry_sdk.crons.consts import MonitorStatus
from sentry_sdk.crons.decorator import monitor
__all__ = [
"capture_checkin",
"MonitorStatus",
"monitor",
]
@@ -0,0 +1,57 @@
import uuid
import sentry_sdk
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional
from sentry_sdk._types import Event, MonitorConfig
def _create_check_in_event(
monitor_slug=None, # type: Optional[str]
check_in_id=None, # type: Optional[str]
status=None, # type: Optional[str]
duration_s=None, # type: Optional[float]
monitor_config=None, # type: Optional[MonitorConfig]
):
# type: (...) -> Event
options = sentry_sdk.get_client().options
check_in_id = check_in_id or uuid.uuid4().hex # type: str
check_in = {
"type": "check_in",
"monitor_slug": monitor_slug,
"check_in_id": check_in_id,
"status": status,
"duration": duration_s,
"environment": options.get("environment", None),
"release": options.get("release", None),
} # type: Event
if monitor_config:
check_in["monitor_config"] = monitor_config
return check_in
def capture_checkin(
monitor_slug=None, # type: Optional[str]
check_in_id=None, # type: Optional[str]
status=None, # type: Optional[str]
duration=None, # type: Optional[float]
monitor_config=None, # type: Optional[MonitorConfig]
):
# type: (...) -> str
check_in_event = _create_check_in_event(
monitor_slug=monitor_slug,
check_in_id=check_in_id,
status=status,
duration_s=duration,
monitor_config=monitor_config,
)
sentry_sdk.capture_event(check_in_event)
return check_in_event["check_in_id"]
@@ -0,0 +1,4 @@
class MonitorStatus:
IN_PROGRESS = "in_progress"
OK = "ok"
ERROR = "error"
@@ -0,0 +1,135 @@
from functools import wraps
from inspect import iscoroutinefunction
from sentry_sdk.crons import capture_checkin
from sentry_sdk.crons.consts import MonitorStatus
from sentry_sdk.utils import now
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from types import TracebackType
from typing import (
Any,
Optional,
ParamSpec,
Type,
TypeVar,
Union,
cast,
overload,
)
from sentry_sdk._types import MonitorConfig
P = ParamSpec("P")
R = TypeVar("R")
class monitor: # noqa: N801
"""
Decorator/context manager to capture checkin events for a monitor.
Usage (as decorator):
```
import sentry_sdk
app = Celery()
@app.task
@sentry_sdk.monitor(monitor_slug='my-fancy-slug')
def test(arg):
print(arg)
```
This does not have to be used with Celery, but if you do use it with celery,
put the `@sentry_sdk.monitor` decorator below Celery's `@app.task` decorator.
Usage (as context manager):
```
import sentry_sdk
def test(arg):
with sentry_sdk.monitor(monitor_slug='my-fancy-slug'):
print(arg)
```
"""
def __init__(self, monitor_slug=None, monitor_config=None):
# type: (Optional[str], Optional[MonitorConfig]) -> None
self.monitor_slug = monitor_slug
self.monitor_config = monitor_config
def __enter__(self):
# type: () -> None
self.start_timestamp = now()
self.check_in_id = capture_checkin(
monitor_slug=self.monitor_slug,
status=MonitorStatus.IN_PROGRESS,
monitor_config=self.monitor_config,
)
def __exit__(self, exc_type, exc_value, traceback):
# type: (Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]) -> None
duration_s = now() - self.start_timestamp
if exc_type is None and exc_value is None and traceback is None:
status = MonitorStatus.OK
else:
status = MonitorStatus.ERROR
capture_checkin(
monitor_slug=self.monitor_slug,
check_in_id=self.check_in_id,
status=status,
duration=duration_s,
monitor_config=self.monitor_config,
)
if TYPE_CHECKING:
@overload
def __call__(self, fn):
# type: (Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[Any]]
# Unfortunately, mypy does not give us any reliable way to type check the
# return value of an Awaitable (i.e. async function) for this overload,
# since calling iscouroutinefunction narrows the type to Callable[P, Awaitable[Any]].
...
@overload
def __call__(self, fn):
# type: (Callable[P, R]) -> Callable[P, R]
...
def __call__(
self,
fn, # type: Union[Callable[P, R], Callable[P, Awaitable[Any]]]
):
# type: (...) -> Union[Callable[P, R], Callable[P, Awaitable[Any]]]
if iscoroutinefunction(fn):
return self._async_wrapper(fn)
else:
if TYPE_CHECKING:
fn = cast("Callable[P, R]", fn)
return self._sync_wrapper(fn)
def _async_wrapper(self, fn):
# type: (Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[Any]]
@wraps(fn)
async def inner(*args: "P.args", **kwargs: "P.kwargs"):
# type: (...) -> R
with self:
return await fn(*args, **kwargs)
return inner
def _sync_wrapper(self, fn):
# type: (Callable[P, R]) -> Callable[P, R]
@wraps(fn)
def inner(*args: "P.args", **kwargs: "P.kwargs"):
# type: (...) -> R
with self:
return fn(*args, **kwargs)
return inner
@@ -0,0 +1,41 @@
import sys
import logging
import warnings
from sentry_sdk import get_client
from sentry_sdk.client import _client_init_debug
from sentry_sdk.utils import logger
from logging import LogRecord
class _DebugFilter(logging.Filter):
def filter(self, record):
# type: (LogRecord) -> bool
if _client_init_debug.get(False):
return True
return get_client().options["debug"]
def init_debug_support():
# type: () -> None
if not logger.handlers:
configure_logger()
def configure_logger():
# type: () -> None
_handler = logging.StreamHandler(sys.stderr)
_handler.setFormatter(logging.Formatter(" [sentry] %(levelname)s: %(message)s"))
logger.addHandler(_handler)
logger.setLevel(logging.DEBUG)
logger.addFilter(_DebugFilter())
def configure_debug_hub():
# type: () -> None
warnings.warn(
"configure_debug_hub is deprecated. Please remove calls to it, as it is a no-op.",
DeprecationWarning,
stacklevel=2,
)
@@ -0,0 +1,349 @@
import io
import json
import mimetypes
from sentry_sdk.session import Session
from sentry_sdk.utils import json_dumps, capture_internal_exceptions
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Optional
from typing import Union
from typing import Dict
from typing import List
from typing import Iterator
from sentry_sdk._types import Event, EventDataCategory
def parse_json(data):
# type: (Union[bytes, str]) -> Any
# on some python 3 versions this needs to be bytes
if isinstance(data, bytes):
data = data.decode("utf-8", "replace")
return json.loads(data)
class Envelope:
"""
Represents a Sentry Envelope. The calling code is responsible for adhering to the constraints
documented in the Sentry docs: https://develop.sentry.dev/sdk/envelopes/#data-model. In particular,
each envelope may have at most one Item with type "event" or "transaction" (but not both).
"""
def __init__(
self,
headers=None, # type: Optional[Dict[str, Any]]
items=None, # type: Optional[List[Item]]
):
# type: (...) -> None
if headers is not None:
headers = dict(headers)
self.headers = headers or {}
if items is None:
items = []
else:
items = list(items)
self.items = items
@property
def description(self):
# type: (...) -> str
return "envelope with %s items (%s)" % (
len(self.items),
", ".join(x.data_category for x in self.items),
)
def add_event(
self, event # type: Event
):
# type: (...) -> None
self.add_item(Item(payload=PayloadRef(json=event), type="event"))
def add_transaction(
self, transaction # type: Event
):
# type: (...) -> None
self.add_item(Item(payload=PayloadRef(json=transaction), type="transaction"))
def add_profile(
self, profile # type: Any
):
# type: (...) -> None
self.add_item(Item(payload=PayloadRef(json=profile), type="profile"))
def add_profile_chunk(
self, profile_chunk # type: Any
):
# type: (...) -> None
self.add_item(
Item(payload=PayloadRef(json=profile_chunk), type="profile_chunk")
)
def add_checkin(
self, checkin # type: Any
):
# type: (...) -> None
self.add_item(Item(payload=PayloadRef(json=checkin), type="check_in"))
def add_session(
self, session # type: Union[Session, Any]
):
# type: (...) -> None
if isinstance(session, Session):
session = session.to_json()
self.add_item(Item(payload=PayloadRef(json=session), type="session"))
def add_sessions(
self, sessions # type: Any
):
# type: (...) -> None
self.add_item(Item(payload=PayloadRef(json=sessions), type="sessions"))
def add_item(
self, item # type: Item
):
# type: (...) -> None
self.items.append(item)
def get_event(self):
# type: (...) -> Optional[Event]
for items in self.items:
event = items.get_event()
if event is not None:
return event
return None
def get_transaction_event(self):
# type: (...) -> Optional[Event]
for item in self.items:
event = item.get_transaction_event()
if event is not None:
return event
return None
def __iter__(self):
# type: (...) -> Iterator[Item]
return iter(self.items)
def serialize_into(
self, f # type: Any
):
# type: (...) -> None
f.write(json_dumps(self.headers))
f.write(b"\n")
for item in self.items:
item.serialize_into(f)
def serialize(self):
# type: (...) -> bytes
out = io.BytesIO()
self.serialize_into(out)
return out.getvalue()
@classmethod
def deserialize_from(
cls, f # type: Any
):
# type: (...) -> Envelope
headers = parse_json(f.readline())
items = []
while 1:
item = Item.deserialize_from(f)
if item is None:
break
items.append(item)
return cls(headers=headers, items=items)
@classmethod
def deserialize(
cls, bytes # type: bytes
):
# type: (...) -> Envelope
return cls.deserialize_from(io.BytesIO(bytes))
def __repr__(self):
# type: (...) -> str
return "<Envelope headers=%r items=%r>" % (self.headers, self.items)
class PayloadRef:
def __init__(
self,
bytes=None, # type: Optional[bytes]
path=None, # type: Optional[Union[bytes, str]]
json=None, # type: Optional[Any]
):
# type: (...) -> None
self.json = json
self.bytes = bytes
self.path = path
def get_bytes(self):
# type: (...) -> bytes
if self.bytes is None:
if self.path is not None:
with capture_internal_exceptions():
with open(self.path, "rb") as f:
self.bytes = f.read()
elif self.json is not None:
self.bytes = json_dumps(self.json)
return self.bytes or b""
@property
def inferred_content_type(self):
# type: (...) -> str
if self.json is not None:
return "application/json"
elif self.path is not None:
path = self.path
if isinstance(path, bytes):
path = path.decode("utf-8", "replace")
ty = mimetypes.guess_type(path)[0]
if ty:
return ty
return "application/octet-stream"
def __repr__(self):
# type: (...) -> str
return "<Payload %r>" % (self.inferred_content_type,)
class Item:
def __init__(
self,
payload, # type: Union[bytes, str, PayloadRef]
headers=None, # type: Optional[Dict[str, Any]]
type=None, # type: Optional[str]
content_type=None, # type: Optional[str]
filename=None, # type: Optional[str]
):
if headers is not None:
headers = dict(headers)
elif headers is None:
headers = {}
self.headers = headers
if isinstance(payload, bytes):
payload = PayloadRef(bytes=payload)
elif isinstance(payload, str):
payload = PayloadRef(bytes=payload.encode("utf-8"))
else:
payload = payload
if filename is not None:
headers["filename"] = filename
if type is not None:
headers["type"] = type
if content_type is not None:
headers["content_type"] = content_type
elif "content_type" not in headers:
headers["content_type"] = payload.inferred_content_type
self.payload = payload
def __repr__(self):
# type: (...) -> str
return "<Item headers=%r payload=%r data_category=%r>" % (
self.headers,
self.payload,
self.data_category,
)
@property
def type(self):
# type: (...) -> Optional[str]
return self.headers.get("type")
@property
def data_category(self):
# type: (...) -> EventDataCategory
ty = self.headers.get("type")
if ty == "session" or ty == "sessions":
return "session"
elif ty == "attachment":
return "attachment"
elif ty == "transaction":
return "transaction"
elif ty == "event":
return "error"
elif ty == "client_report":
return "internal"
elif ty == "profile":
return "profile"
elif ty == "profile_chunk":
return "profile_chunk"
elif ty == "statsd":
return "metric_bucket"
elif ty == "check_in":
return "monitor"
else:
return "default"
def get_bytes(self):
# type: (...) -> bytes
return self.payload.get_bytes()
def get_event(self):
# type: (...) -> Optional[Event]
"""
Returns an error event if there is one.
"""
if self.type == "event" and self.payload.json is not None:
return self.payload.json
return None
def get_transaction_event(self):
# type: (...) -> Optional[Event]
if self.type == "transaction" and self.payload.json is not None:
return self.payload.json
return None
def serialize_into(
self, f # type: Any
):
# type: (...) -> None
headers = dict(self.headers)
bytes = self.get_bytes()
headers["length"] = len(bytes)
f.write(json_dumps(headers))
f.write(b"\n")
f.write(bytes)
f.write(b"\n")
def serialize(self):
# type: (...) -> bytes
out = io.BytesIO()
self.serialize_into(out)
return out.getvalue()
@classmethod
def deserialize_from(
cls, f # type: Any
):
# type: (...) -> Optional[Item]
line = f.readline().rstrip()
if not line:
return None
headers = parse_json(line)
length = headers.get("length")
if length is not None:
payload = f.read(length)
f.readline()
else:
# if no length was specified we need to read up to the end of line
# and remove it (if it is present, i.e. not the very last char in an eof terminated envelope)
payload = f.readline().rstrip(b"\n")
if headers.get("type") in ("event", "transaction", "metric_buckets"):
rv = cls(headers=headers, payload=PayloadRef(json=parse_json(payload)))
else:
rv = cls(headers=headers, payload=payload)
return rv
@classmethod
def deserialize(
cls, bytes # type: bytes
):
# type: (...) -> Optional[Item]
return cls.deserialize_from(io.BytesIO(bytes))
@@ -0,0 +1,68 @@
import copy
import sentry_sdk
from sentry_sdk._lru_cache import LRUCache
from threading import Lock
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from typing import TypedDict
FlagData = TypedDict("FlagData", {"flag": str, "result": bool})
DEFAULT_FLAG_CAPACITY = 100
class FlagBuffer:
def __init__(self, capacity):
# type: (int) -> None
self.capacity = capacity
self.lock = Lock()
# Buffer is private. The name is mangled to discourage use. If you use this attribute
# directly you're on your own!
self.__buffer = LRUCache(capacity)
def clear(self):
# type: () -> None
self.__buffer = LRUCache(self.capacity)
def __deepcopy__(self, memo):
# type: (dict[int, Any]) -> FlagBuffer
with self.lock:
buffer = FlagBuffer(self.capacity)
buffer.__buffer = copy.deepcopy(self.__buffer, memo)
return buffer
def get(self):
# type: () -> list[FlagData]
with self.lock:
return [
{"flag": key, "result": value} for key, value in self.__buffer.get_all()
]
def set(self, flag, result):
# type: (str, bool) -> None
if isinstance(result, FlagBuffer):
# If someone were to insert `self` into `self` this would create a circular dependency
# on the lock. This is of course a deadlock. However, this is far outside the expected
# usage of this class. We guard against it here for completeness and to document this
# expected failure mode.
raise ValueError(
"FlagBuffer instances can not be inserted into the dictionary."
)
with self.lock:
self.__buffer.set(flag, result)
def add_feature_flag(flag, result):
# type: (str, bool) -> None
"""
Records a flag and its value to be sent on subsequent error events.
We recommend you do this on flag evaluations. Flags are buffered per Sentry scope.
"""
flags = sentry_sdk.get_current_scope().flags
flags.set(flag, result)
@@ -0,0 +1,739 @@
import warnings
from contextlib import contextmanager
from sentry_sdk import (
get_client,
get_global_scope,
get_isolation_scope,
get_current_scope,
)
from sentry_sdk._compat import with_metaclass
from sentry_sdk.consts import INSTRUMENTER
from sentry_sdk.scope import _ScopeManager
from sentry_sdk.client import Client
from sentry_sdk.tracing import (
NoOpSpan,
Span,
Transaction,
)
from sentry_sdk.utils import (
logger,
ContextVar,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import ContextManager
from typing import Dict
from typing import Generator
from typing import List
from typing import Optional
from typing import overload
from typing import Tuple
from typing import Type
from typing import TypeVar
from typing import Union
from typing_extensions import Unpack
from sentry_sdk.scope import Scope
from sentry_sdk.client import BaseClient
from sentry_sdk.integrations import Integration
from sentry_sdk._types import (
Event,
Hint,
Breadcrumb,
BreadcrumbHint,
ExcInfo,
LogLevelStr,
SamplingContext,
)
from sentry_sdk.tracing import TransactionKwargs
T = TypeVar("T")
else:
def overload(x):
# type: (T) -> T
return x
class SentryHubDeprecationWarning(DeprecationWarning):
"""
A custom deprecation warning to inform users that the Hub is deprecated.
"""
_MESSAGE = (
"`sentry_sdk.Hub` is deprecated and will be removed in a future major release. "
"Please consult our 1.x to 2.x migration guide for details on how to migrate "
"`Hub` usage to the new API: "
"https://docs.sentry.io/platforms/python/migration/1.x-to-2.x"
)
def __init__(self, *_):
# type: (*object) -> None
super().__init__(self._MESSAGE)
@contextmanager
def _suppress_hub_deprecation_warning():
# type: () -> Generator[None, None, None]
"""Utility function to suppress deprecation warnings for the Hub."""
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=SentryHubDeprecationWarning)
yield
_local = ContextVar("sentry_current_hub")
class HubMeta(type):
@property
def current(cls):
# type: () -> Hub
"""Returns the current instance of the hub."""
warnings.warn(SentryHubDeprecationWarning(), stacklevel=2)
rv = _local.get(None)
if rv is None:
with _suppress_hub_deprecation_warning():
# This will raise a deprecation warning; suppress it since we already warned above.
rv = Hub(GLOBAL_HUB)
_local.set(rv)
return rv
@property
def main(cls):
# type: () -> Hub
"""Returns the main instance of the hub."""
warnings.warn(SentryHubDeprecationWarning(), stacklevel=2)
return GLOBAL_HUB
class Hub(with_metaclass(HubMeta)): # type: ignore
"""
.. deprecated:: 2.0.0
The Hub is deprecated. Its functionality will be merged into :py:class:`sentry_sdk.scope.Scope`.
The hub wraps the concurrency management of the SDK. Each thread has
its own hub but the hub might transfer with the flow of execution if
context vars are available.
If the hub is used with a with statement it's temporarily activated.
"""
_stack = None # type: List[Tuple[Optional[Client], Scope]]
_scope = None # type: Optional[Scope]
# Mypy doesn't pick up on the metaclass.
if TYPE_CHECKING:
current = None # type: Hub
main = None # type: Hub
def __init__(
self,
client_or_hub=None, # type: Optional[Union[Hub, Client]]
scope=None, # type: Optional[Any]
):
# type: (...) -> None
warnings.warn(SentryHubDeprecationWarning(), stacklevel=2)
current_scope = None
if isinstance(client_or_hub, Hub):
client = get_client()
if scope is None:
# hub cloning is going on, we use a fork of the current/isolation scope for context manager
scope = get_isolation_scope().fork()
current_scope = get_current_scope().fork()
else:
client = client_or_hub # type: ignore
get_global_scope().set_client(client)
if scope is None: # so there is no Hub cloning going on
# just the current isolation scope is used for context manager
scope = get_isolation_scope()
current_scope = get_current_scope()
if current_scope is None:
# just the current current scope is used for context manager
current_scope = get_current_scope()
self._stack = [(client, scope)] # type: ignore
self._last_event_id = None # type: Optional[str]
self._old_hubs = [] # type: List[Hub]
self._old_current_scopes = [] # type: List[Scope]
self._old_isolation_scopes = [] # type: List[Scope]
self._current_scope = current_scope # type: Scope
self._scope = scope # type: Scope
def __enter__(self):
# type: () -> Hub
self._old_hubs.append(Hub.current)
_local.set(self)
current_scope = get_current_scope()
self._old_current_scopes.append(current_scope)
scope._current_scope.set(self._current_scope)
isolation_scope = get_isolation_scope()
self._old_isolation_scopes.append(isolation_scope)
scope._isolation_scope.set(self._scope)
return self
def __exit__(
self,
exc_type, # type: Optional[type]
exc_value, # type: Optional[BaseException]
tb, # type: Optional[Any]
):
# type: (...) -> None
old = self._old_hubs.pop()
_local.set(old)
old_current_scope = self._old_current_scopes.pop()
scope._current_scope.set(old_current_scope)
old_isolation_scope = self._old_isolation_scopes.pop()
scope._isolation_scope.set(old_isolation_scope)
def run(
self, callback # type: Callable[[], T]
):
# type: (...) -> T
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Runs a callback in the context of the hub. Alternatively the
with statement can be used on the hub directly.
"""
with self:
return callback()
def get_integration(
self, name_or_class # type: Union[str, Type[Integration]]
):
# type: (...) -> Any
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.client._Client.get_integration` instead.
Returns the integration for this hub by name or class. If there
is no client bound or the client does not have that integration
then `None` is returned.
If the return value is not `None` the hub is guaranteed to have a
client attached.
"""
return get_client().get_integration(name_or_class)
@property
def client(self):
# type: () -> Optional[BaseClient]
"""
.. deprecated:: 2.0.0
This property is deprecated and will be removed in a future release.
Please use :py:func:`sentry_sdk.api.get_client` instead.
Returns the current client on the hub.
"""
client = get_client()
if not client.is_active():
return None
return client
@property
def scope(self):
# type: () -> Scope
"""
.. deprecated:: 2.0.0
This property is deprecated and will be removed in a future release.
Returns the current scope on the hub.
"""
return get_isolation_scope()
def last_event_id(self):
# type: () -> Optional[str]
"""
Returns the last event ID.
.. deprecated:: 1.40.5
This function is deprecated and will be removed in a future release. The functions `capture_event`, `capture_message`, and `capture_exception` return the event ID directly.
"""
logger.warning(
"Deprecated: last_event_id is deprecated. This will be removed in the future. The functions `capture_event`, `capture_message`, and `capture_exception` return the event ID directly."
)
return self._last_event_id
def bind_client(
self, new # type: Optional[BaseClient]
):
# type: (...) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.set_client` instead.
Binds a new client to the hub.
"""
get_global_scope().set_client(new)
def capture_event(self, event, hint=None, scope=None, **scope_kwargs):
# type: (Event, Optional[Hint], Optional[Scope], Any) -> Optional[str]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.capture_event` instead.
Captures an event.
Alias of :py:meth:`sentry_sdk.Scope.capture_event`.
:param event: A ready-made event that can be directly sent to Sentry.
:param hint: Contains metadata about the event that can be read from `before_send`, such as the original exception object or a HTTP request object.
:param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events.
The `scope` and `scope_kwargs` parameters are mutually exclusive.
:param scope_kwargs: Optional data to apply to event.
For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`.
The `scope` and `scope_kwargs` parameters are mutually exclusive.
"""
last_event_id = get_current_scope().capture_event(
event, hint, scope=scope, **scope_kwargs
)
is_transaction = event.get("type") == "transaction"
if last_event_id is not None and not is_transaction:
self._last_event_id = last_event_id
return last_event_id
def capture_message(self, message, level=None, scope=None, **scope_kwargs):
# type: (str, Optional[LogLevelStr], Optional[Scope], Any) -> Optional[str]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.capture_message` instead.
Captures a message.
Alias of :py:meth:`sentry_sdk.Scope.capture_message`.
:param message: The string to send as the message to Sentry.
:param level: If no level is provided, the default level is `info`.
:param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events.
The `scope` and `scope_kwargs` parameters are mutually exclusive.
:param scope_kwargs: Optional data to apply to event.
For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`.
The `scope` and `scope_kwargs` parameters are mutually exclusive.
:returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`).
"""
last_event_id = get_current_scope().capture_message(
message, level=level, scope=scope, **scope_kwargs
)
if last_event_id is not None:
self._last_event_id = last_event_id
return last_event_id
def capture_exception(self, error=None, scope=None, **scope_kwargs):
# type: (Optional[Union[BaseException, ExcInfo]], Optional[Scope], Any) -> Optional[str]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.capture_exception` instead.
Captures an exception.
Alias of :py:meth:`sentry_sdk.Scope.capture_exception`.
:param error: An exception to capture. If `None`, `sys.exc_info()` will be used.
:param scope: An optional :py:class:`sentry_sdk.Scope` to apply to events.
The `scope` and `scope_kwargs` parameters are mutually exclusive.
:param scope_kwargs: Optional data to apply to event.
For supported `**scope_kwargs` see :py:meth:`sentry_sdk.Scope.update_from_kwargs`.
The `scope` and `scope_kwargs` parameters are mutually exclusive.
:returns: An `event_id` if the SDK decided to send the event (see :py:meth:`sentry_sdk.client._Client.capture_event`).
"""
last_event_id = get_current_scope().capture_exception(
error, scope=scope, **scope_kwargs
)
if last_event_id is not None:
self._last_event_id = last_event_id
return last_event_id
def add_breadcrumb(self, crumb=None, hint=None, **kwargs):
# type: (Optional[Breadcrumb], Optional[BreadcrumbHint], Any) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.add_breadcrumb` instead.
Adds a breadcrumb.
:param crumb: Dictionary with the data as the sentry v7/v8 protocol expects.
:param hint: An optional value that can be used by `before_breadcrumb`
to customize the breadcrumbs that are emitted.
"""
get_isolation_scope().add_breadcrumb(crumb, hint, **kwargs)
def start_span(self, instrumenter=INSTRUMENTER.SENTRY, **kwargs):
# type: (str, Any) -> Span
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.start_span` instead.
Start a span whose parent is the currently active span or transaction, if any.
The return value is a :py:class:`sentry_sdk.tracing.Span` instance,
typically used as a context manager to start and stop timing in a `with`
block.
Only spans contained in a transaction are sent to Sentry. Most
integrations start a transaction at the appropriate time, for example
for every incoming HTTP request. Use
:py:meth:`sentry_sdk.start_transaction` to start a new transaction when
one is not already in progress.
For supported `**kwargs` see :py:class:`sentry_sdk.tracing.Span`.
"""
scope = get_current_scope()
return scope.start_span(instrumenter=instrumenter, **kwargs)
def start_transaction(
self,
transaction=None,
instrumenter=INSTRUMENTER.SENTRY,
custom_sampling_context=None,
**kwargs
):
# type: (Optional[Transaction], str, Optional[SamplingContext], Unpack[TransactionKwargs]) -> Union[Transaction, NoOpSpan]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.start_transaction` instead.
Start and return a transaction.
Start an existing transaction if given, otherwise create and start a new
transaction with kwargs.
This is the entry point to manual tracing instrumentation.
A tree structure can be built by adding child spans to the transaction,
and child spans to other spans. To start a new child span within the
transaction or any span, call the respective `.start_child()` method.
Every child span must be finished before the transaction is finished,
otherwise the unfinished spans are discarded.
When used as context managers, spans and transactions are automatically
finished at the end of the `with` block. If not using context managers,
call the `.finish()` method.
When the transaction is finished, it will be sent to Sentry with all its
finished child spans.
For supported `**kwargs` see :py:class:`sentry_sdk.tracing.Transaction`.
"""
scope = get_current_scope()
# For backwards compatibility, we allow passing the scope as the hub.
# We need a major release to make this nice. (if someone searches the code: deprecated)
# Type checking disabled for this line because deprecated keys are not allowed in the type signature.
kwargs["hub"] = scope # type: ignore
return scope.start_transaction(
transaction, instrumenter, custom_sampling_context, **kwargs
)
def continue_trace(self, environ_or_headers, op=None, name=None, source=None):
# type: (Dict[str, Any], Optional[str], Optional[str], Optional[str]) -> Transaction
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.continue_trace` instead.
Sets the propagation context from environment or headers and returns a transaction.
"""
return get_isolation_scope().continue_trace(
environ_or_headers=environ_or_headers, op=op, name=name, source=source
)
@overload
def push_scope(
self, callback=None # type: Optional[None]
):
# type: (...) -> ContextManager[Scope]
pass
@overload
def push_scope( # noqa: F811
self, callback # type: Callable[[Scope], None]
):
# type: (...) -> None
pass
def push_scope( # noqa
self,
callback=None, # type: Optional[Callable[[Scope], None]]
continue_trace=True, # type: bool
):
# type: (...) -> Optional[ContextManager[Scope]]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Pushes a new layer on the scope stack.
:param callback: If provided, this method pushes a scope, calls
`callback`, and pops the scope again.
:returns: If no `callback` is provided, a context manager that should
be used to pop the scope again.
"""
if callback is not None:
with self.push_scope() as scope:
callback(scope)
return None
return _ScopeManager(self)
def pop_scope_unsafe(self):
# type: () -> Tuple[Optional[Client], Scope]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Pops a scope layer from the stack.
Try to use the context manager :py:meth:`push_scope` instead.
"""
rv = self._stack.pop()
assert self._stack, "stack must have at least one layer"
return rv
@overload
def configure_scope(
self, callback=None # type: Optional[None]
):
# type: (...) -> ContextManager[Scope]
pass
@overload
def configure_scope( # noqa: F811
self, callback # type: Callable[[Scope], None]
):
# type: (...) -> None
pass
def configure_scope( # noqa
self,
callback=None, # type: Optional[Callable[[Scope], None]]
continue_trace=True, # type: bool
):
# type: (...) -> Optional[ContextManager[Scope]]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Reconfigures the scope.
:param callback: If provided, call the callback with the current scope.
:returns: If no callback is provided, returns a context manager that returns the scope.
"""
scope = get_isolation_scope()
if continue_trace:
scope.generate_propagation_context()
if callback is not None:
# TODO: used to return None when client is None. Check if this changes behavior.
callback(scope)
return None
@contextmanager
def inner():
# type: () -> Generator[Scope, None, None]
yield scope
return inner()
def start_session(
self, session_mode="application" # type: str
):
# type: (...) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.start_session` instead.
Starts a new session.
"""
get_isolation_scope().start_session(
session_mode=session_mode,
)
def end_session(self):
# type: (...) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.end_session` instead.
Ends the current session if there is one.
"""
get_isolation_scope().end_session()
def stop_auto_session_tracking(self):
# type: (...) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.stop_auto_session_tracking` instead.
Stops automatic session tracking.
This temporarily session tracking for the current scope when called.
To resume session tracking call `resume_auto_session_tracking`.
"""
get_isolation_scope().stop_auto_session_tracking()
def resume_auto_session_tracking(self):
# type: (...) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.resume_auto_session_tracking` instead.
Resumes automatic session tracking for the current scope if
disabled earlier. This requires that generally automatic session
tracking is enabled.
"""
get_isolation_scope().resume_auto_session_tracking()
def flush(
self,
timeout=None, # type: Optional[float]
callback=None, # type: Optional[Callable[[int, float], None]]
):
# type: (...) -> None
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.client._Client.flush` instead.
Alias for :py:meth:`sentry_sdk.client._Client.flush`
"""
return get_client().flush(timeout=timeout, callback=callback)
def get_traceparent(self):
# type: () -> Optional[str]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.get_traceparent` instead.
Returns the traceparent either from the active span or from the scope.
"""
current_scope = get_current_scope()
traceparent = current_scope.get_traceparent()
if traceparent is None:
isolation_scope = get_isolation_scope()
traceparent = isolation_scope.get_traceparent()
return traceparent
def get_baggage(self):
# type: () -> Optional[str]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.get_baggage` instead.
Returns Baggage either from the active span or from the scope.
"""
current_scope = get_current_scope()
baggage = current_scope.get_baggage()
if baggage is None:
isolation_scope = get_isolation_scope()
baggage = isolation_scope.get_baggage()
if baggage is not None:
return baggage.serialize()
return None
def iter_trace_propagation_headers(self, span=None):
# type: (Optional[Span]) -> Generator[Tuple[str, str], None, None]
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.iter_trace_propagation_headers` instead.
Return HTTP headers which allow propagation of trace data. Data taken
from the span representing the request, if available, or the current
span on the scope if not.
"""
return get_current_scope().iter_trace_propagation_headers(
span=span,
)
def trace_propagation_meta(self, span=None):
# type: (Optional[Span]) -> str
"""
.. deprecated:: 2.0.0
This function is deprecated and will be removed in a future release.
Please use :py:meth:`sentry_sdk.Scope.trace_propagation_meta` instead.
Return meta tags which should be injected into HTML templates
to allow propagation of trace information.
"""
if span is not None:
logger.warning(
"The parameter `span` in trace_propagation_meta() is deprecated and will be removed in the future."
)
return get_current_scope().trace_propagation_meta(
span=span,
)
with _suppress_hub_deprecation_warning():
# Suppress deprecation warning for the Hub here, since we still always
# import this module.
GLOBAL_HUB = Hub()
_local.set(GLOBAL_HUB)
# Circular imports
from sentry_sdk import scope
@@ -0,0 +1,293 @@
from abc import ABC, abstractmethod
from threading import Lock
from sentry_sdk.utils import logger
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Sequence
from typing import Callable
from typing import Dict
from typing import Iterator
from typing import List
from typing import Optional
from typing import Set
from typing import Type
from typing import Union
_DEFAULT_FAILED_REQUEST_STATUS_CODES = frozenset(range(500, 600))
_installer_lock = Lock()
# Set of all integration identifiers we have attempted to install
_processed_integrations = set() # type: Set[str]
# Set of all integration identifiers we have actually installed
_installed_integrations = set() # type: Set[str]
def _generate_default_integrations_iterator(
integrations, # type: List[str]
auto_enabling_integrations, # type: List[str]
):
# type: (...) -> Callable[[bool], Iterator[Type[Integration]]]
def iter_default_integrations(with_auto_enabling_integrations):
# type: (bool) -> Iterator[Type[Integration]]
"""Returns an iterator of the default integration classes:"""
from importlib import import_module
if with_auto_enabling_integrations:
all_import_strings = integrations + auto_enabling_integrations
else:
all_import_strings = integrations
for import_string in all_import_strings:
try:
module, cls = import_string.rsplit(".", 1)
yield getattr(import_module(module), cls)
except (DidNotEnable, SyntaxError) as e:
logger.debug(
"Did not import default integration %s: %s", import_string, e
)
if isinstance(iter_default_integrations.__doc__, str):
for import_string in integrations:
iter_default_integrations.__doc__ += "\n- `{}`".format(import_string)
return iter_default_integrations
_DEFAULT_INTEGRATIONS = [
# stdlib/base runtime integrations
"sentry_sdk.integrations.argv.ArgvIntegration",
"sentry_sdk.integrations.atexit.AtexitIntegration",
"sentry_sdk.integrations.dedupe.DedupeIntegration",
"sentry_sdk.integrations.excepthook.ExcepthookIntegration",
"sentry_sdk.integrations.logging.LoggingIntegration",
"sentry_sdk.integrations.modules.ModulesIntegration",
"sentry_sdk.integrations.stdlib.StdlibIntegration",
"sentry_sdk.integrations.threading.ThreadingIntegration",
]
_AUTO_ENABLING_INTEGRATIONS = [
"sentry_sdk.integrations.aiohttp.AioHttpIntegration",
"sentry_sdk.integrations.anthropic.AnthropicIntegration",
"sentry_sdk.integrations.ariadne.AriadneIntegration",
"sentry_sdk.integrations.arq.ArqIntegration",
"sentry_sdk.integrations.asyncpg.AsyncPGIntegration",
"sentry_sdk.integrations.boto3.Boto3Integration",
"sentry_sdk.integrations.bottle.BottleIntegration",
"sentry_sdk.integrations.celery.CeleryIntegration",
"sentry_sdk.integrations.chalice.ChaliceIntegration",
"sentry_sdk.integrations.clickhouse_driver.ClickhouseDriverIntegration",
"sentry_sdk.integrations.cohere.CohereIntegration",
"sentry_sdk.integrations.django.DjangoIntegration",
"sentry_sdk.integrations.falcon.FalconIntegration",
"sentry_sdk.integrations.fastapi.FastApiIntegration",
"sentry_sdk.integrations.flask.FlaskIntegration",
"sentry_sdk.integrations.gql.GQLIntegration",
"sentry_sdk.integrations.graphene.GrapheneIntegration",
"sentry_sdk.integrations.httpx.HttpxIntegration",
"sentry_sdk.integrations.huey.HueyIntegration",
"sentry_sdk.integrations.huggingface_hub.HuggingfaceHubIntegration",
"sentry_sdk.integrations.langchain.LangchainIntegration",
"sentry_sdk.integrations.litestar.LitestarIntegration",
"sentry_sdk.integrations.loguru.LoguruIntegration",
"sentry_sdk.integrations.openai.OpenAIIntegration",
"sentry_sdk.integrations.pymongo.PyMongoIntegration",
"sentry_sdk.integrations.pyramid.PyramidIntegration",
"sentry_sdk.integrations.quart.QuartIntegration",
"sentry_sdk.integrations.redis.RedisIntegration",
"sentry_sdk.integrations.rq.RqIntegration",
"sentry_sdk.integrations.sanic.SanicIntegration",
"sentry_sdk.integrations.sqlalchemy.SqlalchemyIntegration",
"sentry_sdk.integrations.starlette.StarletteIntegration",
"sentry_sdk.integrations.starlite.StarliteIntegration",
"sentry_sdk.integrations.strawberry.StrawberryIntegration",
"sentry_sdk.integrations.tornado.TornadoIntegration",
]
iter_default_integrations = _generate_default_integrations_iterator(
integrations=_DEFAULT_INTEGRATIONS,
auto_enabling_integrations=_AUTO_ENABLING_INTEGRATIONS,
)
del _generate_default_integrations_iterator
_MIN_VERSIONS = {
"aiohttp": (3, 4),
"anthropic": (0, 16),
"ariadne": (0, 20),
"arq": (0, 23),
"asyncpg": (0, 23),
"beam": (2, 12),
"boto3": (1, 12), # botocore
"bottle": (0, 12),
"celery": (4, 4, 7),
"chalice": (1, 16, 0),
"clickhouse_driver": (0, 2, 0),
"django": (1, 8),
"dramatiq": (1, 9),
"falcon": (1, 4),
"fastapi": (0, 79, 0),
"flask": (1, 1, 4),
"gql": (3, 4, 1),
"graphene": (3, 3),
"grpc": (1, 32, 0), # grpcio
"huggingface_hub": (0, 22),
"langchain": (0, 0, 210),
"launchdarkly": (9, 8, 0),
"loguru": (0, 7, 0),
"openai": (1, 0, 0),
"openfeature": (0, 7, 1),
"quart": (0, 16, 0),
"ray": (2, 7, 0),
"requests": (2, 0, 0),
"rq": (0, 6),
"sanic": (0, 8),
"sqlalchemy": (1, 2),
"starlette": (0, 16),
"starlite": (1, 48),
"statsig": (0, 55, 3),
"strawberry": (0, 209, 5),
"tornado": (6, 0),
"typer": (0, 15),
"unleash": (6, 0, 1),
}
def setup_integrations(
integrations,
with_defaults=True,
with_auto_enabling_integrations=False,
disabled_integrations=None,
):
# type: (Sequence[Integration], bool, bool, Optional[Sequence[Union[type[Integration], Integration]]]) -> Dict[str, Integration]
"""
Given a list of integration instances, this installs them all.
When `with_defaults` is set to `True` all default integrations are added
unless they were already provided before.
`disabled_integrations` takes precedence over `with_defaults` and
`with_auto_enabling_integrations`.
"""
integrations = dict(
(integration.identifier, integration) for integration in integrations or ()
)
logger.debug("Setting up integrations (with default = %s)", with_defaults)
# Integrations that will not be enabled
disabled_integrations = [
integration if isinstance(integration, type) else type(integration)
for integration in disabled_integrations or []
]
# Integrations that are not explicitly set up by the user.
used_as_default_integration = set()
if with_defaults:
for integration_cls in iter_default_integrations(
with_auto_enabling_integrations
):
if integration_cls.identifier not in integrations:
instance = integration_cls()
integrations[instance.identifier] = instance
used_as_default_integration.add(instance.identifier)
for identifier, integration in integrations.items():
with _installer_lock:
if identifier not in _processed_integrations:
if type(integration) in disabled_integrations:
logger.debug("Ignoring integration %s", identifier)
else:
logger.debug(
"Setting up previously not enabled integration %s", identifier
)
try:
type(integration).setup_once()
except DidNotEnable as e:
if identifier not in used_as_default_integration:
raise
logger.debug(
"Did not enable default integration %s: %s", identifier, e
)
else:
_installed_integrations.add(identifier)
_processed_integrations.add(identifier)
integrations = {
identifier: integration
for identifier, integration in integrations.items()
if identifier in _installed_integrations
}
for identifier in integrations:
logger.debug("Enabling integration %s", identifier)
return integrations
def _check_minimum_version(integration, version, package=None):
# type: (type[Integration], Optional[tuple[int, ...]], Optional[str]) -> None
package = package or integration.identifier
if version is None:
raise DidNotEnable(f"Unparsable {package} version.")
min_version = _MIN_VERSIONS.get(integration.identifier)
if min_version is None:
return
if version < min_version:
raise DidNotEnable(
f"Integration only supports {package} {'.'.join(map(str, min_version))} or newer."
)
class DidNotEnable(Exception): # noqa: N818
"""
The integration could not be enabled due to a trivial user error like
`flask` not being installed for the `FlaskIntegration`.
This exception is silently swallowed for default integrations, but reraised
for explicitly enabled integrations.
"""
class Integration(ABC):
"""Baseclass for all integrations.
To accept options for an integration, implement your own constructor that
saves those options on `self`.
"""
install = None
"""Legacy method, do not implement."""
identifier = None # type: str
"""String unique ID of integration type"""
@staticmethod
@abstractmethod
def setup_once():
# type: () -> None
"""
Initialize the integration.
This function is only called once, ever. Configuration is not available
at this point, so the only thing to do here is to hook into exception
handlers, and perhaps do monkeypatches.
Inside those hooks `Integration.current` can be used to access the
instance again.
"""
pass
@@ -0,0 +1,108 @@
import urllib
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.integrations._wsgi_common import _filter_headers
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Dict
from typing import Optional
from typing import Union
from typing_extensions import Literal
from sentry_sdk.utils import AnnotatedValue
def _get_headers(asgi_scope):
# type: (Any) -> Dict[str, str]
"""
Extract headers from the ASGI scope, in the format that the Sentry protocol expects.
"""
headers = {} # type: Dict[str, str]
for raw_key, raw_value in asgi_scope["headers"]:
key = raw_key.decode("latin-1")
value = raw_value.decode("latin-1")
if key in headers:
headers[key] = headers[key] + ", " + value
else:
headers[key] = value
return headers
def _get_url(asgi_scope, default_scheme, host):
# type: (Dict[str, Any], Literal["ws", "http"], Optional[Union[AnnotatedValue, str]]) -> str
"""
Extract URL from the ASGI scope, without also including the querystring.
"""
scheme = asgi_scope.get("scheme", default_scheme)
server = asgi_scope.get("server", None)
path = asgi_scope.get("root_path", "") + asgi_scope.get("path", "")
if host:
return "%s://%s%s" % (scheme, host, path)
if server is not None:
host, port = server
default_port = {"http": 80, "https": 443, "ws": 80, "wss": 443}.get(scheme)
if port != default_port:
return "%s://%s:%s%s" % (scheme, host, port, path)
return "%s://%s%s" % (scheme, host, path)
return path
def _get_query(asgi_scope):
# type: (Any) -> Any
"""
Extract querystring from the ASGI scope, in the format that the Sentry protocol expects.
"""
qs = asgi_scope.get("query_string")
if not qs:
return None
return urllib.parse.unquote(qs.decode("latin-1"))
def _get_ip(asgi_scope):
# type: (Any) -> str
"""
Extract IP Address from the ASGI scope based on request headers with fallback to scope client.
"""
headers = _get_headers(asgi_scope)
try:
return headers["x-forwarded-for"].split(",")[0].strip()
except (KeyError, IndexError):
pass
try:
return headers["x-real-ip"]
except KeyError:
pass
return asgi_scope.get("client")[0]
def _get_request_data(asgi_scope):
# type: (Any) -> Dict[str, Any]
"""
Returns data related to the HTTP request from the ASGI scope.
"""
request_data = {} # type: Dict[str, Any]
ty = asgi_scope["type"]
if ty in ("http", "websocket"):
request_data["method"] = asgi_scope.get("method")
request_data["headers"] = headers = _filter_headers(_get_headers(asgi_scope))
request_data["query_string"] = _get_query(asgi_scope)
request_data["url"] = _get_url(
asgi_scope, "http" if ty == "http" else "ws", headers.get("host")
)
client = asgi_scope.get("client")
if client and should_send_default_pii():
request_data["env"] = {"REMOTE_ADDR": _get_ip(asgi_scope)}
return request_data
@@ -0,0 +1,271 @@
from contextlib import contextmanager
import json
from copy import deepcopy
import sentry_sdk
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import AnnotatedValue, logger
try:
from django.http.request import RawPostDataException
except ImportError:
RawPostDataException = None
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Dict
from typing import Iterator
from typing import Mapping
from typing import MutableMapping
from typing import Optional
from typing import Union
from sentry_sdk._types import Event, HttpStatusCodeRange
SENSITIVE_ENV_KEYS = (
"REMOTE_ADDR",
"HTTP_X_FORWARDED_FOR",
"HTTP_SET_COOKIE",
"HTTP_COOKIE",
"HTTP_AUTHORIZATION",
"HTTP_X_API_KEY",
"HTTP_X_FORWARDED_FOR",
"HTTP_X_REAL_IP",
)
SENSITIVE_HEADERS = tuple(
x[len("HTTP_") :] for x in SENSITIVE_ENV_KEYS if x.startswith("HTTP_")
)
DEFAULT_HTTP_METHODS_TO_CAPTURE = (
"CONNECT",
"DELETE",
"GET",
# "HEAD", # do not capture HEAD requests by default
# "OPTIONS", # do not capture OPTIONS requests by default
"PATCH",
"POST",
"PUT",
"TRACE",
)
# This noop context manager can be replaced with "from contextlib import nullcontext" when we drop Python 3.6 support
@contextmanager
def nullcontext():
# type: () -> Iterator[None]
yield
def request_body_within_bounds(client, content_length):
# type: (Optional[sentry_sdk.client.BaseClient], int) -> bool
if client is None:
return False
bodies = client.options["max_request_body_size"]
return not (
bodies == "never"
or (bodies == "small" and content_length > 10**3)
or (bodies == "medium" and content_length > 10**4)
)
class RequestExtractor:
"""
Base class for request extraction.
"""
# It does not make sense to make this class an ABC because it is not used
# for typing, only so that child classes can inherit common methods from
# it. Only some child classes implement all methods that raise
# NotImplementedError in this class.
def __init__(self, request):
# type: (Any) -> None
self.request = request
def extract_into_event(self, event):
# type: (Event) -> None
client = sentry_sdk.get_client()
if not client.is_active():
return
data = None # type: Optional[Union[AnnotatedValue, Dict[str, Any]]]
content_length = self.content_length()
request_info = event.get("request", {})
if should_send_default_pii():
request_info["cookies"] = dict(self.cookies())
if not request_body_within_bounds(client, content_length):
data = AnnotatedValue.removed_because_over_size_limit()
else:
# First read the raw body data
# It is important to read this first because if it is Django
# it will cache the body and then we can read the cached version
# again in parsed_body() (or json() or wherever).
raw_data = None
try:
raw_data = self.raw_data()
except (RawPostDataException, ValueError):
# If DjangoRestFramework is used it already read the body for us
# so reading it here will fail. We can ignore this.
pass
parsed_body = self.parsed_body()
if parsed_body is not None:
data = parsed_body
elif raw_data:
data = AnnotatedValue.removed_because_raw_data()
else:
data = None
if data is not None:
request_info["data"] = data
event["request"] = deepcopy(request_info)
def content_length(self):
# type: () -> int
try:
return int(self.env().get("CONTENT_LENGTH", 0))
except ValueError:
return 0
def cookies(self):
# type: () -> MutableMapping[str, Any]
raise NotImplementedError()
def raw_data(self):
# type: () -> Optional[Union[str, bytes]]
raise NotImplementedError()
def form(self):
# type: () -> Optional[Dict[str, Any]]
raise NotImplementedError()
def parsed_body(self):
# type: () -> Optional[Dict[str, Any]]
try:
form = self.form()
except Exception:
form = None
try:
files = self.files()
except Exception:
files = None
if form or files:
data = {}
if form:
data = dict(form.items())
if files:
for key in files.keys():
data[key] = AnnotatedValue.removed_because_raw_data()
return data
return self.json()
def is_json(self):
# type: () -> bool
return _is_json_content_type(self.env().get("CONTENT_TYPE"))
def json(self):
# type: () -> Optional[Any]
try:
if not self.is_json():
return None
try:
raw_data = self.raw_data()
except (RawPostDataException, ValueError):
# The body might have already been read, in which case this will
# fail
raw_data = None
if raw_data is None:
return None
if isinstance(raw_data, str):
return json.loads(raw_data)
else:
return json.loads(raw_data.decode("utf-8"))
except ValueError:
pass
return None
def files(self):
# type: () -> Optional[Dict[str, Any]]
raise NotImplementedError()
def size_of_file(self, file):
# type: (Any) -> int
raise NotImplementedError()
def env(self):
# type: () -> Dict[str, Any]
raise NotImplementedError()
def _is_json_content_type(ct):
# type: (Optional[str]) -> bool
mt = (ct or "").split(";", 1)[0]
return (
mt == "application/json"
or (mt.startswith("application/"))
and mt.endswith("+json")
)
def _filter_headers(headers):
# type: (Mapping[str, str]) -> Mapping[str, Union[AnnotatedValue, str]]
if should_send_default_pii():
return headers
return {
k: (
v
if k.upper().replace("-", "_") not in SENSITIVE_HEADERS
else AnnotatedValue.removed_because_over_size_limit()
)
for k, v in headers.items()
}
def _in_http_status_code_range(code, code_ranges):
# type: (object, list[HttpStatusCodeRange]) -> bool
for target in code_ranges:
if isinstance(target, int):
if code == target:
return True
continue
try:
if code in target:
return True
except TypeError:
logger.warning(
"failed_request_status_codes has to be a list of integers or containers"
)
return False
class HttpCodeRangeContainer:
"""
Wrapper to make it possible to use list[HttpStatusCodeRange] as a Container[int].
Used for backwards compatibility with the old `failed_request_status_codes` option.
"""
def __init__(self, code_ranges):
# type: (list[HttpStatusCodeRange]) -> None
self._code_ranges = code_ranges
def __contains__(self, item):
# type: (object) -> bool
return _in_http_status_code_range(item, self._code_ranges)
@@ -0,0 +1,357 @@
import sys
import weakref
from functools import wraps
import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
from sentry_sdk.integrations import (
_DEFAULT_FAILED_REQUEST_STATUS_CODES,
_check_minimum_version,
Integration,
DidNotEnable,
)
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.sessions import track_session
from sentry_sdk.integrations._wsgi_common import (
_filter_headers,
request_body_within_bounds,
)
from sentry_sdk.tracing import (
BAGGAGE_HEADER_NAME,
SOURCE_FOR_STYLE,
TRANSACTION_SOURCE_ROUTE,
)
from sentry_sdk.tracing_utils import should_propagate_trace
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
logger,
parse_url,
parse_version,
reraise,
transaction_from_function,
HAS_REAL_CONTEXTVARS,
CONTEXTVARS_ERROR_MESSAGE,
SENSITIVE_DATA_SUBSTITUTE,
AnnotatedValue,
)
try:
import asyncio
from aiohttp import __version__ as AIOHTTP_VERSION
from aiohttp import ClientSession, TraceConfig
from aiohttp.web import Application, HTTPException, UrlDispatcher
except ImportError:
raise DidNotEnable("AIOHTTP not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from aiohttp.web_request import Request
from aiohttp.web_urldispatcher import UrlMappingMatchInfo
from aiohttp import TraceRequestStartParams, TraceRequestEndParams
from collections.abc import Set
from types import SimpleNamespace
from typing import Any
from typing import Optional
from typing import Tuple
from typing import Union
from sentry_sdk.utils import ExcInfo
from sentry_sdk._types import Event, EventProcessor
TRANSACTION_STYLE_VALUES = ("handler_name", "method_and_path_pattern")
class AioHttpIntegration(Integration):
identifier = "aiohttp"
origin = f"auto.http.{identifier}"
def __init__(
self,
transaction_style="handler_name", # type: str
*,
failed_request_status_codes=_DEFAULT_FAILED_REQUEST_STATUS_CODES, # type: Set[int]
):
# type: (...) -> None
if transaction_style not in TRANSACTION_STYLE_VALUES:
raise ValueError(
"Invalid value for transaction_style: %s (must be in %s)"
% (transaction_style, TRANSACTION_STYLE_VALUES)
)
self.transaction_style = transaction_style
self._failed_request_status_codes = failed_request_status_codes
@staticmethod
def setup_once():
# type: () -> None
version = parse_version(AIOHTTP_VERSION)
_check_minimum_version(AioHttpIntegration, version)
if not HAS_REAL_CONTEXTVARS:
# We better have contextvars or we're going to leak state between
# requests.
raise DidNotEnable(
"The aiohttp integration for Sentry requires Python 3.7+ "
" or aiocontextvars package." + CONTEXTVARS_ERROR_MESSAGE
)
ignore_logger("aiohttp.server")
old_handle = Application._handle
async def sentry_app_handle(self, request, *args, **kwargs):
# type: (Any, Request, *Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(AioHttpIntegration)
if integration is None:
return await old_handle(self, request, *args, **kwargs)
weak_request = weakref.ref(request)
with sentry_sdk.isolation_scope() as scope:
with track_session(scope, session_mode="request"):
# Scope data will not leak between requests because aiohttp
# create a task to wrap each request.
scope.generate_propagation_context()
scope.clear_breadcrumbs()
scope.add_event_processor(_make_request_processor(weak_request))
headers = dict(request.headers)
transaction = continue_trace(
headers,
op=OP.HTTP_SERVER,
# If this transaction name makes it to the UI, AIOHTTP's
# URL resolver did not find a route or died trying.
name="generic AIOHTTP request",
source=TRANSACTION_SOURCE_ROUTE,
origin=AioHttpIntegration.origin,
)
with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={"aiohttp_request": request},
):
try:
response = await old_handle(self, request)
except HTTPException as e:
transaction.set_http_status(e.status_code)
if (
e.status_code
in integration._failed_request_status_codes
):
_capture_exception()
raise
except (asyncio.CancelledError, ConnectionResetError):
transaction.set_status(SPANSTATUS.CANCELLED)
raise
except Exception:
# This will probably map to a 500 but seems like we
# have no way to tell. Do not set span status.
reraise(*_capture_exception())
try:
# A valid response handler will return a valid response with a status. But, if the handler
# returns an invalid response (e.g. None), the line below will raise an AttributeError.
# Even though this is likely invalid, we need to handle this case to ensure we don't break
# the application.
response_status = response.status
except AttributeError:
pass
else:
transaction.set_http_status(response_status)
return response
Application._handle = sentry_app_handle
old_urldispatcher_resolve = UrlDispatcher.resolve
@wraps(old_urldispatcher_resolve)
async def sentry_urldispatcher_resolve(self, request):
# type: (UrlDispatcher, Request) -> UrlMappingMatchInfo
rv = await old_urldispatcher_resolve(self, request)
integration = sentry_sdk.get_client().get_integration(AioHttpIntegration)
if integration is None:
return rv
name = None
try:
if integration.transaction_style == "handler_name":
name = transaction_from_function(rv.handler)
elif integration.transaction_style == "method_and_path_pattern":
route_info = rv.get_info()
pattern = route_info.get("path") or route_info.get("formatter")
name = "{} {}".format(request.method, pattern)
except Exception:
pass
if name is not None:
sentry_sdk.get_current_scope().set_transaction_name(
name,
source=SOURCE_FOR_STYLE[integration.transaction_style],
)
return rv
UrlDispatcher.resolve = sentry_urldispatcher_resolve
old_client_session_init = ClientSession.__init__
@ensure_integration_enabled(AioHttpIntegration, old_client_session_init)
def init(*args, **kwargs):
# type: (Any, Any) -> None
client_trace_configs = list(kwargs.get("trace_configs") or ())
trace_config = create_trace_config()
client_trace_configs.append(trace_config)
kwargs["trace_configs"] = client_trace_configs
return old_client_session_init(*args, **kwargs)
ClientSession.__init__ = init
def create_trace_config():
# type: () -> TraceConfig
async def on_request_start(session, trace_config_ctx, params):
# type: (ClientSession, SimpleNamespace, TraceRequestStartParams) -> None
if sentry_sdk.get_client().get_integration(AioHttpIntegration) is None:
return
method = params.method.upper()
parsed_url = None
with capture_internal_exceptions():
parsed_url = parse_url(str(params.url), sanitize=False)
span = sentry_sdk.start_span(
op=OP.HTTP_CLIENT,
name="%s %s"
% (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE),
origin=AioHttpIntegration.origin,
)
span.set_data(SPANDATA.HTTP_METHOD, method)
if parsed_url is not None:
span.set_data("url", parsed_url.url)
span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
client = sentry_sdk.get_client()
if should_propagate_trace(client, str(params.url)):
for (
key,
value,
) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(
span=span
):
logger.debug(
"[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format(
key=key, value=value, url=params.url
)
)
if key == BAGGAGE_HEADER_NAME and params.headers.get(
BAGGAGE_HEADER_NAME
):
# do not overwrite any existing baggage, just append to it
params.headers[key] += "," + value
else:
params.headers[key] = value
trace_config_ctx.span = span
async def on_request_end(session, trace_config_ctx, params):
# type: (ClientSession, SimpleNamespace, TraceRequestEndParams) -> None
if trace_config_ctx.span is None:
return
span = trace_config_ctx.span
span.set_http_status(int(params.response.status))
span.set_data("reason", params.response.reason)
span.finish()
trace_config = TraceConfig()
trace_config.on_request_start.append(on_request_start)
trace_config.on_request_end.append(on_request_end)
return trace_config
def _make_request_processor(weak_request):
# type: (weakref.ReferenceType[Request]) -> EventProcessor
def aiohttp_processor(
event, # type: Event
hint, # type: dict[str, Tuple[type, BaseException, Any]]
):
# type: (...) -> Event
request = weak_request()
if request is None:
return event
with capture_internal_exceptions():
request_info = event.setdefault("request", {})
request_info["url"] = "%s://%s%s" % (
request.scheme,
request.host,
request.path,
)
request_info["query_string"] = request.query_string
request_info["method"] = request.method
request_info["env"] = {"REMOTE_ADDR": request.remote}
request_info["headers"] = _filter_headers(dict(request.headers))
# Just attach raw data here if it is within bounds, if available.
# Unfortunately there's no way to get structured data from aiohttp
# without awaiting on some coroutine.
request_info["data"] = get_aiohttp_request_data(request)
return event
return aiohttp_processor
def _capture_exception():
# type: () -> ExcInfo
exc_info = sys.exc_info()
event, hint = event_from_exception(
exc_info,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "aiohttp", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
return exc_info
BODY_NOT_READ_MESSAGE = "[Can't show request body due to implementation details.]"
def get_aiohttp_request_data(request):
# type: (Request) -> Union[Optional[str], AnnotatedValue]
bytes_body = request._read_bytes
if bytes_body is not None:
# we have body to show
if not request_body_within_bounds(sentry_sdk.get_client(), len(bytes_body)):
return AnnotatedValue.removed_because_over_size_limit()
encoding = request.charset or "utf-8"
return bytes_body.decode(encoding, "replace")
if request.can_read_body:
# body exists but we can't show it
return BODY_NOT_READ_MESSAGE
# request has no body
return None
@@ -0,0 +1,286 @@
from functools import wraps
from typing import TYPE_CHECKING
import sentry_sdk
from sentry_sdk.ai.monitoring import record_token_usage
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import (
capture_internal_exceptions,
event_from_exception,
package_version,
)
try:
from anthropic.resources import AsyncMessages, Messages
if TYPE_CHECKING:
from anthropic.types import MessageStreamEvent
except ImportError:
raise DidNotEnable("Anthropic not installed")
if TYPE_CHECKING:
from typing import Any, AsyncIterator, Iterator
from sentry_sdk.tracing import Span
class AnthropicIntegration(Integration):
identifier = "anthropic"
origin = f"auto.ai.{identifier}"
def __init__(self, include_prompts=True):
# type: (AnthropicIntegration, bool) -> None
self.include_prompts = include_prompts
@staticmethod
def setup_once():
# type: () -> None
version = package_version("anthropic")
_check_minimum_version(AnthropicIntegration, version)
Messages.create = _wrap_message_create(Messages.create)
AsyncMessages.create = _wrap_message_create_async(AsyncMessages.create)
def _capture_exception(exc):
# type: (Any) -> None
event, hint = event_from_exception(
exc,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "anthropic", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _calculate_token_usage(result, span):
# type: (Messages, Span) -> None
input_tokens = 0
output_tokens = 0
if hasattr(result, "usage"):
usage = result.usage
if hasattr(usage, "input_tokens") and isinstance(usage.input_tokens, int):
input_tokens = usage.input_tokens
if hasattr(usage, "output_tokens") and isinstance(usage.output_tokens, int):
output_tokens = usage.output_tokens
total_tokens = input_tokens + output_tokens
record_token_usage(span, input_tokens, output_tokens, total_tokens)
def _get_responses(content):
# type: (list[Any]) -> list[dict[str, Any]]
"""
Get JSON of a Anthropic responses.
"""
responses = []
for item in content:
if hasattr(item, "text"):
responses.append(
{
"type": item.type,
"text": item.text,
}
)
return responses
def _collect_ai_data(event, input_tokens, output_tokens, content_blocks):
# type: (MessageStreamEvent, int, int, list[str]) -> tuple[int, int, list[str]]
"""
Count token usage and collect content blocks from the AI streaming response.
"""
with capture_internal_exceptions():
if hasattr(event, "type"):
if event.type == "message_start":
usage = event.message.usage
input_tokens += usage.input_tokens
output_tokens += usage.output_tokens
elif event.type == "content_block_start":
pass
elif event.type == "content_block_delta":
if hasattr(event.delta, "text"):
content_blocks.append(event.delta.text)
elif event.type == "content_block_stop":
pass
elif event.type == "message_delta":
output_tokens += event.usage.output_tokens
return input_tokens, output_tokens, content_blocks
def _add_ai_data_to_span(
span, integration, input_tokens, output_tokens, content_blocks
):
# type: (Span, AnthropicIntegration, int, int, list[str]) -> None
"""
Add token usage and content blocks from the AI streaming response to the span.
"""
with capture_internal_exceptions():
if should_send_default_pii() and integration.include_prompts:
complete_message = "".join(content_blocks)
span.set_data(
SPANDATA.AI_RESPONSES,
[{"type": "text", "text": complete_message}],
)
total_tokens = input_tokens + output_tokens
record_token_usage(span, input_tokens, output_tokens, total_tokens)
span.set_data(SPANDATA.AI_STREAMING, True)
def _sentry_patched_create_common(f, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
integration = kwargs.pop("integration")
if integration is None:
return f(*args, **kwargs)
if "messages" not in kwargs:
return f(*args, **kwargs)
try:
iter(kwargs["messages"])
except TypeError:
return f(*args, **kwargs)
span = sentry_sdk.start_span(
op=OP.ANTHROPIC_MESSAGES_CREATE,
description="Anthropic messages create",
origin=AnthropicIntegration.origin,
)
span.__enter__()
result = yield f, args, kwargs
# add data to span and finish it
messages = list(kwargs["messages"])
model = kwargs.get("model")
with capture_internal_exceptions():
span.set_data(SPANDATA.AI_MODEL_ID, model)
span.set_data(SPANDATA.AI_STREAMING, False)
if should_send_default_pii() and integration.include_prompts:
span.set_data(SPANDATA.AI_INPUT_MESSAGES, messages)
if hasattr(result, "content"):
if should_send_default_pii() and integration.include_prompts:
span.set_data(SPANDATA.AI_RESPONSES, _get_responses(result.content))
_calculate_token_usage(result, span)
span.__exit__(None, None, None)
# Streaming response
elif hasattr(result, "_iterator"):
old_iterator = result._iterator
def new_iterator():
# type: () -> Iterator[MessageStreamEvent]
input_tokens = 0
output_tokens = 0
content_blocks = [] # type: list[str]
for event in old_iterator:
input_tokens, output_tokens, content_blocks = _collect_ai_data(
event, input_tokens, output_tokens, content_blocks
)
if event.type != "message_stop":
yield event
_add_ai_data_to_span(
span, integration, input_tokens, output_tokens, content_blocks
)
span.__exit__(None, None, None)
async def new_iterator_async():
# type: () -> AsyncIterator[MessageStreamEvent]
input_tokens = 0
output_tokens = 0
content_blocks = [] # type: list[str]
async for event in old_iterator:
input_tokens, output_tokens, content_blocks = _collect_ai_data(
event, input_tokens, output_tokens, content_blocks
)
if event.type != "message_stop":
yield event
_add_ai_data_to_span(
span, integration, input_tokens, output_tokens, content_blocks
)
span.__exit__(None, None, None)
if str(type(result._iterator)) == "<class 'async_generator'>":
result._iterator = new_iterator_async()
else:
result._iterator = new_iterator()
else:
span.set_data("unknown_response", True)
span.__exit__(None, None, None)
return result
def _wrap_message_create(f):
# type: (Any) -> Any
def _execute_sync(f, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
gen = _sentry_patched_create_common(f, *args, **kwargs)
try:
f, args, kwargs = next(gen)
except StopIteration as e:
return e.value
try:
try:
result = f(*args, **kwargs)
except Exception as exc:
_capture_exception(exc)
raise exc from None
return gen.send(result)
except StopIteration as e:
return e.value
@wraps(f)
def _sentry_patched_create_sync(*args, **kwargs):
# type: (*Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
kwargs["integration"] = integration
return _execute_sync(f, *args, **kwargs)
return _sentry_patched_create_sync
def _wrap_message_create_async(f):
# type: (Any) -> Any
async def _execute_async(f, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
gen = _sentry_patched_create_common(f, *args, **kwargs)
try:
f, args, kwargs = next(gen)
except StopIteration as e:
return await e.value
try:
try:
result = await f(*args, **kwargs)
except Exception as exc:
_capture_exception(exc)
raise exc from None
return gen.send(result)
except StopIteration as e:
return e.value
@wraps(f)
async def _sentry_patched_create_async(*args, **kwargs):
# type: (*Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
kwargs["integration"] = integration
return await _execute_async(f, *args, **kwargs)
return _sentry_patched_create_async
@@ -0,0 +1,31 @@
import sys
import sentry_sdk
from sentry_sdk.integrations import Integration
from sentry_sdk.scope import add_global_event_processor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional
from sentry_sdk._types import Event, Hint
class ArgvIntegration(Integration):
identifier = "argv"
@staticmethod
def setup_once():
# type: () -> None
@add_global_event_processor
def processor(event, hint):
# type: (Event, Optional[Hint]) -> Optional[Event]
if sentry_sdk.get_client().get_integration(ArgvIntegration) is not None:
extra = event.setdefault("extra", {})
# If some event processor decided to set extra to e.g. an
# `int`, don't crash. Not here.
if isinstance(extra, dict):
extra["sys.argv"] = sys.argv
return event
@@ -0,0 +1,161 @@
from importlib import import_module
import sentry_sdk
from sentry_sdk import get_client, capture_event
from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.integrations._wsgi_common import request_body_within_bounds
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
package_version,
)
try:
# importing like this is necessary due to name shadowing in ariadne
# (ariadne.graphql is also a function)
ariadne_graphql = import_module("ariadne.graphql")
except ImportError:
raise DidNotEnable("ariadne is not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Dict, List, Optional
from ariadne.types import GraphQLError, GraphQLResult, GraphQLSchema, QueryParser # type: ignore
from graphql.language.ast import DocumentNode
from sentry_sdk._types import Event, EventProcessor
class AriadneIntegration(Integration):
identifier = "ariadne"
@staticmethod
def setup_once():
# type: () -> None
version = package_version("ariadne")
_check_minimum_version(AriadneIntegration, version)
ignore_logger("ariadne")
_patch_graphql()
def _patch_graphql():
# type: () -> None
old_parse_query = ariadne_graphql.parse_query
old_handle_errors = ariadne_graphql.handle_graphql_errors
old_handle_query_result = ariadne_graphql.handle_query_result
@ensure_integration_enabled(AriadneIntegration, old_parse_query)
def _sentry_patched_parse_query(context_value, query_parser, data):
# type: (Optional[Any], Optional[QueryParser], Any) -> DocumentNode
event_processor = _make_request_event_processor(data)
sentry_sdk.get_isolation_scope().add_event_processor(event_processor)
result = old_parse_query(context_value, query_parser, data)
return result
@ensure_integration_enabled(AriadneIntegration, old_handle_errors)
def _sentry_patched_handle_graphql_errors(errors, *args, **kwargs):
# type: (List[GraphQLError], Any, Any) -> GraphQLResult
result = old_handle_errors(errors, *args, **kwargs)
event_processor = _make_response_event_processor(result[1])
sentry_sdk.get_isolation_scope().add_event_processor(event_processor)
client = get_client()
if client.is_active():
with capture_internal_exceptions():
for error in errors:
event, hint = event_from_exception(
error,
client_options=client.options,
mechanism={
"type": AriadneIntegration.identifier,
"handled": False,
},
)
capture_event(event, hint=hint)
return result
@ensure_integration_enabled(AriadneIntegration, old_handle_query_result)
def _sentry_patched_handle_query_result(result, *args, **kwargs):
# type: (Any, Any, Any) -> GraphQLResult
query_result = old_handle_query_result(result, *args, **kwargs)
event_processor = _make_response_event_processor(query_result[1])
sentry_sdk.get_isolation_scope().add_event_processor(event_processor)
client = get_client()
if client.is_active():
with capture_internal_exceptions():
for error in result.errors or []:
event, hint = event_from_exception(
error,
client_options=client.options,
mechanism={
"type": AriadneIntegration.identifier,
"handled": False,
},
)
capture_event(event, hint=hint)
return query_result
ariadne_graphql.parse_query = _sentry_patched_parse_query # type: ignore
ariadne_graphql.handle_graphql_errors = _sentry_patched_handle_graphql_errors # type: ignore
ariadne_graphql.handle_query_result = _sentry_patched_handle_query_result # type: ignore
def _make_request_event_processor(data):
# type: (GraphQLSchema) -> EventProcessor
"""Add request data and api_target to events."""
def inner(event, hint):
# type: (Event, dict[str, Any]) -> Event
if not isinstance(data, dict):
return event
with capture_internal_exceptions():
try:
content_length = int(
(data.get("headers") or {}).get("Content-Length", 0)
)
except (TypeError, ValueError):
return event
if should_send_default_pii() and request_body_within_bounds(
get_client(), content_length
):
request_info = event.setdefault("request", {})
request_info["api_target"] = "graphql"
request_info["data"] = data
elif event.get("request", {}).get("data"):
del event["request"]["data"]
return event
return inner
def _make_response_event_processor(response):
# type: (Dict[str, Any]) -> EventProcessor
"""Add response data to the event's response context."""
def inner(event, hint):
# type: (Event, dict[str, Any]) -> Event
with capture_internal_exceptions():
if should_send_default_pii() and response.get("errors"):
contexts = event.setdefault("contexts", {})
contexts["response"] = {
"data": response,
}
return event
return inner
@@ -0,0 +1,245 @@
import sys
import sentry_sdk
from sentry_sdk.consts import OP, SPANSTATUS
from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_TASK
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
SENSITIVE_DATA_SUBSTITUTE,
parse_version,
reraise,
)
try:
import arq.worker
from arq.version import VERSION as ARQ_VERSION
from arq.connections import ArqRedis
from arq.worker import JobExecutionFailed, Retry, RetryJob, Worker
except ImportError:
raise DidNotEnable("Arq is not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Dict, Optional, Union
from sentry_sdk._types import EventProcessor, Event, ExcInfo, Hint
from arq.cron import CronJob
from arq.jobs import Job
from arq.typing import WorkerCoroutine
from arq.worker import Function
ARQ_CONTROL_FLOW_EXCEPTIONS = (JobExecutionFailed, Retry, RetryJob)
class ArqIntegration(Integration):
identifier = "arq"
origin = f"auto.queue.{identifier}"
@staticmethod
def setup_once():
# type: () -> None
try:
if isinstance(ARQ_VERSION, str):
version = parse_version(ARQ_VERSION)
else:
version = ARQ_VERSION.version[:2]
except (TypeError, ValueError):
version = None
_check_minimum_version(ArqIntegration, version)
patch_enqueue_job()
patch_run_job()
patch_create_worker()
ignore_logger("arq.worker")
def patch_enqueue_job():
# type: () -> None
old_enqueue_job = ArqRedis.enqueue_job
original_kwdefaults = old_enqueue_job.__kwdefaults__
async def _sentry_enqueue_job(self, function, *args, **kwargs):
# type: (ArqRedis, str, *Any, **Any) -> Optional[Job]
integration = sentry_sdk.get_client().get_integration(ArqIntegration)
if integration is None:
return await old_enqueue_job(self, function, *args, **kwargs)
with sentry_sdk.start_span(
op=OP.QUEUE_SUBMIT_ARQ, name=function, origin=ArqIntegration.origin
):
return await old_enqueue_job(self, function, *args, **kwargs)
_sentry_enqueue_job.__kwdefaults__ = original_kwdefaults
ArqRedis.enqueue_job = _sentry_enqueue_job
def patch_run_job():
# type: () -> None
old_run_job = Worker.run_job
async def _sentry_run_job(self, job_id, score):
# type: (Worker, str, int) -> None
integration = sentry_sdk.get_client().get_integration(ArqIntegration)
if integration is None:
return await old_run_job(self, job_id, score)
with sentry_sdk.isolation_scope() as scope:
scope._name = "arq"
scope.clear_breadcrumbs()
transaction = Transaction(
name="unknown arq task",
status="ok",
op=OP.QUEUE_TASK_ARQ,
source=TRANSACTION_SOURCE_TASK,
origin=ArqIntegration.origin,
)
with sentry_sdk.start_transaction(transaction):
return await old_run_job(self, job_id, score)
Worker.run_job = _sentry_run_job
def _capture_exception(exc_info):
# type: (ExcInfo) -> None
scope = sentry_sdk.get_current_scope()
if scope.transaction is not None:
if exc_info[0] in ARQ_CONTROL_FLOW_EXCEPTIONS:
scope.transaction.set_status(SPANSTATUS.ABORTED)
return
scope.transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
event, hint = event_from_exception(
exc_info,
client_options=sentry_sdk.get_client().options,
mechanism={"type": ArqIntegration.identifier, "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _make_event_processor(ctx, *args, **kwargs):
# type: (Dict[Any, Any], *Any, **Any) -> EventProcessor
def event_processor(event, hint):
# type: (Event, Hint) -> Optional[Event]
with capture_internal_exceptions():
scope = sentry_sdk.get_current_scope()
if scope.transaction is not None:
scope.transaction.name = ctx["job_name"]
event["transaction"] = ctx["job_name"]
tags = event.setdefault("tags", {})
tags["arq_task_id"] = ctx["job_id"]
tags["arq_task_retry"] = ctx["job_try"] > 1
extra = event.setdefault("extra", {})
extra["arq-job"] = {
"task": ctx["job_name"],
"args": (
args if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
),
"kwargs": (
kwargs if should_send_default_pii() else SENSITIVE_DATA_SUBSTITUTE
),
"retry": ctx["job_try"],
}
return event
return event_processor
def _wrap_coroutine(name, coroutine):
# type: (str, WorkerCoroutine) -> WorkerCoroutine
async def _sentry_coroutine(ctx, *args, **kwargs):
# type: (Dict[Any, Any], *Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(ArqIntegration)
if integration is None:
return await coroutine(ctx, *args, **kwargs)
sentry_sdk.get_isolation_scope().add_event_processor(
_make_event_processor({**ctx, "job_name": name}, *args, **kwargs)
)
try:
result = await coroutine(ctx, *args, **kwargs)
except Exception:
exc_info = sys.exc_info()
_capture_exception(exc_info)
reraise(*exc_info)
return result
return _sentry_coroutine
def patch_create_worker():
# type: () -> None
old_create_worker = arq.worker.create_worker
@ensure_integration_enabled(ArqIntegration, old_create_worker)
def _sentry_create_worker(*args, **kwargs):
# type: (*Any, **Any) -> Worker
settings_cls = args[0]
if isinstance(settings_cls, dict):
if "functions" in settings_cls:
settings_cls["functions"] = [
_get_arq_function(func) for func in settings_cls["functions"]
]
if "cron_jobs" in settings_cls:
settings_cls["cron_jobs"] = [
_get_arq_cron_job(cron_job)
for cron_job in settings_cls["cron_jobs"]
]
if hasattr(settings_cls, "functions"):
settings_cls.functions = [
_get_arq_function(func) for func in settings_cls.functions
]
if hasattr(settings_cls, "cron_jobs"):
settings_cls.cron_jobs = [
_get_arq_cron_job(cron_job) for cron_job in settings_cls.cron_jobs
]
if "functions" in kwargs:
kwargs["functions"] = [
_get_arq_function(func) for func in kwargs["functions"]
]
if "cron_jobs" in kwargs:
kwargs["cron_jobs"] = [
_get_arq_cron_job(cron_job) for cron_job in kwargs["cron_jobs"]
]
return old_create_worker(*args, **kwargs)
arq.worker.create_worker = _sentry_create_worker
def _get_arq_function(func):
# type: (Union[str, Function, WorkerCoroutine]) -> Function
arq_func = arq.worker.func(func)
arq_func.coroutine = _wrap_coroutine(arq_func.name, arq_func.coroutine)
return arq_func
def _get_arq_cron_job(cron_job):
# type: (CronJob) -> CronJob
cron_job.coroutine = _wrap_coroutine(cron_job.name, cron_job.coroutine)
return cron_job
@@ -0,0 +1,335 @@
"""
An ASGI middleware.
Based on Tom Christie's `sentry-asgi <https://github.com/encode/sentry-asgi>`.
"""
import asyncio
import inspect
from copy import deepcopy
from functools import partial
import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP
from sentry_sdk.integrations._asgi_common import (
_get_headers,
_get_request_data,
_get_url,
)
from sentry_sdk.integrations._wsgi_common import (
DEFAULT_HTTP_METHODS_TO_CAPTURE,
nullcontext,
)
from sentry_sdk.sessions import track_session
from sentry_sdk.tracing import (
SOURCE_FOR_STYLE,
TRANSACTION_SOURCE_ROUTE,
TRANSACTION_SOURCE_URL,
TRANSACTION_SOURCE_COMPONENT,
TRANSACTION_SOURCE_CUSTOM,
)
from sentry_sdk.utils import (
ContextVar,
event_from_exception,
HAS_REAL_CONTEXTVARS,
CONTEXTVARS_ERROR_MESSAGE,
logger,
transaction_from_function,
_get_installed_modules,
)
from sentry_sdk.tracing import Transaction
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import Dict
from typing import Optional
from typing import Tuple
from sentry_sdk._types import Event, Hint
_asgi_middleware_applied = ContextVar("sentry_asgi_middleware_applied")
_DEFAULT_TRANSACTION_NAME = "generic ASGI request"
TRANSACTION_STYLE_VALUES = ("endpoint", "url")
def _capture_exception(exc, mechanism_type="asgi"):
# type: (Any, str) -> None
event, hint = event_from_exception(
exc,
client_options=sentry_sdk.get_client().options,
mechanism={"type": mechanism_type, "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _looks_like_asgi3(app):
# type: (Any) -> bool
"""
Try to figure out if an application object supports ASGI3.
This is how uvicorn figures out the application version as well.
"""
if inspect.isclass(app):
return hasattr(app, "__await__")
elif inspect.isfunction(app):
return asyncio.iscoroutinefunction(app)
else:
call = getattr(app, "__call__", None) # noqa
return asyncio.iscoroutinefunction(call)
class SentryAsgiMiddleware:
__slots__ = (
"app",
"__call__",
"transaction_style",
"mechanism_type",
"span_origin",
"http_methods_to_capture",
)
def __init__(
self,
app, # type: Any
unsafe_context_data=False, # type: bool
transaction_style="endpoint", # type: str
mechanism_type="asgi", # type: str
span_origin="manual", # type: str
http_methods_to_capture=DEFAULT_HTTP_METHODS_TO_CAPTURE, # type: Tuple[str, ...]
):
# type: (...) -> None
"""
Instrument an ASGI application with Sentry. Provides HTTP/websocket
data to sent events and basic handling for exceptions bubbling up
through the middleware.
:param unsafe_context_data: Disable errors when a proper contextvars installation could not be found. We do not recommend changing this from the default.
"""
if not unsafe_context_data and not HAS_REAL_CONTEXTVARS:
# We better have contextvars or we're going to leak state between
# requests.
raise RuntimeError(
"The ASGI middleware for Sentry requires Python 3.7+ "
"or the aiocontextvars package." + CONTEXTVARS_ERROR_MESSAGE
)
if transaction_style not in TRANSACTION_STYLE_VALUES:
raise ValueError(
"Invalid value for transaction_style: %s (must be in %s)"
% (transaction_style, TRANSACTION_STYLE_VALUES)
)
asgi_middleware_while_using_starlette_or_fastapi = (
mechanism_type == "asgi" and "starlette" in _get_installed_modules()
)
if asgi_middleware_while_using_starlette_or_fastapi:
logger.warning(
"The Sentry Python SDK can now automatically support ASGI frameworks like Starlette and FastAPI. "
"Please remove 'SentryAsgiMiddleware' from your project. "
"See https://docs.sentry.io/platforms/python/guides/asgi/ for more information."
)
self.transaction_style = transaction_style
self.mechanism_type = mechanism_type
self.span_origin = span_origin
self.app = app
self.http_methods_to_capture = http_methods_to_capture
if _looks_like_asgi3(app):
self.__call__ = self._run_asgi3 # type: Callable[..., Any]
else:
self.__call__ = self._run_asgi2
def _run_asgi2(self, scope):
# type: (Any) -> Any
async def inner(receive, send):
# type: (Any, Any) -> Any
return await self._run_app(scope, receive, send, asgi_version=2)
return inner
async def _run_asgi3(self, scope, receive, send):
# type: (Any, Any, Any) -> Any
return await self._run_app(scope, receive, send, asgi_version=3)
async def _run_app(self, scope, receive, send, asgi_version):
# type: (Any, Any, Any, Any, int) -> Any
is_recursive_asgi_middleware = _asgi_middleware_applied.get(False)
is_lifespan = scope["type"] == "lifespan"
if is_recursive_asgi_middleware or is_lifespan:
try:
if asgi_version == 2:
return await self.app(scope)(receive, send)
else:
return await self.app(scope, receive, send)
except Exception as exc:
_capture_exception(exc, mechanism_type=self.mechanism_type)
raise exc from None
_asgi_middleware_applied.set(True)
try:
with sentry_sdk.isolation_scope() as sentry_scope:
with track_session(sentry_scope, session_mode="request"):
sentry_scope.clear_breadcrumbs()
sentry_scope._name = "asgi"
processor = partial(self.event_processor, asgi_scope=scope)
sentry_scope.add_event_processor(processor)
ty = scope["type"]
(
transaction_name,
transaction_source,
) = self._get_transaction_name_and_source(
self.transaction_style,
scope,
)
method = scope.get("method", "").upper()
transaction = None
if method in self.http_methods_to_capture:
if ty in ("http", "websocket"):
transaction = continue_trace(
_get_headers(scope),
op="{}.server".format(ty),
name=transaction_name,
source=transaction_source,
origin=self.span_origin,
)
logger.debug(
"[ASGI] Created transaction (continuing trace): %s",
transaction,
)
else:
transaction = Transaction(
op=OP.HTTP_SERVER,
name=transaction_name,
source=transaction_source,
origin=self.span_origin,
)
logger.debug(
"[ASGI] Created transaction (new): %s", transaction
)
transaction.set_tag("asgi.type", ty)
logger.debug(
"[ASGI] Set transaction name and source on transaction: '%s' / '%s'",
transaction.name,
transaction.source,
)
with (
sentry_sdk.start_transaction(
transaction,
custom_sampling_context={"asgi_scope": scope},
)
if transaction is not None
else nullcontext()
):
logger.debug("[ASGI] Started transaction: %s", transaction)
try:
async def _sentry_wrapped_send(event):
# type: (Dict[str, Any]) -> Any
if transaction is not None:
is_http_response = (
event.get("type") == "http.response.start"
and "status" in event
)
if is_http_response:
transaction.set_http_status(event["status"])
return await send(event)
if asgi_version == 2:
return await self.app(scope)(
receive, _sentry_wrapped_send
)
else:
return await self.app(
scope, receive, _sentry_wrapped_send
)
except Exception as exc:
_capture_exception(exc, mechanism_type=self.mechanism_type)
raise exc from None
finally:
_asgi_middleware_applied.set(False)
def event_processor(self, event, hint, asgi_scope):
# type: (Event, Hint, Any) -> Optional[Event]
request_data = event.get("request", {})
request_data.update(_get_request_data(asgi_scope))
event["request"] = deepcopy(request_data)
# Only set transaction name if not already set by Starlette or FastAPI (or other frameworks)
already_set = event["transaction"] != _DEFAULT_TRANSACTION_NAME and event[
"transaction_info"
].get("source") in [
TRANSACTION_SOURCE_COMPONENT,
TRANSACTION_SOURCE_ROUTE,
TRANSACTION_SOURCE_CUSTOM,
]
if not already_set:
name, source = self._get_transaction_name_and_source(
self.transaction_style, asgi_scope
)
event["transaction"] = name
event["transaction_info"] = {"source": source}
logger.debug(
"[ASGI] Set transaction name and source in event_processor: '%s' / '%s'",
event["transaction"],
event["transaction_info"]["source"],
)
return event
# Helper functions.
#
# Note: Those functions are not public API. If you want to mutate request
# data to your liking it's recommended to use the `before_send` callback
# for that.
def _get_transaction_name_and_source(self, transaction_style, asgi_scope):
# type: (SentryAsgiMiddleware, str, Any) -> Tuple[str, str]
name = None
source = SOURCE_FOR_STYLE[transaction_style]
ty = asgi_scope.get("type")
if transaction_style == "endpoint":
endpoint = asgi_scope.get("endpoint")
# Webframeworks like Starlette mutate the ASGI env once routing is
# done, which is sometime after the request has started. If we have
# an endpoint, overwrite our generic transaction name.
if endpoint:
name = transaction_from_function(endpoint) or ""
else:
name = _get_url(asgi_scope, "http" if ty == "http" else "ws", host=None)
source = TRANSACTION_SOURCE_URL
elif transaction_style == "url":
# FastAPI includes the route object in the scope to let Sentry extract the
# path from it for the transaction name
route = asgi_scope.get("route")
if route:
path = getattr(route, "path", None)
if path is not None:
name = path
else:
name = _get_url(asgi_scope, "http" if ty == "http" else "ws", host=None)
source = TRANSACTION_SOURCE_URL
if name is None:
name = _DEFAULT_TRANSACTION_NAME
source = TRANSACTION_SOURCE_ROUTE
return name, source
return name, source
@@ -0,0 +1,107 @@
import sys
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import Integration, DidNotEnable
from sentry_sdk.utils import event_from_exception, reraise
try:
import asyncio
from asyncio.tasks import Task
except ImportError:
raise DidNotEnable("asyncio not available")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from collections.abc import Coroutine
from sentry_sdk._types import ExcInfo
def get_name(coro):
# type: (Any) -> str
return (
getattr(coro, "__qualname__", None)
or getattr(coro, "__name__", None)
or "coroutine without __name__"
)
def patch_asyncio():
# type: () -> None
orig_task_factory = None
try:
loop = asyncio.get_running_loop()
orig_task_factory = loop.get_task_factory()
def _sentry_task_factory(loop, coro, **kwargs):
# type: (asyncio.AbstractEventLoop, Coroutine[Any, Any, Any], Any) -> asyncio.Future[Any]
async def _coro_creating_hub_and_span():
# type: () -> Any
result = None
with sentry_sdk.isolation_scope():
with sentry_sdk.start_span(
op=OP.FUNCTION,
name=get_name(coro),
origin=AsyncioIntegration.origin,
):
try:
result = await coro
except Exception:
reraise(*_capture_exception())
return result
# Trying to use user set task factory (if there is one)
if orig_task_factory:
return orig_task_factory(loop, _coro_creating_hub_and_span(), **kwargs)
# The default task factory in `asyncio` does not have its own function
# but is just a couple of lines in `asyncio.base_events.create_task()`
# Those lines are copied here.
# WARNING:
# If the default behavior of the task creation in asyncio changes,
# this will break!
task = Task(_coro_creating_hub_and_span(), loop=loop, **kwargs)
if task._source_traceback: # type: ignore
del task._source_traceback[-1] # type: ignore
return task
loop.set_task_factory(_sentry_task_factory) # type: ignore
except RuntimeError:
# When there is no running loop, we have nothing to patch.
pass
def _capture_exception():
# type: () -> ExcInfo
exc_info = sys.exc_info()
client = sentry_sdk.get_client()
integration = client.get_integration(AsyncioIntegration)
if integration is not None:
event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "asyncio", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
return exc_info
class AsyncioIntegration(Integration):
identifier = "asyncio"
origin = f"auto.function.{identifier}"
@staticmethod
def setup_once():
# type: () -> None
patch_asyncio()
@@ -0,0 +1,208 @@
from __future__ import annotations
import contextlib
from typing import Any, TypeVar, Callable, Awaitable, Iterator
import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
from sentry_sdk.tracing import Span
from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
from sentry_sdk.utils import (
ensure_integration_enabled,
parse_version,
capture_internal_exceptions,
)
try:
import asyncpg # type: ignore[import-not-found]
from asyncpg.cursor import BaseCursor # type: ignore
except ImportError:
raise DidNotEnable("asyncpg not installed.")
class AsyncPGIntegration(Integration):
identifier = "asyncpg"
origin = f"auto.db.{identifier}"
_record_params = False
def __init__(self, *, record_params: bool = False):
AsyncPGIntegration._record_params = record_params
@staticmethod
def setup_once() -> None:
# asyncpg.__version__ is a string containing the semantic version in the form of "<major>.<minor>.<patch>"
asyncpg_version = parse_version(asyncpg.__version__)
_check_minimum_version(AsyncPGIntegration, asyncpg_version)
asyncpg.Connection.execute = _wrap_execute(
asyncpg.Connection.execute,
)
asyncpg.Connection._execute = _wrap_connection_method(
asyncpg.Connection._execute
)
asyncpg.Connection._executemany = _wrap_connection_method(
asyncpg.Connection._executemany, executemany=True
)
asyncpg.Connection.cursor = _wrap_cursor_creation(asyncpg.Connection.cursor)
asyncpg.Connection.prepare = _wrap_connection_method(asyncpg.Connection.prepare)
asyncpg.connect_utils._connect_addr = _wrap_connect_addr(
asyncpg.connect_utils._connect_addr
)
T = TypeVar("T")
def _wrap_execute(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
async def _inner(*args: Any, **kwargs: Any) -> T:
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
return await f(*args, **kwargs)
# Avoid recording calls to _execute twice.
# Calls to Connection.execute with args also call
# Connection._execute, which is recorded separately
# args[0] = the connection object, args[1] is the query
if len(args) > 2:
return await f(*args, **kwargs)
query = args[1]
with record_sql_queries(
cursor=None,
query=query,
params_list=None,
paramstyle=None,
executemany=False,
span_origin=AsyncPGIntegration.origin,
) as span:
res = await f(*args, **kwargs)
with capture_internal_exceptions():
add_query_source(span)
return res
return _inner
SubCursor = TypeVar("SubCursor", bound=BaseCursor)
@contextlib.contextmanager
def _record(
cursor: SubCursor | None,
query: str,
params_list: tuple[Any, ...] | None,
*,
executemany: bool = False,
) -> Iterator[Span]:
integration = sentry_sdk.get_client().get_integration(AsyncPGIntegration)
if integration is not None and not integration._record_params:
params_list = None
param_style = "pyformat" if params_list else None
with record_sql_queries(
cursor=cursor,
query=query,
params_list=params_list,
paramstyle=param_style,
executemany=executemany,
record_cursor_repr=cursor is not None,
span_origin=AsyncPGIntegration.origin,
) as span:
yield span
def _wrap_connection_method(
f: Callable[..., Awaitable[T]], *, executemany: bool = False
) -> Callable[..., Awaitable[T]]:
async def _inner(*args: Any, **kwargs: Any) -> T:
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
return await f(*args, **kwargs)
query = args[1]
params_list = args[2] if len(args) > 2 else None
with _record(None, query, params_list, executemany=executemany) as span:
_set_db_data(span, args[0])
res = await f(*args, **kwargs)
return res
return _inner
def _wrap_cursor_creation(f: Callable[..., T]) -> Callable[..., T]:
@ensure_integration_enabled(AsyncPGIntegration, f)
def _inner(*args: Any, **kwargs: Any) -> T: # noqa: N807
query = args[1]
params_list = args[2] if len(args) > 2 else None
with _record(
None,
query,
params_list,
executemany=False,
) as span:
_set_db_data(span, args[0])
res = f(*args, **kwargs)
span.set_data("db.cursor", res)
return res
return _inner
def _wrap_connect_addr(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
async def _inner(*args: Any, **kwargs: Any) -> T:
if sentry_sdk.get_client().get_integration(AsyncPGIntegration) is None:
return await f(*args, **kwargs)
user = kwargs["params"].user
database = kwargs["params"].database
with sentry_sdk.start_span(
op=OP.DB,
name="connect",
origin=AsyncPGIntegration.origin,
) as span:
span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
addr = kwargs.get("addr")
if addr:
try:
span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
span.set_data(SPANDATA.SERVER_PORT, addr[1])
except IndexError:
pass
span.set_data(SPANDATA.DB_NAME, database)
span.set_data(SPANDATA.DB_USER, user)
with capture_internal_exceptions():
sentry_sdk.add_breadcrumb(
message="connect", category="query", data=span._data
)
res = await f(*args, **kwargs)
return res
return _inner
def _set_db_data(span: Span, conn: Any) -> None:
span.set_data(SPANDATA.DB_SYSTEM, "postgresql")
addr = conn._addr
if addr:
try:
span.set_data(SPANDATA.SERVER_ADDRESS, addr[0])
span.set_data(SPANDATA.SERVER_PORT, addr[1])
except IndexError:
pass
database = conn._params.database
if database:
span.set_data(SPANDATA.DB_NAME, database)
user = conn._params.user
if user:
span.set_data(SPANDATA.DB_USER, user)
@@ -0,0 +1,57 @@
import os
import sys
import atexit
import sentry_sdk
from sentry_sdk.utils import logger
from sentry_sdk.integrations import Integration
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Optional
def default_callback(pending, timeout):
# type: (int, int) -> None
"""This is the default shutdown callback that is set on the options.
It prints out a message to stderr that informs the user that some events
are still pending and the process is waiting for them to flush out.
"""
def echo(msg):
# type: (str) -> None
sys.stderr.write(msg + "\n")
echo("Sentry is attempting to send %i pending events" % pending)
echo("Waiting up to %s seconds" % timeout)
echo("Press Ctrl-%s to quit" % (os.name == "nt" and "Break" or "C"))
sys.stderr.flush()
class AtexitIntegration(Integration):
identifier = "atexit"
def __init__(self, callback=None):
# type: (Optional[Any]) -> None
if callback is None:
callback = default_callback
self.callback = callback
@staticmethod
def setup_once():
# type: () -> None
@atexit.register
def _shutdown():
# type: () -> None
client = sentry_sdk.get_client()
integration = client.get_integration(AtexitIntegration)
if integration is None:
return
logger.debug("atexit: got shutdown signal")
logger.debug("atexit: shutting down client")
sentry_sdk.get_isolation_scope().end_session()
client.close(callback=integration.callback)
@@ -0,0 +1,496 @@
import functools
import json
import re
import sys
from copy import deepcopy
from datetime import datetime, timedelta, timezone
from os import environ
import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
from sentry_sdk.utils import (
AnnotatedValue,
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
logger,
TimeoutThread,
reraise,
)
from sentry_sdk.integrations import Integration
from sentry_sdk.integrations._wsgi_common import _filter_headers
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import TypeVar
from typing import Callable
from typing import Optional
from sentry_sdk._types import EventProcessor, Event, Hint
F = TypeVar("F", bound=Callable[..., Any])
# Constants
TIMEOUT_WARNING_BUFFER = 1500 # Buffer time required to send timeout warning to Sentry
MILLIS_TO_SECONDS = 1000.0
def _wrap_init_error(init_error):
# type: (F) -> F
@ensure_integration_enabled(AwsLambdaIntegration, init_error)
def sentry_init_error(*args, **kwargs):
# type: (*Any, **Any) -> Any
client = sentry_sdk.get_client()
with capture_internal_exceptions():
sentry_sdk.get_isolation_scope().clear_breadcrumbs()
exc_info = sys.exc_info()
if exc_info and all(exc_info):
sentry_event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "aws_lambda", "handled": False},
)
sentry_sdk.capture_event(sentry_event, hint=hint)
else:
# Fall back to AWS lambdas JSON representation of the error
sentry_event = _event_from_error_json(json.loads(args[1]))
sentry_sdk.capture_event(sentry_event)
return init_error(*args, **kwargs)
return sentry_init_error # type: ignore
def _wrap_handler(handler):
# type: (F) -> F
@functools.wraps(handler)
def sentry_handler(aws_event, aws_context, *args, **kwargs):
# type: (Any, Any, *Any, **Any) -> Any
# Per https://docs.aws.amazon.com/lambda/latest/dg/python-handler.html,
# `event` here is *likely* a dictionary, but also might be a number of
# other types (str, int, float, None).
#
# In some cases, it is a list (if the user is batch-invoking their
# function, for example), in which case we'll use the first entry as a
# representative from which to try pulling request data. (Presumably it
# will be the same for all events in the list, since they're all hitting
# the lambda in the same request.)
client = sentry_sdk.get_client()
integration = client.get_integration(AwsLambdaIntegration)
if integration is None:
return handler(aws_event, aws_context, *args, **kwargs)
if isinstance(aws_event, list) and len(aws_event) >= 1:
request_data = aws_event[0]
batch_size = len(aws_event)
else:
request_data = aws_event
batch_size = 1
if not isinstance(request_data, dict):
# If we're not dealing with a dictionary, we won't be able to get
# headers, path, http method, etc in any case, so it's fine that
# this is empty
request_data = {}
configured_time = aws_context.get_remaining_time_in_millis()
with sentry_sdk.isolation_scope() as scope:
timeout_thread = None
with capture_internal_exceptions():
scope.clear_breadcrumbs()
scope.add_event_processor(
_make_request_event_processor(
request_data, aws_context, configured_time
)
)
scope.set_tag(
"aws_region", aws_context.invoked_function_arn.split(":")[3]
)
if batch_size > 1:
scope.set_tag("batch_request", True)
scope.set_tag("batch_size", batch_size)
# Starting the Timeout thread only if the configured time is greater than Timeout warning
# buffer and timeout_warning parameter is set True.
if (
integration.timeout_warning
and configured_time > TIMEOUT_WARNING_BUFFER
):
waiting_time = (
configured_time - TIMEOUT_WARNING_BUFFER
) / MILLIS_TO_SECONDS
timeout_thread = TimeoutThread(
waiting_time,
configured_time / MILLIS_TO_SECONDS,
)
# Starting the thread to raise timeout warning exception
timeout_thread.start()
headers = request_data.get("headers", {})
# Some AWS Services (ie. EventBridge) set headers as a list
# or None, so we must ensure it is a dict
if not isinstance(headers, dict):
headers = {}
transaction = continue_trace(
headers,
op=OP.FUNCTION_AWS,
name=aws_context.function_name,
source=TRANSACTION_SOURCE_COMPONENT,
origin=AwsLambdaIntegration.origin,
)
with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={
"aws_event": aws_event,
"aws_context": aws_context,
},
):
try:
return handler(aws_event, aws_context, *args, **kwargs)
except Exception:
exc_info = sys.exc_info()
sentry_event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "aws_lambda", "handled": False},
)
sentry_sdk.capture_event(sentry_event, hint=hint)
reraise(*exc_info)
finally:
if timeout_thread:
timeout_thread.stop()
return sentry_handler # type: ignore
def _drain_queue():
# type: () -> None
with capture_internal_exceptions():
client = sentry_sdk.get_client()
integration = client.get_integration(AwsLambdaIntegration)
if integration is not None:
# Flush out the event queue before AWS kills the
# process.
client.flush()
class AwsLambdaIntegration(Integration):
identifier = "aws_lambda"
origin = f"auto.function.{identifier}"
def __init__(self, timeout_warning=False):
# type: (bool) -> None
self.timeout_warning = timeout_warning
@staticmethod
def setup_once():
# type: () -> None
lambda_bootstrap = get_lambda_bootstrap()
if not lambda_bootstrap:
logger.warning(
"Not running in AWS Lambda environment, "
"AwsLambdaIntegration disabled (could not find bootstrap module)"
)
return
if not hasattr(lambda_bootstrap, "handle_event_request"):
logger.warning(
"Not running in AWS Lambda environment, "
"AwsLambdaIntegration disabled (could not find handle_event_request)"
)
return
pre_37 = hasattr(lambda_bootstrap, "handle_http_request") # Python 3.6
if pre_37:
old_handle_event_request = lambda_bootstrap.handle_event_request
def sentry_handle_event_request(request_handler, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
request_handler = _wrap_handler(request_handler)
return old_handle_event_request(request_handler, *args, **kwargs)
lambda_bootstrap.handle_event_request = sentry_handle_event_request
old_handle_http_request = lambda_bootstrap.handle_http_request
def sentry_handle_http_request(request_handler, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
request_handler = _wrap_handler(request_handler)
return old_handle_http_request(request_handler, *args, **kwargs)
lambda_bootstrap.handle_http_request = sentry_handle_http_request
# Patch to_json to drain the queue. This should work even when the
# SDK is initialized inside of the handler
old_to_json = lambda_bootstrap.to_json
def sentry_to_json(*args, **kwargs):
# type: (*Any, **Any) -> Any
_drain_queue()
return old_to_json(*args, **kwargs)
lambda_bootstrap.to_json = sentry_to_json
else:
lambda_bootstrap.LambdaRuntimeClient.post_init_error = _wrap_init_error(
lambda_bootstrap.LambdaRuntimeClient.post_init_error
)
old_handle_event_request = lambda_bootstrap.handle_event_request
def sentry_handle_event_request( # type: ignore
lambda_runtime_client, request_handler, *args, **kwargs
):
request_handler = _wrap_handler(request_handler)
return old_handle_event_request(
lambda_runtime_client, request_handler, *args, **kwargs
)
lambda_bootstrap.handle_event_request = sentry_handle_event_request
# Patch the runtime client to drain the queue. This should work
# even when the SDK is initialized inside of the handler
def _wrap_post_function(f):
# type: (F) -> F
def inner(*args, **kwargs):
# type: (*Any, **Any) -> Any
_drain_queue()
return f(*args, **kwargs)
return inner # type: ignore
lambda_bootstrap.LambdaRuntimeClient.post_invocation_result = (
_wrap_post_function(
lambda_bootstrap.LambdaRuntimeClient.post_invocation_result
)
)
lambda_bootstrap.LambdaRuntimeClient.post_invocation_error = (
_wrap_post_function(
lambda_bootstrap.LambdaRuntimeClient.post_invocation_error
)
)
def get_lambda_bootstrap():
# type: () -> Optional[Any]
# Python 3.7: If the bootstrap module is *already imported*, it is the
# one we actually want to use (no idea what's in __main__)
#
# Python 3.8: bootstrap is also importable, but will be the same file
# as __main__ imported under a different name:
#
# sys.modules['__main__'].__file__ == sys.modules['bootstrap'].__file__
# sys.modules['__main__'] is not sys.modules['bootstrap']
#
# Python 3.9: bootstrap is in __main__.awslambdaricmain
#
# On container builds using the `aws-lambda-python-runtime-interface-client`
# (awslamdaric) module, bootstrap is located in sys.modules['__main__'].bootstrap
#
# Such a setup would then make all monkeypatches useless.
if "bootstrap" in sys.modules:
return sys.modules["bootstrap"]
elif "__main__" in sys.modules:
module = sys.modules["__main__"]
# python3.9 runtime
if hasattr(module, "awslambdaricmain") and hasattr(
module.awslambdaricmain, "bootstrap"
):
return module.awslambdaricmain.bootstrap
elif hasattr(module, "bootstrap"):
# awslambdaric python module in container builds
return module.bootstrap
# python3.8 runtime
return module
else:
return None
def _make_request_event_processor(aws_event, aws_context, configured_timeout):
# type: (Any, Any, Any) -> EventProcessor
start_time = datetime.now(timezone.utc)
def event_processor(sentry_event, hint, start_time=start_time):
# type: (Event, Hint, datetime) -> Optional[Event]
remaining_time_in_milis = aws_context.get_remaining_time_in_millis()
exec_duration = configured_timeout - remaining_time_in_milis
extra = sentry_event.setdefault("extra", {})
extra["lambda"] = {
"function_name": aws_context.function_name,
"function_version": aws_context.function_version,
"invoked_function_arn": aws_context.invoked_function_arn,
"aws_request_id": aws_context.aws_request_id,
"execution_duration_in_millis": exec_duration,
"remaining_time_in_millis": remaining_time_in_milis,
}
extra["cloudwatch logs"] = {
"url": _get_cloudwatch_logs_url(aws_context, start_time),
"log_group": aws_context.log_group_name,
"log_stream": aws_context.log_stream_name,
}
request = sentry_event.get("request", {})
if "httpMethod" in aws_event:
request["method"] = aws_event["httpMethod"]
request["url"] = _get_url(aws_event, aws_context)
if "queryStringParameters" in aws_event:
request["query_string"] = aws_event["queryStringParameters"]
if "headers" in aws_event:
request["headers"] = _filter_headers(aws_event["headers"])
if should_send_default_pii():
user_info = sentry_event.setdefault("user", {})
identity = aws_event.get("identity")
if identity is None:
identity = {}
id = identity.get("userArn")
if id is not None:
user_info.setdefault("id", id)
ip = identity.get("sourceIp")
if ip is not None:
user_info.setdefault("ip_address", ip)
if "body" in aws_event:
request["data"] = aws_event.get("body", "")
else:
if aws_event.get("body", None):
# Unfortunately couldn't find a way to get structured body from AWS
# event. Meaning every body is unstructured to us.
request["data"] = AnnotatedValue.removed_because_raw_data()
sentry_event["request"] = deepcopy(request)
return sentry_event
return event_processor
def _get_url(aws_event, aws_context):
# type: (Any, Any) -> str
path = aws_event.get("path", None)
headers = aws_event.get("headers")
if headers is None:
headers = {}
host = headers.get("Host", None)
proto = headers.get("X-Forwarded-Proto", None)
if proto and host and path:
return "{}://{}{}".format(proto, host, path)
return "awslambda:///{}".format(aws_context.function_name)
def _get_cloudwatch_logs_url(aws_context, start_time):
# type: (Any, datetime) -> str
"""
Generates a CloudWatchLogs console URL based on the context object
Arguments:
aws_context {Any} -- context from lambda handler
Returns:
str -- AWS Console URL to logs.
"""
formatstring = "%Y-%m-%dT%H:%M:%SZ"
region = environ.get("AWS_REGION", "")
url = (
"https://console.{domain}/cloudwatch/home?region={region}"
"#logEventViewer:group={log_group};stream={log_stream}"
";start={start_time};end={end_time}"
).format(
domain="amazonaws.cn" if region.startswith("cn-") else "aws.amazon.com",
region=region,
log_group=aws_context.log_group_name,
log_stream=aws_context.log_stream_name,
start_time=(start_time - timedelta(seconds=1)).strftime(formatstring),
end_time=(datetime.now(timezone.utc) + timedelta(seconds=2)).strftime(
formatstring
),
)
return url
def _parse_formatted_traceback(formatted_tb):
# type: (list[str]) -> list[dict[str, Any]]
frames = []
for frame in formatted_tb:
match = re.match(r'File "(.+)", line (\d+), in (.+)', frame.strip())
if match:
file_name, line_number, func_name = match.groups()
line_number = int(line_number)
frames.append(
{
"filename": file_name,
"function": func_name,
"lineno": line_number,
"vars": None,
"pre_context": None,
"context_line": None,
"post_context": None,
}
)
return frames
def _event_from_error_json(error_json):
# type: (dict[str, Any]) -> Event
"""
Converts the error JSON from AWS Lambda into a Sentry error event.
This is not a full fletched event, but better than nothing.
This is an example of where AWS creates the error JSON:
https://github.com/aws/aws-lambda-python-runtime-interface-client/blob/2.2.1/awslambdaric/bootstrap.py#L479
"""
event = {
"level": "error",
"exception": {
"values": [
{
"type": error_json.get("errorType"),
"value": error_json.get("errorMessage"),
"stacktrace": {
"frames": _parse_formatted_traceback(
error_json.get("stackTrace", [])
),
},
"mechanism": {
"type": "aws_lambda",
"handled": False,
},
}
],
},
} # type: Event
return event
@@ -0,0 +1,176 @@
import sys
import types
from functools import wraps
import sentry_sdk
from sentry_sdk.integrations import Integration
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
reraise,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Iterator
from typing import TypeVar
from typing import Callable
from sentry_sdk._types import ExcInfo
T = TypeVar("T")
F = TypeVar("F", bound=Callable[..., Any])
WRAPPED_FUNC = "_wrapped_{}_"
INSPECT_FUNC = "_inspect_{}" # Required format per apache_beam/transforms/core.py
USED_FUNC = "_sentry_used_"
class BeamIntegration(Integration):
identifier = "beam"
@staticmethod
def setup_once():
# type: () -> None
from apache_beam.transforms.core import DoFn, ParDo # type: ignore
ignore_logger("root")
ignore_logger("bundle_processor.create")
function_patches = ["process", "start_bundle", "finish_bundle", "setup"]
for func_name in function_patches:
setattr(
DoFn,
INSPECT_FUNC.format(func_name),
_wrap_inspect_call(DoFn, func_name),
)
old_init = ParDo.__init__
def sentry_init_pardo(self, fn, *args, **kwargs):
# type: (ParDo, Any, *Any, **Any) -> Any
# Do not monkey patch init twice
if not getattr(self, "_sentry_is_patched", False):
for func_name in function_patches:
if not hasattr(fn, func_name):
continue
wrapped_func = WRAPPED_FUNC.format(func_name)
# Check to see if inspect is set and process is not
# to avoid monkey patching process twice.
# Check to see if function is part of object for
# backwards compatibility.
process_func = getattr(fn, func_name)
inspect_func = getattr(fn, INSPECT_FUNC.format(func_name))
if not getattr(inspect_func, USED_FUNC, False) and not getattr(
process_func, USED_FUNC, False
):
setattr(fn, wrapped_func, process_func)
setattr(fn, func_name, _wrap_task_call(process_func))
self._sentry_is_patched = True
old_init(self, fn, *args, **kwargs)
ParDo.__init__ = sentry_init_pardo
def _wrap_inspect_call(cls, func_name):
# type: (Any, Any) -> Any
if not hasattr(cls, func_name):
return None
def _inspect(self):
# type: (Any) -> Any
"""
Inspect function overrides the way Beam gets argspec.
"""
wrapped_func = WRAPPED_FUNC.format(func_name)
if hasattr(self, wrapped_func):
process_func = getattr(self, wrapped_func)
else:
process_func = getattr(self, func_name)
setattr(self, func_name, _wrap_task_call(process_func))
setattr(self, wrapped_func, process_func)
# getfullargspec is deprecated in more recent beam versions and get_function_args_defaults
# (which uses Signatures internally) should be used instead.
try:
from apache_beam.transforms.core import get_function_args_defaults
return get_function_args_defaults(process_func)
except ImportError:
from apache_beam.typehints.decorators import getfullargspec # type: ignore
return getfullargspec(process_func)
setattr(_inspect, USED_FUNC, True)
return _inspect
def _wrap_task_call(func):
# type: (F) -> F
"""
Wrap task call with a try catch to get exceptions.
"""
@wraps(func)
def _inner(*args, **kwargs):
# type: (*Any, **Any) -> Any
try:
gen = func(*args, **kwargs)
except Exception:
raise_exception()
if not isinstance(gen, types.GeneratorType):
return gen
return _wrap_generator_call(gen)
setattr(_inner, USED_FUNC, True)
return _inner # type: ignore
@ensure_integration_enabled(BeamIntegration)
def _capture_exception(exc_info):
# type: (ExcInfo) -> None
"""
Send Beam exception to Sentry.
"""
client = sentry_sdk.get_client()
event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "beam", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def raise_exception():
# type: () -> None
"""
Raise an exception.
"""
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc_info)
reraise(*exc_info)
def _wrap_generator_call(gen):
# type: (Iterator[T]) -> Iterator[T]
"""
Wrap the generator to handle any failures.
"""
while True:
try:
yield next(gen)
except StopIteration:
break
except Exception:
raise_exception()
@@ -0,0 +1,137 @@
from functools import partial
import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
from sentry_sdk.tracing import Span
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
parse_url,
parse_version,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Dict
from typing import Optional
from typing import Type
try:
from botocore import __version__ as BOTOCORE_VERSION # type: ignore
from botocore.client import BaseClient # type: ignore
from botocore.response import StreamingBody # type: ignore
from botocore.awsrequest import AWSRequest # type: ignore
except ImportError:
raise DidNotEnable("botocore is not installed")
class Boto3Integration(Integration):
identifier = "boto3"
origin = f"auto.http.{identifier}"
@staticmethod
def setup_once():
# type: () -> None
version = parse_version(BOTOCORE_VERSION)
_check_minimum_version(Boto3Integration, version, "botocore")
orig_init = BaseClient.__init__
def sentry_patched_init(self, *args, **kwargs):
# type: (Type[BaseClient], *Any, **Any) -> None
orig_init(self, *args, **kwargs)
meta = self.meta
service_id = meta.service_model.service_id.hyphenize()
meta.events.register(
"request-created",
partial(_sentry_request_created, service_id=service_id),
)
meta.events.register("after-call", _sentry_after_call)
meta.events.register("after-call-error", _sentry_after_call_error)
BaseClient.__init__ = sentry_patched_init
@ensure_integration_enabled(Boto3Integration)
def _sentry_request_created(service_id, request, operation_name, **kwargs):
# type: (str, AWSRequest, str, **Any) -> None
description = "aws.%s.%s" % (service_id, operation_name)
span = sentry_sdk.start_span(
op=OP.HTTP_CLIENT,
name=description,
origin=Boto3Integration.origin,
)
with capture_internal_exceptions():
parsed_url = parse_url(request.url, sanitize=False)
span.set_data("aws.request.url", parsed_url.url)
span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
span.set_tag("aws.service_id", service_id)
span.set_tag("aws.operation_name", operation_name)
span.set_data(SPANDATA.HTTP_METHOD, request.method)
# We do it in order for subsequent http calls/retries be
# attached to this span.
span.__enter__()
# request.context is an open-ended data-structure
# where we can add anything useful in request life cycle.
request.context["_sentrysdk_span"] = span
def _sentry_after_call(context, parsed, **kwargs):
# type: (Dict[str, Any], Dict[str, Any], **Any) -> None
span = context.pop("_sentrysdk_span", None) # type: Optional[Span]
# Span could be absent if the integration is disabled.
if span is None:
return
span.__exit__(None, None, None)
body = parsed.get("Body")
if not isinstance(body, StreamingBody):
return
streaming_span = span.start_child(
op=OP.HTTP_CLIENT_STREAM,
name=span.description,
origin=Boto3Integration.origin,
)
orig_read = body.read
orig_close = body.close
def sentry_streaming_body_read(*args, **kwargs):
# type: (*Any, **Any) -> bytes
try:
ret = orig_read(*args, **kwargs)
if not ret:
streaming_span.finish()
return ret
except Exception:
streaming_span.finish()
raise
body.read = sentry_streaming_body_read
def sentry_streaming_body_close(*args, **kwargs):
# type: (*Any, **Any) -> None
streaming_span.finish()
orig_close(*args, **kwargs)
body.close = sentry_streaming_body_close
def _sentry_after_call_error(context, exception, **kwargs):
# type: (Dict[str, Any], Type[BaseException], **Any) -> None
span = context.pop("_sentrysdk_span", None) # type: Optional[Span]
# Span could be absent if the integration is disabled.
if span is None:
return
span.__exit__(type(exception), exception, None)
@@ -0,0 +1,215 @@
import functools
import sentry_sdk
from sentry_sdk.tracing import SOURCE_FOR_STYLE
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
parse_version,
transaction_from_function,
)
from sentry_sdk.integrations import (
Integration,
DidNotEnable,
_DEFAULT_FAILED_REQUEST_STATUS_CODES,
_check_minimum_version,
)
from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
from sentry_sdk.integrations._wsgi_common import RequestExtractor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Set
from sentry_sdk.integrations.wsgi import _ScopedResponse
from typing import Any
from typing import Dict
from typing import Callable
from typing import Optional
from bottle import FileUpload, FormsDict, LocalRequest # type: ignore
from sentry_sdk._types import EventProcessor, Event
try:
from bottle import (
Bottle,
HTTPResponse,
Route,
request as bottle_request,
__version__ as BOTTLE_VERSION,
)
except ImportError:
raise DidNotEnable("Bottle not installed")
TRANSACTION_STYLE_VALUES = ("endpoint", "url")
class BottleIntegration(Integration):
identifier = "bottle"
origin = f"auto.http.{identifier}"
transaction_style = ""
def __init__(
self,
transaction_style="endpoint", # type: str
*,
failed_request_status_codes=_DEFAULT_FAILED_REQUEST_STATUS_CODES, # type: Set[int]
):
# type: (...) -> None
if transaction_style not in TRANSACTION_STYLE_VALUES:
raise ValueError(
"Invalid value for transaction_style: %s (must be in %s)"
% (transaction_style, TRANSACTION_STYLE_VALUES)
)
self.transaction_style = transaction_style
self.failed_request_status_codes = failed_request_status_codes
@staticmethod
def setup_once():
# type: () -> None
version = parse_version(BOTTLE_VERSION)
_check_minimum_version(BottleIntegration, version)
old_app = Bottle.__call__
@ensure_integration_enabled(BottleIntegration, old_app)
def sentry_patched_wsgi_app(self, environ, start_response):
# type: (Any, Dict[str, str], Callable[..., Any]) -> _ScopedResponse
middleware = SentryWsgiMiddleware(
lambda *a, **kw: old_app(self, *a, **kw),
span_origin=BottleIntegration.origin,
)
return middleware(environ, start_response)
Bottle.__call__ = sentry_patched_wsgi_app
old_handle = Bottle._handle
@functools.wraps(old_handle)
def _patched_handle(self, environ):
# type: (Bottle, Dict[str, Any]) -> Any
integration = sentry_sdk.get_client().get_integration(BottleIntegration)
if integration is None:
return old_handle(self, environ)
scope = sentry_sdk.get_isolation_scope()
scope._name = "bottle"
scope.add_event_processor(
_make_request_event_processor(self, bottle_request, integration)
)
res = old_handle(self, environ)
return res
Bottle._handle = _patched_handle
old_make_callback = Route._make_callback
@functools.wraps(old_make_callback)
def patched_make_callback(self, *args, **kwargs):
# type: (Route, *object, **object) -> Any
prepared_callback = old_make_callback(self, *args, **kwargs)
integration = sentry_sdk.get_client().get_integration(BottleIntegration)
if integration is None:
return prepared_callback
def wrapped_callback(*args, **kwargs):
# type: (*object, **object) -> Any
try:
res = prepared_callback(*args, **kwargs)
except Exception as exception:
_capture_exception(exception, handled=False)
raise exception
if (
isinstance(res, HTTPResponse)
and res.status_code in integration.failed_request_status_codes
):
_capture_exception(res, handled=True)
return res
return wrapped_callback
Route._make_callback = patched_make_callback
class BottleRequestExtractor(RequestExtractor):
def env(self):
# type: () -> Dict[str, str]
return self.request.environ
def cookies(self):
# type: () -> Dict[str, str]
return self.request.cookies
def raw_data(self):
# type: () -> bytes
return self.request.body.read()
def form(self):
# type: () -> FormsDict
if self.is_json():
return None
return self.request.forms.decode()
def files(self):
# type: () -> Optional[Dict[str, str]]
if self.is_json():
return None
return self.request.files
def size_of_file(self, file):
# type: (FileUpload) -> int
return file.content_length
def _set_transaction_name_and_source(event, transaction_style, request):
# type: (Event, str, Any) -> None
name = ""
if transaction_style == "url":
name = request.route.rule or ""
elif transaction_style == "endpoint":
name = (
request.route.name
or transaction_from_function(request.route.callback)
or ""
)
event["transaction"] = name
event["transaction_info"] = {"source": SOURCE_FOR_STYLE[transaction_style]}
def _make_request_event_processor(app, request, integration):
# type: (Bottle, LocalRequest, BottleIntegration) -> EventProcessor
def event_processor(event, hint):
# type: (Event, dict[str, Any]) -> Event
_set_transaction_name_and_source(event, integration.transaction_style, request)
with capture_internal_exceptions():
BottleRequestExtractor(request).extract_into_event(event)
return event
return event_processor
def _capture_exception(exception, handled):
# type: (BaseException, bool) -> None
event, hint = event_from_exception(
exception,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "bottle", "handled": handled},
)
sentry_sdk.capture_event(event, hint=hint)
@@ -0,0 +1,528 @@
import sys
from collections.abc import Mapping
from functools import wraps
import sentry_sdk
from sentry_sdk import isolation_scope
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
from sentry_sdk.integrations.celery.beat import (
_patch_beat_apply_entry,
_patch_redbeat_maybe_due,
_setup_celery_beat_signals,
)
from sentry_sdk.integrations.celery.utils import _now_seconds_since_epoch
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TRANSACTION_SOURCE_TASK
from sentry_sdk.tracing_utils import Baggage
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
reraise,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import List
from typing import Optional
from typing import TypeVar
from typing import Union
from sentry_sdk._types import EventProcessor, Event, Hint, ExcInfo
from sentry_sdk.tracing import Span
F = TypeVar("F", bound=Callable[..., Any])
try:
from celery import VERSION as CELERY_VERSION # type: ignore
from celery.app.task import Task # type: ignore
from celery.app.trace import task_has_custom
from celery.exceptions import ( # type: ignore
Ignore,
Reject,
Retry,
SoftTimeLimitExceeded,
)
from kombu import Producer # type: ignore
except ImportError:
raise DidNotEnable("Celery not installed")
CELERY_CONTROL_FLOW_EXCEPTIONS = (Retry, Ignore, Reject)
class CeleryIntegration(Integration):
identifier = "celery"
origin = f"auto.queue.{identifier}"
def __init__(
self,
propagate_traces=True,
monitor_beat_tasks=False,
exclude_beat_tasks=None,
):
# type: (bool, bool, Optional[List[str]]) -> None
self.propagate_traces = propagate_traces
self.monitor_beat_tasks = monitor_beat_tasks
self.exclude_beat_tasks = exclude_beat_tasks
_patch_beat_apply_entry()
_patch_redbeat_maybe_due()
_setup_celery_beat_signals(monitor_beat_tasks)
@staticmethod
def setup_once():
# type: () -> None
_check_minimum_version(CeleryIntegration, CELERY_VERSION)
_patch_build_tracer()
_patch_task_apply_async()
_patch_celery_send_task()
_patch_worker_exit()
_patch_producer_publish()
# This logger logs every status of every task that ran on the worker.
# Meaning that every task's breadcrumbs are full of stuff like "Task
# <foo> raised unexpected <bar>".
ignore_logger("celery.worker.job")
ignore_logger("celery.app.trace")
# This is stdout/err redirected to a logger, can't deal with this
# (need event_level=logging.WARN to reproduce)
ignore_logger("celery.redirected")
def _set_status(status):
# type: (str) -> None
with capture_internal_exceptions():
scope = sentry_sdk.get_current_scope()
if scope.span is not None:
scope.span.set_status(status)
def _capture_exception(task, exc_info):
# type: (Any, ExcInfo) -> None
client = sentry_sdk.get_client()
if client.get_integration(CeleryIntegration) is None:
return
if isinstance(exc_info[1], CELERY_CONTROL_FLOW_EXCEPTIONS):
# ??? Doesn't map to anything
_set_status("aborted")
return
_set_status("internal_error")
if hasattr(task, "throws") and isinstance(exc_info[1], task.throws):
return
event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "celery", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _make_event_processor(task, uuid, args, kwargs, request=None):
# type: (Any, Any, Any, Any, Optional[Any]) -> EventProcessor
def event_processor(event, hint):
# type: (Event, Hint) -> Optional[Event]
with capture_internal_exceptions():
tags = event.setdefault("tags", {})
tags["celery_task_id"] = uuid
extra = event.setdefault("extra", {})
extra["celery-job"] = {
"task_name": task.name,
"args": args,
"kwargs": kwargs,
}
if "exc_info" in hint:
with capture_internal_exceptions():
if issubclass(hint["exc_info"][0], SoftTimeLimitExceeded):
event["fingerprint"] = [
"celery",
"SoftTimeLimitExceeded",
getattr(task, "name", task),
]
return event
return event_processor
def _update_celery_task_headers(original_headers, span, monitor_beat_tasks):
# type: (dict[str, Any], Optional[Span], bool) -> dict[str, Any]
"""
Updates the headers of the Celery task with the tracing information
and eventually Sentry Crons monitoring information for beat tasks.
"""
updated_headers = original_headers.copy()
with capture_internal_exceptions():
# if span is None (when the task was started by Celery Beat)
# this will return the trace headers from the scope.
headers = dict(
sentry_sdk.get_isolation_scope().iter_trace_propagation_headers(span=span)
)
if monitor_beat_tasks:
headers.update(
{
"sentry-monitor-start-timestamp-s": "%.9f"
% _now_seconds_since_epoch(),
}
)
# Add the time the task was enqueued to the headers
# This is used in the consumer to calculate the latency
updated_headers.update(
{"sentry-task-enqueued-time": _now_seconds_since_epoch()}
)
if headers:
existing_baggage = updated_headers.get(BAGGAGE_HEADER_NAME)
sentry_baggage = headers.get(BAGGAGE_HEADER_NAME)
combined_baggage = sentry_baggage or existing_baggage
if sentry_baggage and existing_baggage:
# Merge incoming and sentry baggage, where the sentry trace information
# in the incoming baggage takes precedence and the third-party items
# are concatenated.
incoming = Baggage.from_incoming_header(existing_baggage)
combined = Baggage.from_incoming_header(sentry_baggage)
combined.sentry_items.update(incoming.sentry_items)
combined.third_party_items = ",".join(
[
x
for x in [
combined.third_party_items,
incoming.third_party_items,
]
if x is not None and x != ""
]
)
combined_baggage = combined.serialize(include_third_party=True)
updated_headers.update(headers)
if combined_baggage:
updated_headers[BAGGAGE_HEADER_NAME] = combined_baggage
# https://github.com/celery/celery/issues/4875
#
# Need to setdefault the inner headers too since other
# tracing tools (dd-trace-py) also employ this exact
# workaround and we don't want to break them.
updated_headers.setdefault("headers", {}).update(headers)
if combined_baggage:
updated_headers["headers"][BAGGAGE_HEADER_NAME] = combined_baggage
# Add the Sentry options potentially added in `sentry_apply_entry`
# to the headers (done when auto-instrumenting Celery Beat tasks)
for key, value in updated_headers.items():
if key.startswith("sentry-"):
updated_headers["headers"][key] = value
return updated_headers
class NoOpMgr:
def __enter__(self):
# type: () -> None
return None
def __exit__(self, exc_type, exc_value, traceback):
# type: (Any, Any, Any) -> None
return None
def _wrap_task_run(f):
# type: (F) -> F
@wraps(f)
def apply_async(*args, **kwargs):
# type: (*Any, **Any) -> Any
# Note: kwargs can contain headers=None, so no setdefault!
# Unsure which backend though.
integration = sentry_sdk.get_client().get_integration(CeleryIntegration)
if integration is None:
return f(*args, **kwargs)
kwarg_headers = kwargs.get("headers") or {}
propagate_traces = kwarg_headers.pop(
"sentry-propagate-traces", integration.propagate_traces
)
if not propagate_traces:
return f(*args, **kwargs)
if isinstance(args[0], Task):
task_name = args[0].name # type: str
elif len(args) > 1 and isinstance(args[1], str):
task_name = args[1]
else:
task_name = "<unknown Celery task>"
task_started_from_beat = sentry_sdk.get_isolation_scope()._name == "celery-beat"
span_mgr = (
sentry_sdk.start_span(
op=OP.QUEUE_SUBMIT_CELERY,
name=task_name,
origin=CeleryIntegration.origin,
)
if not task_started_from_beat
else NoOpMgr()
) # type: Union[Span, NoOpMgr]
with span_mgr as span:
kwargs["headers"] = _update_celery_task_headers(
kwarg_headers, span, integration.monitor_beat_tasks
)
return f(*args, **kwargs)
return apply_async # type: ignore
def _wrap_tracer(task, f):
# type: (Any, F) -> F
# Need to wrap tracer for pushing the scope before prerun is sent, and
# popping it after postrun is sent.
#
# This is the reason we don't use signals for hooking in the first place.
# Also because in Celery 3, signal dispatch returns early if one handler
# crashes.
@wraps(f)
@ensure_integration_enabled(CeleryIntegration, f)
def _inner(*args, **kwargs):
# type: (*Any, **Any) -> Any
with isolation_scope() as scope:
scope._name = "celery"
scope.clear_breadcrumbs()
scope.add_event_processor(_make_event_processor(task, *args, **kwargs))
transaction = None
# Celery task objects are not a thing to be trusted. Even
# something such as attribute access can fail.
with capture_internal_exceptions():
headers = args[3].get("headers") or {}
transaction = continue_trace(
headers,
op=OP.QUEUE_TASK_CELERY,
name="unknown celery task",
source=TRANSACTION_SOURCE_TASK,
origin=CeleryIntegration.origin,
)
transaction.name = task.name
transaction.set_status(SPANSTATUS.OK)
if transaction is None:
return f(*args, **kwargs)
with sentry_sdk.start_transaction(
transaction,
custom_sampling_context={
"celery_job": {
"task": task.name,
# for some reason, args[1] is a list if non-empty but a
# tuple if empty
"args": list(args[1]),
"kwargs": args[2],
}
},
):
return f(*args, **kwargs)
return _inner # type: ignore
def _set_messaging_destination_name(task, span):
# type: (Any, Span) -> None
"""Set "messaging.destination.name" tag for span"""
with capture_internal_exceptions():
delivery_info = task.request.delivery_info
if delivery_info:
routing_key = delivery_info.get("routing_key")
if delivery_info.get("exchange") == "" and routing_key is not None:
# Empty exchange indicates the default exchange, meaning the tasks
# are sent to the queue with the same name as the routing key.
span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
def _wrap_task_call(task, f):
# type: (Any, F) -> F
# Need to wrap task call because the exception is caught before we get to
# see it. Also celery's reported stacktrace is untrustworthy.
# functools.wraps is important here because celery-once looks at this
# method's name. @ensure_integration_enabled internally calls functools.wraps,
# but if we ever remove the @ensure_integration_enabled decorator, we need
# to add @functools.wraps(f) here.
# https://github.com/getsentry/sentry-python/issues/421
@ensure_integration_enabled(CeleryIntegration, f)
def _inner(*args, **kwargs):
# type: (*Any, **Any) -> Any
try:
with sentry_sdk.start_span(
op=OP.QUEUE_PROCESS,
name=task.name,
origin=CeleryIntegration.origin,
) as span:
_set_messaging_destination_name(task, span)
latency = None
with capture_internal_exceptions():
if (
task.request.headers is not None
and "sentry-task-enqueued-time" in task.request.headers
):
latency = _now_seconds_since_epoch() - task.request.headers.pop(
"sentry-task-enqueued-time"
)
if latency is not None:
span.set_data(SPANDATA.MESSAGING_MESSAGE_RECEIVE_LATENCY, latency)
with capture_internal_exceptions():
span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, task.request.id)
with capture_internal_exceptions():
span.set_data(
SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, task.request.retries
)
with capture_internal_exceptions():
span.set_data(
SPANDATA.MESSAGING_SYSTEM,
task.app.connection().transport.driver_type,
)
return f(*args, **kwargs)
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(task, exc_info)
reraise(*exc_info)
return _inner # type: ignore
def _patch_build_tracer():
# type: () -> None
import celery.app.trace as trace # type: ignore
original_build_tracer = trace.build_tracer
def sentry_build_tracer(name, task, *args, **kwargs):
# type: (Any, Any, *Any, **Any) -> Any
if not getattr(task, "_sentry_is_patched", False):
# determine whether Celery will use __call__ or run and patch
# accordingly
if task_has_custom(task, "__call__"):
type(task).__call__ = _wrap_task_call(task, type(task).__call__)
else:
task.run = _wrap_task_call(task, task.run)
# `build_tracer` is apparently called for every task
# invocation. Can't wrap every celery task for every invocation
# or we will get infinitely nested wrapper functions.
task._sentry_is_patched = True
return _wrap_tracer(task, original_build_tracer(name, task, *args, **kwargs))
trace.build_tracer = sentry_build_tracer
def _patch_task_apply_async():
# type: () -> None
Task.apply_async = _wrap_task_run(Task.apply_async)
def _patch_celery_send_task():
# type: () -> None
from celery import Celery
Celery.send_task = _wrap_task_run(Celery.send_task)
def _patch_worker_exit():
# type: () -> None
# Need to flush queue before worker shutdown because a crashing worker will
# call os._exit
from billiard.pool import Worker # type: ignore
original_workloop = Worker.workloop
def sentry_workloop(*args, **kwargs):
# type: (*Any, **Any) -> Any
try:
return original_workloop(*args, **kwargs)
finally:
with capture_internal_exceptions():
if (
sentry_sdk.get_client().get_integration(CeleryIntegration)
is not None
):
sentry_sdk.flush()
Worker.workloop = sentry_workloop
def _patch_producer_publish():
# type: () -> None
original_publish = Producer.publish
@ensure_integration_enabled(CeleryIntegration, original_publish)
def sentry_publish(self, *args, **kwargs):
# type: (Producer, *Any, **Any) -> Any
kwargs_headers = kwargs.get("headers", {})
if not isinstance(kwargs_headers, Mapping):
# Ensure kwargs_headers is a Mapping, so we can safely call get().
# We don't expect this to happen, but it's better to be safe. Even
# if it does happen, only our instrumentation breaks. This line
# does not overwrite kwargs["headers"], so the original publish
# method will still work.
kwargs_headers = {}
task_name = kwargs_headers.get("task")
task_id = kwargs_headers.get("id")
retries = kwargs_headers.get("retries")
routing_key = kwargs.get("routing_key")
exchange = kwargs.get("exchange")
with sentry_sdk.start_span(
op=OP.QUEUE_PUBLISH,
name=task_name,
origin=CeleryIntegration.origin,
) as span:
if task_id is not None:
span.set_data(SPANDATA.MESSAGING_MESSAGE_ID, task_id)
if exchange == "" and routing_key is not None:
# Empty exchange indicates the default exchange, meaning messages are
# routed to the queue with the same name as the routing key.
span.set_data(SPANDATA.MESSAGING_DESTINATION_NAME, routing_key)
if retries is not None:
span.set_data(SPANDATA.MESSAGING_MESSAGE_RETRY_COUNT, retries)
with capture_internal_exceptions():
span.set_data(
SPANDATA.MESSAGING_SYSTEM, self.connection.transport.driver_type
)
return original_publish(self, *args, **kwargs)
Producer.publish = sentry_publish
@@ -0,0 +1,293 @@
import sentry_sdk
from sentry_sdk.crons import capture_checkin, MonitorStatus
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.integrations.celery.utils import (
_get_humanized_interval,
_now_seconds_since_epoch,
)
from sentry_sdk.utils import (
logger,
match_regex_list,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Callable
from typing import Any, Optional, TypeVar, Union
from sentry_sdk._types import (
MonitorConfig,
MonitorConfigScheduleType,
MonitorConfigScheduleUnit,
)
F = TypeVar("F", bound=Callable[..., Any])
try:
from celery import Task, Celery # type: ignore
from celery.beat import Scheduler # type: ignore
from celery.schedules import crontab, schedule # type: ignore
from celery.signals import ( # type: ignore
task_failure,
task_success,
task_retry,
)
except ImportError:
raise DidNotEnable("Celery not installed")
try:
from redbeat.schedulers import RedBeatScheduler # type: ignore
except ImportError:
RedBeatScheduler = None
def _get_headers(task):
# type: (Task) -> dict[str, Any]
headers = task.request.get("headers") or {}
# flatten nested headers
if "headers" in headers:
headers.update(headers["headers"])
del headers["headers"]
headers.update(task.request.get("properties") or {})
return headers
def _get_monitor_config(celery_schedule, app, monitor_name):
# type: (Any, Celery, str) -> MonitorConfig
monitor_config = {} # type: MonitorConfig
schedule_type = None # type: Optional[MonitorConfigScheduleType]
schedule_value = None # type: Optional[Union[str, int]]
schedule_unit = None # type: Optional[MonitorConfigScheduleUnit]
if isinstance(celery_schedule, crontab):
schedule_type = "crontab"
schedule_value = (
"{0._orig_minute} "
"{0._orig_hour} "
"{0._orig_day_of_month} "
"{0._orig_month_of_year} "
"{0._orig_day_of_week}".format(celery_schedule)
)
elif isinstance(celery_schedule, schedule):
schedule_type = "interval"
(schedule_value, schedule_unit) = _get_humanized_interval(
celery_schedule.seconds
)
if schedule_unit == "second":
logger.warning(
"Intervals shorter than one minute are not supported by Sentry Crons. Monitor '%s' has an interval of %s seconds. Use the `exclude_beat_tasks` option in the celery integration to exclude it.",
monitor_name,
schedule_value,
)
return {}
else:
logger.warning(
"Celery schedule type '%s' not supported by Sentry Crons.",
type(celery_schedule),
)
return {}
monitor_config["schedule"] = {}
monitor_config["schedule"]["type"] = schedule_type
monitor_config["schedule"]["value"] = schedule_value
if schedule_unit is not None:
monitor_config["schedule"]["unit"] = schedule_unit
monitor_config["timezone"] = (
(
hasattr(celery_schedule, "tz")
and celery_schedule.tz is not None
and str(celery_schedule.tz)
)
or app.timezone
or "UTC"
)
return monitor_config
def _apply_crons_data_to_schedule_entry(scheduler, schedule_entry, integration):
# type: (Any, Any, sentry_sdk.integrations.celery.CeleryIntegration) -> None
"""
Add Sentry Crons information to the schedule_entry headers.
"""
if not integration.monitor_beat_tasks:
return
monitor_name = schedule_entry.name
task_should_be_excluded = match_regex_list(
monitor_name, integration.exclude_beat_tasks
)
if task_should_be_excluded:
return
celery_schedule = schedule_entry.schedule
app = scheduler.app
monitor_config = _get_monitor_config(celery_schedule, app, monitor_name)
is_supported_schedule = bool(monitor_config)
if not is_supported_schedule:
return
headers = schedule_entry.options.pop("headers", {})
headers.update(
{
"sentry-monitor-slug": monitor_name,
"sentry-monitor-config": monitor_config,
}
)
check_in_id = capture_checkin(
monitor_slug=monitor_name,
monitor_config=monitor_config,
status=MonitorStatus.IN_PROGRESS,
)
headers.update({"sentry-monitor-check-in-id": check_in_id})
# Set the Sentry configuration in the options of the ScheduleEntry.
# Those will be picked up in `apply_async` and added to the headers.
schedule_entry.options["headers"] = headers
def _wrap_beat_scheduler(original_function):
# type: (Callable[..., Any]) -> Callable[..., Any]
"""
Makes sure that:
- a new Sentry trace is started for each task started by Celery Beat and
it is propagated to the task.
- the Sentry Crons information is set in the Celery Beat task's
headers so that is is monitored with Sentry Crons.
After the patched function is called,
Celery Beat will call apply_async to put the task in the queue.
"""
# Patch only once
# Can't use __name__ here, because some of our tests mock original_apply_entry
already_patched = "sentry_patched_scheduler" in str(original_function)
if already_patched:
return original_function
from sentry_sdk.integrations.celery import CeleryIntegration
def sentry_patched_scheduler(*args, **kwargs):
# type: (*Any, **Any) -> None
integration = sentry_sdk.get_client().get_integration(CeleryIntegration)
if integration is None:
return original_function(*args, **kwargs)
# Tasks started by Celery Beat start a new Trace
scope = sentry_sdk.get_isolation_scope()
scope.set_new_propagation_context()
scope._name = "celery-beat"
scheduler, schedule_entry = args
_apply_crons_data_to_schedule_entry(scheduler, schedule_entry, integration)
return original_function(*args, **kwargs)
return sentry_patched_scheduler
def _patch_beat_apply_entry():
# type: () -> None
Scheduler.apply_entry = _wrap_beat_scheduler(Scheduler.apply_entry)
def _patch_redbeat_maybe_due():
# type: () -> None
if RedBeatScheduler is None:
return
RedBeatScheduler.maybe_due = _wrap_beat_scheduler(RedBeatScheduler.maybe_due)
def _setup_celery_beat_signals(monitor_beat_tasks):
# type: (bool) -> None
if monitor_beat_tasks:
task_success.connect(crons_task_success)
task_failure.connect(crons_task_failure)
task_retry.connect(crons_task_retry)
def crons_task_success(sender, **kwargs):
# type: (Task, dict[Any, Any]) -> None
logger.debug("celery_task_success %s", sender)
headers = _get_headers(sender)
if "sentry-monitor-slug" not in headers:
return
monitor_config = headers.get("sentry-monitor-config", {})
start_timestamp_s = headers.get("sentry-monitor-start-timestamp-s")
capture_checkin(
monitor_slug=headers["sentry-monitor-slug"],
monitor_config=monitor_config,
check_in_id=headers["sentry-monitor-check-in-id"],
duration=(
_now_seconds_since_epoch() - float(start_timestamp_s)
if start_timestamp_s
else None
),
status=MonitorStatus.OK,
)
def crons_task_failure(sender, **kwargs):
# type: (Task, dict[Any, Any]) -> None
logger.debug("celery_task_failure %s", sender)
headers = _get_headers(sender)
if "sentry-monitor-slug" not in headers:
return
monitor_config = headers.get("sentry-monitor-config", {})
start_timestamp_s = headers.get("sentry-monitor-start-timestamp-s")
capture_checkin(
monitor_slug=headers["sentry-monitor-slug"],
monitor_config=monitor_config,
check_in_id=headers["sentry-monitor-check-in-id"],
duration=(
_now_seconds_since_epoch() - float(start_timestamp_s)
if start_timestamp_s
else None
),
status=MonitorStatus.ERROR,
)
def crons_task_retry(sender, **kwargs):
# type: (Task, dict[Any, Any]) -> None
logger.debug("celery_task_retry %s", sender)
headers = _get_headers(sender)
if "sentry-monitor-slug" not in headers:
return
monitor_config = headers.get("sentry-monitor-config", {})
start_timestamp_s = headers.get("sentry-monitor-start-timestamp-s")
capture_checkin(
monitor_slug=headers["sentry-monitor-slug"],
monitor_config=monitor_config,
check_in_id=headers["sentry-monitor-check-in-id"],
duration=(
_now_seconds_since_epoch() - float(start_timestamp_s)
if start_timestamp_s
else None
),
status=MonitorStatus.ERROR,
)
@@ -0,0 +1,43 @@
import time
from typing import TYPE_CHECKING, cast
if TYPE_CHECKING:
from typing import Any, Tuple
from sentry_sdk._types import MonitorConfigScheduleUnit
def _now_seconds_since_epoch():
# type: () -> float
# We cannot use `time.perf_counter()` when dealing with the duration
# of a Celery task, because the start of a Celery task and
# the end are recorded in different processes.
# Start happens in the Celery Beat process,
# the end in a Celery Worker process.
return time.time()
def _get_humanized_interval(seconds):
# type: (float) -> Tuple[int, MonitorConfigScheduleUnit]
TIME_UNITS = ( # noqa: N806
("day", 60 * 60 * 24.0),
("hour", 60 * 60.0),
("minute", 60.0),
)
seconds = float(seconds)
for unit, divider in TIME_UNITS:
if seconds >= divider:
interval = int(seconds / divider)
return (interval, cast("MonitorConfigScheduleUnit", unit))
return (int(seconds), "second")
class NoOpMgr:
def __enter__(self):
# type: () -> None
return None
def __exit__(self, exc_type, exc_value, traceback):
# type: (Any, Any, Any) -> None
return None
@@ -0,0 +1,134 @@
import sys
from functools import wraps
import sentry_sdk
from sentry_sdk.integrations import Integration, DidNotEnable
from sentry_sdk.integrations.aws_lambda import _make_request_event_processor
from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
from sentry_sdk.utils import (
capture_internal_exceptions,
event_from_exception,
parse_version,
reraise,
)
try:
import chalice # type: ignore
from chalice import __version__ as CHALICE_VERSION
from chalice import Chalice, ChaliceViewError
from chalice.app import EventSourceHandler as ChaliceEventSourceHandler # type: ignore
except ImportError:
raise DidNotEnable("Chalice is not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Dict
from typing import TypeVar
from typing import Callable
F = TypeVar("F", bound=Callable[..., Any])
class EventSourceHandler(ChaliceEventSourceHandler): # type: ignore
def __call__(self, event, context):
# type: (Any, Any) -> Any
client = sentry_sdk.get_client()
with sentry_sdk.isolation_scope() as scope:
with capture_internal_exceptions():
configured_time = context.get_remaining_time_in_millis()
scope.add_event_processor(
_make_request_event_processor(event, context, configured_time)
)
try:
return ChaliceEventSourceHandler.__call__(self, event, context)
except Exception:
exc_info = sys.exc_info()
event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "chalice", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
client.flush()
reraise(*exc_info)
def _get_view_function_response(app, view_function, function_args):
# type: (Any, F, Any) -> F
@wraps(view_function)
def wrapped_view_function(**function_args):
# type: (**Any) -> Any
client = sentry_sdk.get_client()
with sentry_sdk.isolation_scope() as scope:
with capture_internal_exceptions():
configured_time = app.lambda_context.get_remaining_time_in_millis()
scope.set_transaction_name(
app.lambda_context.function_name,
source=TRANSACTION_SOURCE_COMPONENT,
)
scope.add_event_processor(
_make_request_event_processor(
app.current_request.to_dict(),
app.lambda_context,
configured_time,
)
)
try:
return view_function(**function_args)
except Exception as exc:
if isinstance(exc, ChaliceViewError):
raise
exc_info = sys.exc_info()
event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "chalice", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
client.flush()
raise
return wrapped_view_function # type: ignore
class ChaliceIntegration(Integration):
identifier = "chalice"
@staticmethod
def setup_once():
# type: () -> None
version = parse_version(CHALICE_VERSION)
if version is None:
raise DidNotEnable("Unparsable Chalice version: {}".format(CHALICE_VERSION))
if version < (1, 20):
old_get_view_function_response = Chalice._get_view_function_response
else:
from chalice.app import RestAPIEventHandler
old_get_view_function_response = (
RestAPIEventHandler._get_view_function_response
)
def sentry_event_response(app, view_function, function_args):
# type: (Any, F, Dict[str, Any]) -> Any
wrapped_view_function = _get_view_function_response(
app, view_function, function_args
)
return old_get_view_function_response(
app, wrapped_view_function, function_args
)
if version < (1, 20):
Chalice._get_view_function_response = sentry_event_response
else:
RestAPIEventHandler._get_view_function_response = sentry_event_response
# for everything else (like events)
chalice.app.EventSourceHandler = EventSourceHandler
@@ -0,0 +1,157 @@
import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
from sentry_sdk.tracing import Span
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import capture_internal_exceptions, ensure_integration_enabled
from typing import TYPE_CHECKING, TypeVar
# Hack to get new Python features working in older versions
# without introducing a hard dependency on `typing_extensions`
# from: https://stackoverflow.com/a/71944042/300572
if TYPE_CHECKING:
from typing import ParamSpec, Callable
else:
# Fake ParamSpec
class ParamSpec:
def __init__(self, _):
self.args = None
self.kwargs = None
# Callable[anything] will return None
class _Callable:
def __getitem__(self, _):
return None
# Make instances
Callable = _Callable()
try:
import clickhouse_driver # type: ignore[import-not-found]
except ImportError:
raise DidNotEnable("clickhouse-driver not installed.")
class ClickhouseDriverIntegration(Integration):
identifier = "clickhouse_driver"
origin = f"auto.db.{identifier}"
@staticmethod
def setup_once() -> None:
_check_minimum_version(ClickhouseDriverIntegration, clickhouse_driver.VERSION)
# Every query is done using the Connection's `send_query` function
clickhouse_driver.connection.Connection.send_query = _wrap_start(
clickhouse_driver.connection.Connection.send_query
)
# If the query contains parameters then the send_data function is used to send those parameters to clickhouse
clickhouse_driver.client.Client.send_data = _wrap_send_data(
clickhouse_driver.client.Client.send_data
)
# Every query ends either with the Client's `receive_end_of_query` (no result expected)
# or its `receive_result` (result expected)
clickhouse_driver.client.Client.receive_end_of_query = _wrap_end(
clickhouse_driver.client.Client.receive_end_of_query
)
if hasattr(clickhouse_driver.client.Client, "receive_end_of_insert_query"):
# In 0.2.7, insert queries are handled separately via `receive_end_of_insert_query`
clickhouse_driver.client.Client.receive_end_of_insert_query = _wrap_end(
clickhouse_driver.client.Client.receive_end_of_insert_query
)
clickhouse_driver.client.Client.receive_result = _wrap_end(
clickhouse_driver.client.Client.receive_result
)
P = ParamSpec("P")
T = TypeVar("T")
def _wrap_start(f: Callable[P, T]) -> Callable[P, T]:
@ensure_integration_enabled(ClickhouseDriverIntegration, f)
def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
connection = args[0]
query = args[1]
query_id = args[2] if len(args) > 2 else kwargs.get("query_id")
params = args[3] if len(args) > 3 else kwargs.get("params")
span = sentry_sdk.start_span(
op=OP.DB,
name=query,
origin=ClickhouseDriverIntegration.origin,
)
connection._sentry_span = span # type: ignore[attr-defined]
_set_db_data(span, connection)
span.set_data("query", query)
if query_id:
span.set_data("db.query_id", query_id)
if params and should_send_default_pii():
span.set_data("db.params", params)
# run the original code
ret = f(*args, **kwargs)
return ret
return _inner
def _wrap_end(f: Callable[P, T]) -> Callable[P, T]:
def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
res = f(*args, **kwargs)
instance = args[0]
span = getattr(instance.connection, "_sentry_span", None) # type: ignore[attr-defined]
if span is not None:
if res is not None and should_send_default_pii():
span.set_data("db.result", res)
with capture_internal_exceptions():
span.scope.add_breadcrumb(
message=span._data.pop("query"), category="query", data=span._data
)
span.finish()
return res
return _inner_end
def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]:
def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
instance = args[0] # type: clickhouse_driver.client.Client
data = args[2]
span = getattr(instance.connection, "_sentry_span", None)
if span is not None:
_set_db_data(span, instance.connection)
if should_send_default_pii():
db_params = span._data.get("db.params", [])
db_params.extend(data)
span.set_data("db.params", db_params)
return f(*args, **kwargs)
return _inner_send_data
def _set_db_data(
span: Span, connection: clickhouse_driver.connection.Connection
) -> None:
span.set_data(SPANDATA.DB_SYSTEM, "clickhouse")
span.set_data(SPANDATA.SERVER_ADDRESS, connection.host)
span.set_data(SPANDATA.SERVER_PORT, connection.port)
span.set_data(SPANDATA.DB_NAME, connection.database)
span.set_data(SPANDATA.DB_USER, connection.user)
@@ -0,0 +1,258 @@
import json
import urllib3
from sentry_sdk.integrations import Integration
from sentry_sdk.api import set_context
from sentry_sdk.utils import logger
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Dict
CONTEXT_TYPE = "cloud_resource"
AWS_METADATA_HOST = "169.254.169.254"
AWS_TOKEN_URL = "http://{}/latest/api/token".format(AWS_METADATA_HOST)
AWS_METADATA_URL = "http://{}/latest/dynamic/instance-identity/document".format(
AWS_METADATA_HOST
)
GCP_METADATA_HOST = "metadata.google.internal"
GCP_METADATA_URL = "http://{}/computeMetadata/v1/?recursive=true".format(
GCP_METADATA_HOST
)
class CLOUD_PROVIDER: # noqa: N801
"""
Name of the cloud provider.
see https://opentelemetry.io/docs/reference/specification/resource/semantic_conventions/cloud/
"""
ALIBABA = "alibaba_cloud"
AWS = "aws"
AZURE = "azure"
GCP = "gcp"
IBM = "ibm_cloud"
TENCENT = "tencent_cloud"
class CLOUD_PLATFORM: # noqa: N801
"""
The cloud platform.
see https://opentelemetry.io/docs/reference/specification/resource/semantic_conventions/cloud/
"""
AWS_EC2 = "aws_ec2"
GCP_COMPUTE_ENGINE = "gcp_compute_engine"
class CloudResourceContextIntegration(Integration):
"""
Adds cloud resource context to the Senty scope
"""
identifier = "cloudresourcecontext"
cloud_provider = ""
aws_token = ""
http = urllib3.PoolManager()
gcp_metadata = None
def __init__(self, cloud_provider=""):
# type: (str) -> None
CloudResourceContextIntegration.cloud_provider = cloud_provider
@classmethod
def _is_aws(cls):
# type: () -> bool
try:
r = cls.http.request(
"PUT",
AWS_TOKEN_URL,
headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
)
if r.status != 200:
return False
cls.aws_token = r.data.decode()
return True
except Exception:
return False
@classmethod
def _get_aws_context(cls):
# type: () -> Dict[str, str]
ctx = {
"cloud.provider": CLOUD_PROVIDER.AWS,
"cloud.platform": CLOUD_PLATFORM.AWS_EC2,
}
try:
r = cls.http.request(
"GET",
AWS_METADATA_URL,
headers={"X-aws-ec2-metadata-token": cls.aws_token},
)
if r.status != 200:
return ctx
data = json.loads(r.data.decode("utf-8"))
try:
ctx["cloud.account.id"] = data["accountId"]
except Exception:
pass
try:
ctx["cloud.availability_zone"] = data["availabilityZone"]
except Exception:
pass
try:
ctx["cloud.region"] = data["region"]
except Exception:
pass
try:
ctx["host.id"] = data["instanceId"]
except Exception:
pass
try:
ctx["host.type"] = data["instanceType"]
except Exception:
pass
except Exception:
pass
return ctx
@classmethod
def _is_gcp(cls):
# type: () -> bool
try:
r = cls.http.request(
"GET",
GCP_METADATA_URL,
headers={"Metadata-Flavor": "Google"},
)
if r.status != 200:
return False
cls.gcp_metadata = json.loads(r.data.decode("utf-8"))
return True
except Exception:
return False
@classmethod
def _get_gcp_context(cls):
# type: () -> Dict[str, str]
ctx = {
"cloud.provider": CLOUD_PROVIDER.GCP,
"cloud.platform": CLOUD_PLATFORM.GCP_COMPUTE_ENGINE,
}
try:
if cls.gcp_metadata is None:
r = cls.http.request(
"GET",
GCP_METADATA_URL,
headers={"Metadata-Flavor": "Google"},
)
if r.status != 200:
return ctx
cls.gcp_metadata = json.loads(r.data.decode("utf-8"))
try:
ctx["cloud.account.id"] = cls.gcp_metadata["project"]["projectId"]
except Exception:
pass
try:
ctx["cloud.availability_zone"] = cls.gcp_metadata["instance"][
"zone"
].split("/")[-1]
except Exception:
pass
try:
# only populated in google cloud run
ctx["cloud.region"] = cls.gcp_metadata["instance"]["region"].split("/")[
-1
]
except Exception:
pass
try:
ctx["host.id"] = cls.gcp_metadata["instance"]["id"]
except Exception:
pass
except Exception:
pass
return ctx
@classmethod
def _get_cloud_provider(cls):
# type: () -> str
if cls._is_aws():
return CLOUD_PROVIDER.AWS
if cls._is_gcp():
return CLOUD_PROVIDER.GCP
return ""
@classmethod
def _get_cloud_resource_context(cls):
# type: () -> Dict[str, str]
cloud_provider = (
cls.cloud_provider
if cls.cloud_provider != ""
else CloudResourceContextIntegration._get_cloud_provider()
)
if cloud_provider in context_getters.keys():
return context_getters[cloud_provider]()
return {}
@staticmethod
def setup_once():
# type: () -> None
cloud_provider = CloudResourceContextIntegration.cloud_provider
unsupported_cloud_provider = (
cloud_provider != "" and cloud_provider not in context_getters.keys()
)
if unsupported_cloud_provider:
logger.warning(
"Invalid value for cloud_provider: %s (must be in %s). Falling back to autodetection...",
CloudResourceContextIntegration.cloud_provider,
list(context_getters.keys()),
)
context = CloudResourceContextIntegration._get_cloud_resource_context()
if context != {}:
set_context(CONTEXT_TYPE, context)
# Map with the currently supported cloud providers
# mapping to functions extracting the context
context_getters = {
CLOUD_PROVIDER.AWS: CloudResourceContextIntegration._get_aws_context,
CLOUD_PROVIDER.GCP: CloudResourceContextIntegration._get_gcp_context,
}
@@ -0,0 +1,270 @@
from functools import wraps
from sentry_sdk import consts
from sentry_sdk.ai.monitoring import record_token_usage
from sentry_sdk.consts import SPANDATA
from sentry_sdk.ai.utils import set_data_normalized
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Iterator
from sentry_sdk.tracing import Span
import sentry_sdk
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.utils import capture_internal_exceptions, event_from_exception
try:
from cohere.client import Client
from cohere.base_client import BaseCohere
from cohere import (
ChatStreamEndEvent,
NonStreamedChatResponse,
)
if TYPE_CHECKING:
from cohere import StreamedChatResponse
except ImportError:
raise DidNotEnable("Cohere not installed")
try:
# cohere 5.9.3+
from cohere import StreamEndStreamedChatResponse
except ImportError:
from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse
COLLECTED_CHAT_PARAMS = {
"model": SPANDATA.AI_MODEL_ID,
"k": SPANDATA.AI_TOP_K,
"p": SPANDATA.AI_TOP_P,
"seed": SPANDATA.AI_SEED,
"frequency_penalty": SPANDATA.AI_FREQUENCY_PENALTY,
"presence_penalty": SPANDATA.AI_PRESENCE_PENALTY,
"raw_prompting": SPANDATA.AI_RAW_PROMPTING,
}
COLLECTED_PII_CHAT_PARAMS = {
"tools": SPANDATA.AI_TOOLS,
"preamble": SPANDATA.AI_PREAMBLE,
}
COLLECTED_CHAT_RESP_ATTRS = {
"generation_id": "ai.generation_id",
"is_search_required": "ai.is_search_required",
"finish_reason": "ai.finish_reason",
}
COLLECTED_PII_CHAT_RESP_ATTRS = {
"citations": "ai.citations",
"documents": "ai.documents",
"search_queries": "ai.search_queries",
"search_results": "ai.search_results",
"tool_calls": "ai.tool_calls",
}
class CohereIntegration(Integration):
identifier = "cohere"
origin = f"auto.ai.{identifier}"
def __init__(self, include_prompts=True):
# type: (CohereIntegration, bool) -> None
self.include_prompts = include_prompts
@staticmethod
def setup_once():
# type: () -> None
BaseCohere.chat = _wrap_chat(BaseCohere.chat, streaming=False)
Client.embed = _wrap_embed(Client.embed)
BaseCohere.chat_stream = _wrap_chat(BaseCohere.chat_stream, streaming=True)
def _capture_exception(exc):
# type: (Any) -> None
event, hint = event_from_exception(
exc,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "cohere", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _wrap_chat(f, streaming):
# type: (Callable[..., Any], bool) -> Callable[..., Any]
def collect_chat_response_fields(span, res, include_pii):
# type: (Span, NonStreamedChatResponse, bool) -> None
if include_pii:
if hasattr(res, "text"):
set_data_normalized(
span,
SPANDATA.AI_RESPONSES,
[res.text],
)
for pii_attr in COLLECTED_PII_CHAT_RESP_ATTRS:
if hasattr(res, pii_attr):
set_data_normalized(span, "ai." + pii_attr, getattr(res, pii_attr))
for attr in COLLECTED_CHAT_RESP_ATTRS:
if hasattr(res, attr):
set_data_normalized(span, "ai." + attr, getattr(res, attr))
if hasattr(res, "meta"):
if hasattr(res.meta, "billed_units"):
record_token_usage(
span,
prompt_tokens=res.meta.billed_units.input_tokens,
completion_tokens=res.meta.billed_units.output_tokens,
)
elif hasattr(res.meta, "tokens"):
record_token_usage(
span,
prompt_tokens=res.meta.tokens.input_tokens,
completion_tokens=res.meta.tokens.output_tokens,
)
if hasattr(res.meta, "warnings"):
set_data_normalized(span, "ai.warnings", res.meta.warnings)
@wraps(f)
def new_chat(*args, **kwargs):
# type: (*Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(CohereIntegration)
if (
integration is None
or "message" not in kwargs
or not isinstance(kwargs.get("message"), str)
):
return f(*args, **kwargs)
message = kwargs.get("message")
span = sentry_sdk.start_span(
op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE,
name="cohere.client.Chat",
origin=CohereIntegration.origin,
)
span.__enter__()
try:
res = f(*args, **kwargs)
except Exception as e:
_capture_exception(e)
span.__exit__(None, None, None)
raise e from None
with capture_internal_exceptions():
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(
span,
SPANDATA.AI_INPUT_MESSAGES,
list(
map(
lambda x: {
"role": getattr(x, "role", "").lower(),
"content": getattr(x, "message", ""),
},
kwargs.get("chat_history", []),
)
)
+ [{"role": "user", "content": message}],
)
for k, v in COLLECTED_PII_CHAT_PARAMS.items():
if k in kwargs:
set_data_normalized(span, v, kwargs[k])
for k, v in COLLECTED_CHAT_PARAMS.items():
if k in kwargs:
set_data_normalized(span, v, kwargs[k])
set_data_normalized(span, SPANDATA.AI_STREAMING, False)
if streaming:
old_iterator = res
def new_iterator():
# type: () -> Iterator[StreamedChatResponse]
with capture_internal_exceptions():
for x in old_iterator:
if isinstance(x, ChatStreamEndEvent) or isinstance(
x, StreamEndStreamedChatResponse
):
collect_chat_response_fields(
span,
x.response,
include_pii=should_send_default_pii()
and integration.include_prompts,
)
yield x
span.__exit__(None, None, None)
return new_iterator()
elif isinstance(res, NonStreamedChatResponse):
collect_chat_response_fields(
span,
res,
include_pii=should_send_default_pii()
and integration.include_prompts,
)
span.__exit__(None, None, None)
else:
set_data_normalized(span, "unknown_response", True)
span.__exit__(None, None, None)
return res
return new_chat
def _wrap_embed(f):
# type: (Callable[..., Any]) -> Callable[..., Any]
@wraps(f)
def new_embed(*args, **kwargs):
# type: (*Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(CohereIntegration)
if integration is None:
return f(*args, **kwargs)
with sentry_sdk.start_span(
op=consts.OP.COHERE_EMBEDDINGS_CREATE,
name="Cohere Embedding Creation",
origin=CohereIntegration.origin,
) as span:
if "texts" in kwargs and (
should_send_default_pii() and integration.include_prompts
):
if isinstance(kwargs["texts"], str):
set_data_normalized(span, "ai.texts", [kwargs["texts"]])
elif (
isinstance(kwargs["texts"], list)
and len(kwargs["texts"]) > 0
and isinstance(kwargs["texts"][0], str)
):
set_data_normalized(
span, SPANDATA.AI_INPUT_MESSAGES, kwargs["texts"]
)
if "model" in kwargs:
set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"])
try:
res = f(*args, **kwargs)
except Exception as e:
_capture_exception(e)
raise e from None
if (
hasattr(res, "meta")
and hasattr(res.meta, "billed_units")
and hasattr(res.meta.billed_units, "input_tokens")
):
record_token_usage(
span,
prompt_tokens=res.meta.billed_units.input_tokens,
total_tokens=res.meta.billed_units.input_tokens,
)
return res
return new_embed
@@ -0,0 +1,42 @@
import sentry_sdk
from sentry_sdk.utils import ContextVar
from sentry_sdk.integrations import Integration
from sentry_sdk.scope import add_global_event_processor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional
from sentry_sdk._types import Event, Hint
class DedupeIntegration(Integration):
identifier = "dedupe"
def __init__(self):
# type: () -> None
self._last_seen = ContextVar("last-seen")
@staticmethod
def setup_once():
# type: () -> None
@add_global_event_processor
def processor(event, hint):
# type: (Event, Optional[Hint]) -> Optional[Event]
if hint is None:
return event
integration = sentry_sdk.get_client().get_integration(DedupeIntegration)
if integration is None:
return event
exc_info = hint.get("exc_info", None)
if exc_info is None:
return event
exc = exc_info[1]
if integration._last_seen.get(None) is exc:
return None
integration._last_seen.set(exc)
return event
@@ -0,0 +1,747 @@
import inspect
import sys
import threading
import weakref
from importlib import import_module
import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.scope import add_global_event_processor, should_send_default_pii
from sentry_sdk.serializer import add_global_repr_processor
from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_URL
from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
from sentry_sdk.utils import (
AnnotatedValue,
HAS_REAL_CONTEXTVARS,
CONTEXTVARS_ERROR_MESSAGE,
SENSITIVE_DATA_SUBSTITUTE,
logger,
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
transaction_from_function,
walk_exception_chain,
)
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
from sentry_sdk.integrations._wsgi_common import (
DEFAULT_HTTP_METHODS_TO_CAPTURE,
RequestExtractor,
)
try:
from django import VERSION as DJANGO_VERSION
from django.conf import settings as django_settings
from django.core import signals
from django.conf import settings
try:
from django.urls import resolve
except ImportError:
from django.core.urlresolvers import resolve
try:
from django.urls import Resolver404
except ImportError:
from django.core.urlresolvers import Resolver404
# Only available in Django 3.0+
try:
from django.core.handlers.asgi import ASGIRequest
except Exception:
ASGIRequest = None
except ImportError:
raise DidNotEnable("Django not installed")
from sentry_sdk.integrations.django.transactions import LEGACY_RESOLVER
from sentry_sdk.integrations.django.templates import (
get_template_frame_from_exception,
patch_templates,
)
from sentry_sdk.integrations.django.middleware import patch_django_middlewares
from sentry_sdk.integrations.django.signals_handlers import patch_signals
from sentry_sdk.integrations.django.views import patch_views
if DJANGO_VERSION[:2] > (1, 8):
from sentry_sdk.integrations.django.caching import patch_caching
else:
patch_caching = None # type: ignore
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import Dict
from typing import Optional
from typing import Union
from typing import List
from django.core.handlers.wsgi import WSGIRequest
from django.http.response import HttpResponse
from django.http.request import QueryDict
from django.utils.datastructures import MultiValueDict
from sentry_sdk.tracing import Span
from sentry_sdk.integrations.wsgi import _ScopedResponse
from sentry_sdk._types import Event, Hint, EventProcessor, NotImplementedType
if DJANGO_VERSION < (1, 10):
def is_authenticated(request_user):
# type: (Any) -> bool
return request_user.is_authenticated()
else:
def is_authenticated(request_user):
# type: (Any) -> bool
return request_user.is_authenticated
TRANSACTION_STYLE_VALUES = ("function_name", "url")
class DjangoIntegration(Integration):
"""
Auto instrument a Django application.
:param transaction_style: How to derive transaction names. Either `"function_name"` or `"url"`. Defaults to `"url"`.
:param middleware_spans: Whether to create spans for middleware. Defaults to `True`.
:param signals_spans: Whether to create spans for signals. Defaults to `True`.
:param signals_denylist: A list of signals to ignore when creating spans.
:param cache_spans: Whether to create spans for cache operations. Defaults to `False`.
"""
identifier = "django"
origin = f"auto.http.{identifier}"
origin_db = f"auto.db.{identifier}"
transaction_style = ""
middleware_spans = None
signals_spans = None
cache_spans = None
signals_denylist = [] # type: list[signals.Signal]
def __init__(
self,
transaction_style="url", # type: str
middleware_spans=True, # type: bool
signals_spans=True, # type: bool
cache_spans=False, # type: bool
signals_denylist=None, # type: Optional[list[signals.Signal]]
http_methods_to_capture=DEFAULT_HTTP_METHODS_TO_CAPTURE, # type: tuple[str, ...]
):
# type: (...) -> None
if transaction_style not in TRANSACTION_STYLE_VALUES:
raise ValueError(
"Invalid value for transaction_style: %s (must be in %s)"
% (transaction_style, TRANSACTION_STYLE_VALUES)
)
self.transaction_style = transaction_style
self.middleware_spans = middleware_spans
self.signals_spans = signals_spans
self.signals_denylist = signals_denylist or []
self.cache_spans = cache_spans
self.http_methods_to_capture = tuple(map(str.upper, http_methods_to_capture))
@staticmethod
def setup_once():
# type: () -> None
_check_minimum_version(DjangoIntegration, DJANGO_VERSION)
install_sql_hook()
# Patch in our custom middleware.
# logs an error for every 500
ignore_logger("django.server")
ignore_logger("django.request")
from django.core.handlers.wsgi import WSGIHandler
old_app = WSGIHandler.__call__
@ensure_integration_enabled(DjangoIntegration, old_app)
def sentry_patched_wsgi_handler(self, environ, start_response):
# type: (Any, Dict[str, str], Callable[..., Any]) -> _ScopedResponse
bound_old_app = old_app.__get__(self, WSGIHandler)
from django.conf import settings
use_x_forwarded_for = settings.USE_X_FORWARDED_HOST
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
middleware = SentryWsgiMiddleware(
bound_old_app,
use_x_forwarded_for,
span_origin=DjangoIntegration.origin,
http_methods_to_capture=(
integration.http_methods_to_capture
if integration
else DEFAULT_HTTP_METHODS_TO_CAPTURE
),
)
return middleware(environ, start_response)
WSGIHandler.__call__ = sentry_patched_wsgi_handler
_patch_get_response()
_patch_django_asgi_handler()
signals.got_request_exception.connect(_got_request_exception)
@add_global_event_processor
def process_django_templates(event, hint):
# type: (Event, Optional[Hint]) -> Optional[Event]
if hint is None:
return event
exc_info = hint.get("exc_info", None)
if exc_info is None:
return event
exception = event.get("exception", None)
if exception is None:
return event
values = exception.get("values", None)
if values is None:
return event
for exception, (_, exc_value, _) in zip(
reversed(values), walk_exception_chain(exc_info)
):
frame = get_template_frame_from_exception(exc_value)
if frame is not None:
frames = exception.get("stacktrace", {}).get("frames", [])
for i in reversed(range(len(frames))):
f = frames[i]
if (
f.get("function") in ("Parser.parse", "parse", "render")
and f.get("module") == "django.template.base"
):
i += 1
break
else:
i = len(frames)
frames.insert(i, frame)
return event
@add_global_repr_processor
def _django_queryset_repr(value, hint):
# type: (Any, Dict[str, Any]) -> Union[NotImplementedType, str]
try:
# Django 1.6 can fail to import `QuerySet` when Django settings
# have not yet been initialized.
#
# If we fail to import, return `NotImplemented`. It's at least
# unlikely that we have a query set in `value` when importing
# `QuerySet` fails.
from django.db.models.query import QuerySet
except Exception:
return NotImplemented
if not isinstance(value, QuerySet) or value._result_cache:
return NotImplemented
return "<%s from %s at 0x%x>" % (
value.__class__.__name__,
value.__module__,
id(value),
)
_patch_channels()
patch_django_middlewares()
patch_views()
patch_templates()
patch_signals()
if patch_caching is not None:
patch_caching()
_DRF_PATCHED = False
_DRF_PATCH_LOCK = threading.Lock()
def _patch_drf():
# type: () -> None
"""
Patch Django Rest Framework for more/better request data. DRF's request
type is a wrapper around Django's request type. The attribute we're
interested in is `request.data`, which is a cached property containing a
parsed request body. Reading a request body from that property is more
reliable than reading from any of Django's own properties, as those don't
hold payloads in memory and therefore can only be accessed once.
We patch the Django request object to include a weak backreference to the
DRF request object, such that we can later use either in
`DjangoRequestExtractor`.
This function is not called directly on SDK setup, because importing almost
any part of Django Rest Framework will try to access Django settings (where
`sentry_sdk.init()` might be called from in the first place). Instead we
run this function on every request and do the patching on the first
request.
"""
global _DRF_PATCHED
if _DRF_PATCHED:
# Double-checked locking
return
with _DRF_PATCH_LOCK:
if _DRF_PATCHED:
return
# We set this regardless of whether the code below succeeds or fails.
# There is no point in trying to patch again on the next request.
_DRF_PATCHED = True
with capture_internal_exceptions():
try:
from rest_framework.views import APIView # type: ignore
except ImportError:
pass
else:
old_drf_initial = APIView.initial
def sentry_patched_drf_initial(self, request, *args, **kwargs):
# type: (APIView, Any, *Any, **Any) -> Any
with capture_internal_exceptions():
request._request._sentry_drf_request_backref = weakref.ref(
request
)
pass
return old_drf_initial(self, request, *args, **kwargs)
APIView.initial = sentry_patched_drf_initial
def _patch_channels():
# type: () -> None
try:
from channels.http import AsgiHandler # type: ignore
except ImportError:
return
if not HAS_REAL_CONTEXTVARS:
# We better have contextvars or we're going to leak state between
# requests.
#
# We cannot hard-raise here because channels may not be used at all in
# the current process. That is the case when running traditional WSGI
# workers in gunicorn+gevent and the websocket stuff in a separate
# process.
logger.warning(
"We detected that you are using Django channels 2.0."
+ CONTEXTVARS_ERROR_MESSAGE
)
from sentry_sdk.integrations.django.asgi import patch_channels_asgi_handler_impl
patch_channels_asgi_handler_impl(AsgiHandler)
def _patch_django_asgi_handler():
# type: () -> None
try:
from django.core.handlers.asgi import ASGIHandler
except ImportError:
return
if not HAS_REAL_CONTEXTVARS:
# We better have contextvars or we're going to leak state between
# requests.
#
# We cannot hard-raise here because Django's ASGI stuff may not be used
# at all.
logger.warning(
"We detected that you are using Django 3." + CONTEXTVARS_ERROR_MESSAGE
)
from sentry_sdk.integrations.django.asgi import patch_django_asgi_handler_impl
patch_django_asgi_handler_impl(ASGIHandler)
def _set_transaction_name_and_source(scope, transaction_style, request):
# type: (sentry_sdk.Scope, str, WSGIRequest) -> None
try:
transaction_name = None
if transaction_style == "function_name":
fn = resolve(request.path).func
transaction_name = transaction_from_function(getattr(fn, "view_class", fn))
elif transaction_style == "url":
if hasattr(request, "urlconf"):
transaction_name = LEGACY_RESOLVER.resolve(
request.path_info, urlconf=request.urlconf
)
else:
transaction_name = LEGACY_RESOLVER.resolve(request.path_info)
if transaction_name is None:
transaction_name = request.path_info
source = TRANSACTION_SOURCE_URL
else:
source = SOURCE_FOR_STYLE[transaction_style]
scope.set_transaction_name(
transaction_name,
source=source,
)
except Resolver404:
urlconf = import_module(settings.ROOT_URLCONF)
# This exception only gets thrown when transaction_style is `function_name`
# So we don't check here what style is configured
if hasattr(urlconf, "handler404"):
handler = urlconf.handler404
if isinstance(handler, str):
scope.transaction = handler
else:
scope.transaction = transaction_from_function(
getattr(handler, "view_class", handler)
)
except Exception:
pass
def _before_get_response(request):
# type: (WSGIRequest) -> None
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is None:
return
_patch_drf()
scope = sentry_sdk.get_current_scope()
# Rely on WSGI middleware to start a trace
_set_transaction_name_and_source(scope, integration.transaction_style, request)
scope.add_event_processor(
_make_wsgi_request_event_processor(weakref.ref(request), integration)
)
def _attempt_resolve_again(request, scope, transaction_style):
# type: (WSGIRequest, sentry_sdk.Scope, str) -> None
"""
Some django middlewares overwrite request.urlconf
so we need to respect that contract,
so we try to resolve the url again.
"""
if not hasattr(request, "urlconf"):
return
_set_transaction_name_and_source(scope, transaction_style, request)
def _after_get_response(request):
# type: (WSGIRequest) -> None
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is None or integration.transaction_style != "url":
return
scope = sentry_sdk.get_current_scope()
_attempt_resolve_again(request, scope, integration.transaction_style)
def _patch_get_response():
# type: () -> None
"""
patch get_response, because at that point we have the Django request object
"""
from django.core.handlers.base import BaseHandler
old_get_response = BaseHandler.get_response
def sentry_patched_get_response(self, request):
# type: (Any, WSGIRequest) -> Union[HttpResponse, BaseException]
_before_get_response(request)
rv = old_get_response(self, request)
_after_get_response(request)
return rv
BaseHandler.get_response = sentry_patched_get_response
if hasattr(BaseHandler, "get_response_async"):
from sentry_sdk.integrations.django.asgi import patch_get_response_async
patch_get_response_async(BaseHandler, _before_get_response)
def _make_wsgi_request_event_processor(weak_request, integration):
# type: (Callable[[], WSGIRequest], DjangoIntegration) -> EventProcessor
def wsgi_request_event_processor(event, hint):
# type: (Event, dict[str, Any]) -> Event
# if the request is gone we are fine not logging the data from
# it. This might happen if the processor is pushed away to
# another thread.
request = weak_request()
if request is None:
return event
django_3 = ASGIRequest is not None
if django_3 and type(request) == ASGIRequest:
# We have a `asgi_request_event_processor` for this.
return event
with capture_internal_exceptions():
DjangoRequestExtractor(request).extract_into_event(event)
if should_send_default_pii():
with capture_internal_exceptions():
_set_user_info(request, event)
return event
return wsgi_request_event_processor
def _got_request_exception(request=None, **kwargs):
# type: (WSGIRequest, **Any) -> None
client = sentry_sdk.get_client()
integration = client.get_integration(DjangoIntegration)
if integration is None:
return
if request is not None and integration.transaction_style == "url":
scope = sentry_sdk.get_current_scope()
_attempt_resolve_again(request, scope, integration.transaction_style)
event, hint = event_from_exception(
sys.exc_info(),
client_options=client.options,
mechanism={"type": "django", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
class DjangoRequestExtractor(RequestExtractor):
def __init__(self, request):
# type: (Union[WSGIRequest, ASGIRequest]) -> None
try:
drf_request = request._sentry_drf_request_backref()
if drf_request is not None:
request = drf_request
except AttributeError:
pass
self.request = request
def env(self):
# type: () -> Dict[str, str]
return self.request.META
def cookies(self):
# type: () -> Dict[str, Union[str, AnnotatedValue]]
privacy_cookies = [
django_settings.CSRF_COOKIE_NAME,
django_settings.SESSION_COOKIE_NAME,
]
clean_cookies = {} # type: Dict[str, Union[str, AnnotatedValue]]
for key, val in self.request.COOKIES.items():
if key in privacy_cookies:
clean_cookies[key] = SENSITIVE_DATA_SUBSTITUTE
else:
clean_cookies[key] = val
return clean_cookies
def raw_data(self):
# type: () -> bytes
return self.request.body
def form(self):
# type: () -> QueryDict
return self.request.POST
def files(self):
# type: () -> MultiValueDict
return self.request.FILES
def size_of_file(self, file):
# type: (Any) -> int
return file.size
def parsed_body(self):
# type: () -> Optional[Dict[str, Any]]
try:
return self.request.data
except AttributeError:
return RequestExtractor.parsed_body(self)
def _set_user_info(request, event):
# type: (WSGIRequest, Event) -> None
user_info = event.setdefault("user", {})
user = getattr(request, "user", None)
if user is None or not is_authenticated(user):
return
try:
user_info.setdefault("id", str(user.pk))
except Exception:
pass
try:
user_info.setdefault("email", user.email)
except Exception:
pass
try:
user_info.setdefault("username", user.get_username())
except Exception:
pass
def install_sql_hook():
# type: () -> None
"""If installed this causes Django's queries to be captured."""
try:
from django.db.backends.utils import CursorWrapper
except ImportError:
from django.db.backends.util import CursorWrapper
try:
# django 1.6 and 1.7 compatability
from django.db.backends import BaseDatabaseWrapper
except ImportError:
# django 1.8 or later
from django.db.backends.base.base import BaseDatabaseWrapper
try:
real_execute = CursorWrapper.execute
real_executemany = CursorWrapper.executemany
real_connect = BaseDatabaseWrapper.connect
except AttributeError:
# This won't work on Django versions < 1.6
return
@ensure_integration_enabled(DjangoIntegration, real_execute)
def execute(self, sql, params=None):
# type: (CursorWrapper, Any, Optional[Any]) -> Any
with record_sql_queries(
cursor=self.cursor,
query=sql,
params_list=params,
paramstyle="format",
executemany=False,
span_origin=DjangoIntegration.origin_db,
) as span:
_set_db_data(span, self)
result = real_execute(self, sql, params)
with capture_internal_exceptions():
add_query_source(span)
return result
@ensure_integration_enabled(DjangoIntegration, real_executemany)
def executemany(self, sql, param_list):
# type: (CursorWrapper, Any, List[Any]) -> Any
with record_sql_queries(
cursor=self.cursor,
query=sql,
params_list=param_list,
paramstyle="format",
executemany=True,
span_origin=DjangoIntegration.origin_db,
) as span:
_set_db_data(span, self)
result = real_executemany(self, sql, param_list)
with capture_internal_exceptions():
add_query_source(span)
return result
@ensure_integration_enabled(DjangoIntegration, real_connect)
def connect(self):
# type: (BaseDatabaseWrapper) -> None
with capture_internal_exceptions():
sentry_sdk.add_breadcrumb(message="connect", category="query")
with sentry_sdk.start_span(
op=OP.DB,
name="connect",
origin=DjangoIntegration.origin_db,
) as span:
_set_db_data(span, self)
return real_connect(self)
CursorWrapper.execute = execute
CursorWrapper.executemany = executemany
BaseDatabaseWrapper.connect = connect
ignore_logger("django.db.backends")
def _set_db_data(span, cursor_or_db):
# type: (Span, Any) -> None
db = cursor_or_db.db if hasattr(cursor_or_db, "db") else cursor_or_db
vendor = db.vendor
span.set_data(SPANDATA.DB_SYSTEM, vendor)
# Some custom backends override `__getattr__`, making it look like `cursor_or_db`
# actually has a `connection` and the `connection` has a `get_dsn_parameters`
# attribute, only to throw an error once you actually want to call it.
# Hence the `inspect` check whether `get_dsn_parameters` is an actual callable
# function.
is_psycopg2 = (
hasattr(cursor_or_db, "connection")
and hasattr(cursor_or_db.connection, "get_dsn_parameters")
and inspect.isroutine(cursor_or_db.connection.get_dsn_parameters)
)
if is_psycopg2:
connection_params = cursor_or_db.connection.get_dsn_parameters()
else:
try:
# psycopg3, only extract needed params as get_parameters
# can be slow because of the additional logic to filter out default
# values
connection_params = {
"dbname": cursor_or_db.connection.info.dbname,
"port": cursor_or_db.connection.info.port,
}
# PGhost returns host or base dir of UNIX socket as an absolute path
# starting with /, use it only when it contains host
pg_host = cursor_or_db.connection.info.host
if pg_host and not pg_host.startswith("/"):
connection_params["host"] = pg_host
except Exception:
connection_params = db.get_connection_params()
db_name = connection_params.get("dbname") or connection_params.get("database")
if db_name is not None:
span.set_data(SPANDATA.DB_NAME, db_name)
server_address = connection_params.get("host")
if server_address is not None:
span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
server_port = connection_params.get("port")
if server_port is not None:
span.set_data(SPANDATA.SERVER_PORT, str(server_port))
server_socket_address = connection_params.get("unix_socket")
if server_socket_address is not None:
span.set_data(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)
@@ -0,0 +1,245 @@
"""
Instrumentation for Django 3.0
Since this file contains `async def` it is conditionally imported in
`sentry_sdk.integrations.django` (depending on the existence of
`django.core.handlers.asgi`.
"""
import asyncio
import functools
import inspect
from django.core.handlers.wsgi import WSGIRequest
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Union, TypeVar
from django.core.handlers.asgi import ASGIRequest
from django.http.response import HttpResponse
from sentry_sdk._types import Event, EventProcessor
_F = TypeVar("_F", bound=Callable[..., Any])
# Python 3.12 deprecates asyncio.iscoroutinefunction() as an alias for
# inspect.iscoroutinefunction(), whilst also removing the _is_coroutine marker.
# The latter is replaced with the inspect.markcoroutinefunction decorator.
# Until 3.12 is the minimum supported Python version, provide a shim.
# This was copied from https://github.com/django/asgiref/blob/main/asgiref/sync.py
if hasattr(inspect, "markcoroutinefunction"):
iscoroutinefunction = inspect.iscoroutinefunction
markcoroutinefunction = inspect.markcoroutinefunction
else:
iscoroutinefunction = asyncio.iscoroutinefunction # type: ignore[assignment]
def markcoroutinefunction(func: "_F") -> "_F":
func._is_coroutine = asyncio.coroutines._is_coroutine # type: ignore
return func
def _make_asgi_request_event_processor(request):
# type: (ASGIRequest) -> EventProcessor
def asgi_request_event_processor(event, hint):
# type: (Event, dict[str, Any]) -> Event
# if the request is gone we are fine not logging the data from
# it. This might happen if the processor is pushed away to
# another thread.
from sentry_sdk.integrations.django import (
DjangoRequestExtractor,
_set_user_info,
)
if request is None:
return event
if type(request) == WSGIRequest:
return event
with capture_internal_exceptions():
DjangoRequestExtractor(request).extract_into_event(event)
if should_send_default_pii():
with capture_internal_exceptions():
_set_user_info(request, event)
return event
return asgi_request_event_processor
def patch_django_asgi_handler_impl(cls):
# type: (Any) -> None
from sentry_sdk.integrations.django import DjangoIntegration
old_app = cls.__call__
async def sentry_patched_asgi_handler(self, scope, receive, send):
# type: (Any, Any, Any, Any) -> Any
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is None:
return await old_app(self, scope, receive, send)
middleware = SentryAsgiMiddleware(
old_app.__get__(self, cls),
unsafe_context_data=True,
span_origin=DjangoIntegration.origin,
http_methods_to_capture=integration.http_methods_to_capture,
)._run_asgi3
return await middleware(scope, receive, send)
cls.__call__ = sentry_patched_asgi_handler
modern_django_asgi_support = hasattr(cls, "create_request")
if modern_django_asgi_support:
old_create_request = cls.create_request
@ensure_integration_enabled(DjangoIntegration, old_create_request)
def sentry_patched_create_request(self, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
request, error_response = old_create_request(self, *args, **kwargs)
scope = sentry_sdk.get_isolation_scope()
scope.add_event_processor(_make_asgi_request_event_processor(request))
return request, error_response
cls.create_request = sentry_patched_create_request
def patch_get_response_async(cls, _before_get_response):
# type: (Any, Any) -> None
old_get_response_async = cls.get_response_async
async def sentry_patched_get_response_async(self, request):
# type: (Any, Any) -> Union[HttpResponse, BaseException]
_before_get_response(request)
return await old_get_response_async(self, request)
cls.get_response_async = sentry_patched_get_response_async
def patch_channels_asgi_handler_impl(cls):
# type: (Any) -> None
import channels # type: ignore
from sentry_sdk.integrations.django import DjangoIntegration
if channels.__version__ < "3.0.0":
old_app = cls.__call__
async def sentry_patched_asgi_handler(self, receive, send):
# type: (Any, Any, Any) -> Any
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is None:
return await old_app(self, receive, send)
middleware = SentryAsgiMiddleware(
lambda _scope: old_app.__get__(self, cls),
unsafe_context_data=True,
span_origin=DjangoIntegration.origin,
http_methods_to_capture=integration.http_methods_to_capture,
)
return await middleware(self.scope)(receive, send)
cls.__call__ = sentry_patched_asgi_handler
else:
# The ASGI handler in Channels >= 3 has the same signature as
# the Django handler.
patch_django_asgi_handler_impl(cls)
def wrap_async_view(callback):
# type: (Any) -> Any
from sentry_sdk.integrations.django import DjangoIntegration
@functools.wraps(callback)
async def sentry_wrapped_callback(request, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
current_scope = sentry_sdk.get_current_scope()
if current_scope.transaction is not None:
current_scope.transaction.update_active_thread()
sentry_scope = sentry_sdk.get_isolation_scope()
if sentry_scope.profile is not None:
sentry_scope.profile.update_active_thread_id()
with sentry_sdk.start_span(
op=OP.VIEW_RENDER,
name=request.resolver_match.view_name,
origin=DjangoIntegration.origin,
):
return await callback(request, *args, **kwargs)
return sentry_wrapped_callback
def _asgi_middleware_mixin_factory(_check_middleware_span):
# type: (Callable[..., Any]) -> Any
"""
Mixin class factory that generates a middleware mixin for handling requests
in async mode.
"""
class SentryASGIMixin:
if TYPE_CHECKING:
_inner = None
def __init__(self, get_response):
# type: (Callable[..., Any]) -> None
self.get_response = get_response
self._acall_method = None
self._async_check()
def _async_check(self):
# type: () -> None
"""
If get_response is a coroutine function, turns us into async mode so
a thread is not consumed during a whole request.
Taken from django.utils.deprecation::MiddlewareMixin._async_check
"""
if iscoroutinefunction(self.get_response):
markcoroutinefunction(self)
def async_route_check(self):
# type: () -> bool
"""
Function that checks if we are in async mode,
and if we are forwards the handling of requests to __acall__
"""
return iscoroutinefunction(self.get_response)
async def __acall__(self, *args, **kwargs):
# type: (*Any, **Any) -> Any
f = self._acall_method
if f is None:
if hasattr(self._inner, "__acall__"):
self._acall_method = f = self._inner.__acall__ # type: ignore
else:
self._acall_method = f = self._inner
middleware_span = _check_middleware_span(old_method=f)
if middleware_span is None:
return await f(*args, **kwargs)
with middleware_span:
return await f(*args, **kwargs)
return SentryASGIMixin
@@ -0,0 +1,191 @@
import functools
from typing import TYPE_CHECKING
from sentry_sdk.integrations.redis.utils import _get_safe_key, _key_as_string
from urllib3.util import parse_url as urlparse
from django import VERSION as DJANGO_VERSION
from django.core.cache import CacheHandler
import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
)
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import Optional
METHODS_TO_INSTRUMENT = [
"set",
"set_many",
"get",
"get_many",
]
def _get_span_description(method_name, args, kwargs):
# type: (str, tuple[Any], dict[str, Any]) -> str
return _key_as_string(_get_safe_key(method_name, args, kwargs))
def _patch_cache_method(cache, method_name, address, port):
# type: (CacheHandler, str, Optional[str], Optional[int]) -> None
from sentry_sdk.integrations.django import DjangoIntegration
original_method = getattr(cache, method_name)
@ensure_integration_enabled(DjangoIntegration, original_method)
def _instrument_call(
cache, method_name, original_method, args, kwargs, address, port
):
# type: (CacheHandler, str, Callable[..., Any], tuple[Any, ...], dict[str, Any], Optional[str], Optional[int]) -> Any
is_set_operation = method_name.startswith("set")
is_get_operation = not is_set_operation
op = OP.CACHE_PUT if is_set_operation else OP.CACHE_GET
description = _get_span_description(method_name, args, kwargs)
with sentry_sdk.start_span(
op=op,
name=description,
origin=DjangoIntegration.origin,
) as span:
value = original_method(*args, **kwargs)
with capture_internal_exceptions():
if address is not None:
span.set_data(SPANDATA.NETWORK_PEER_ADDRESS, address)
if port is not None:
span.set_data(SPANDATA.NETWORK_PEER_PORT, port)
key = _get_safe_key(method_name, args, kwargs)
if key is not None:
span.set_data(SPANDATA.CACHE_KEY, key)
item_size = None
if is_get_operation:
if value:
item_size = len(str(value))
span.set_data(SPANDATA.CACHE_HIT, True)
else:
span.set_data(SPANDATA.CACHE_HIT, False)
else: # TODO: We don't handle `get_or_set` which we should
arg_count = len(args)
if arg_count >= 2:
# 'set' command
item_size = len(str(args[1]))
elif arg_count == 1:
# 'set_many' command
item_size = len(str(args[0]))
if item_size is not None:
span.set_data(SPANDATA.CACHE_ITEM_SIZE, item_size)
return value
@functools.wraps(original_method)
def sentry_method(*args, **kwargs):
# type: (*Any, **Any) -> Any
return _instrument_call(
cache, method_name, original_method, args, kwargs, address, port
)
setattr(cache, method_name, sentry_method)
def _patch_cache(cache, address=None, port=None):
# type: (CacheHandler, Optional[str], Optional[int]) -> None
if not hasattr(cache, "_sentry_patched"):
for method_name in METHODS_TO_INSTRUMENT:
_patch_cache_method(cache, method_name, address, port)
cache._sentry_patched = True
def _get_address_port(settings):
# type: (dict[str, Any]) -> tuple[Optional[str], Optional[int]]
location = settings.get("LOCATION")
# TODO: location can also be an array of locations
# see: https://docs.djangoproject.com/en/5.0/topics/cache/#redis
# GitHub issue: https://github.com/getsentry/sentry-python/issues/3062
if not isinstance(location, str):
return None, None
if "://" in location:
parsed_url = urlparse(location)
# remove the username and password from URL to not leak sensitive data.
address = "{}://{}{}".format(
parsed_url.scheme or "",
parsed_url.hostname or "",
parsed_url.path or "",
)
port = parsed_url.port
else:
address = location
port = None
return address, int(port) if port is not None else None
def should_enable_cache_spans():
# type: () -> bool
from sentry_sdk.integrations.django import DjangoIntegration
client = sentry_sdk.get_client()
integration = client.get_integration(DjangoIntegration)
from django.conf import settings
return integration is not None and (
(client.spotlight is not None and settings.DEBUG is True)
or integration.cache_spans is True
)
def patch_caching():
# type: () -> None
if not hasattr(CacheHandler, "_sentry_patched"):
if DJANGO_VERSION < (3, 2):
original_get_item = CacheHandler.__getitem__
@functools.wraps(original_get_item)
def sentry_get_item(self, alias):
# type: (CacheHandler, str) -> Any
cache = original_get_item(self, alias)
if should_enable_cache_spans():
from django.conf import settings
address, port = _get_address_port(
settings.CACHES[alias or "default"]
)
_patch_cache(cache, address, port)
return cache
CacheHandler.__getitem__ = sentry_get_item
CacheHandler._sentry_patched = True
else:
original_create_connection = CacheHandler.create_connection
@functools.wraps(original_create_connection)
def sentry_create_connection(self, alias):
# type: (CacheHandler, str) -> Any
cache = original_create_connection(self, alias)
if should_enable_cache_spans():
address, port = _get_address_port(self.settings[alias or "default"])
_patch_cache(cache, address, port)
return cache
CacheHandler.create_connection = sentry_create_connection
CacheHandler._sentry_patched = True
@@ -0,0 +1,187 @@
"""
Create spans from Django middleware invocations
"""
from functools import wraps
from django import VERSION as DJANGO_VERSION
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.utils import (
ContextVar,
transaction_from_function,
capture_internal_exceptions,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Callable
from typing import Optional
from typing import TypeVar
from sentry_sdk.tracing import Span
F = TypeVar("F", bound=Callable[..., Any])
_import_string_should_wrap_middleware = ContextVar(
"import_string_should_wrap_middleware"
)
DJANGO_SUPPORTS_ASYNC_MIDDLEWARE = DJANGO_VERSION >= (3, 1)
if not DJANGO_SUPPORTS_ASYNC_MIDDLEWARE:
_asgi_middleware_mixin_factory = lambda _: object
else:
from .asgi import _asgi_middleware_mixin_factory
def patch_django_middlewares():
# type: () -> None
from django.core.handlers import base
old_import_string = base.import_string
def sentry_patched_import_string(dotted_path):
# type: (str) -> Any
rv = old_import_string(dotted_path)
if _import_string_should_wrap_middleware.get(None):
rv = _wrap_middleware(rv, dotted_path)
return rv
base.import_string = sentry_patched_import_string
old_load_middleware = base.BaseHandler.load_middleware
def sentry_patched_load_middleware(*args, **kwargs):
# type: (Any, Any) -> Any
_import_string_should_wrap_middleware.set(True)
try:
return old_load_middleware(*args, **kwargs)
finally:
_import_string_should_wrap_middleware.set(False)
base.BaseHandler.load_middleware = sentry_patched_load_middleware
def _wrap_middleware(middleware, middleware_name):
# type: (Any, str) -> Any
from sentry_sdk.integrations.django import DjangoIntegration
def _check_middleware_span(old_method):
# type: (Callable[..., Any]) -> Optional[Span]
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is None or not integration.middleware_spans:
return None
function_name = transaction_from_function(old_method)
description = middleware_name
function_basename = getattr(old_method, "__name__", None)
if function_basename:
description = "{}.{}".format(description, function_basename)
middleware_span = sentry_sdk.start_span(
op=OP.MIDDLEWARE_DJANGO,
name=description,
origin=DjangoIntegration.origin,
)
middleware_span.set_tag("django.function_name", function_name)
middleware_span.set_tag("django.middleware_name", middleware_name)
return middleware_span
def _get_wrapped_method(old_method):
# type: (F) -> F
with capture_internal_exceptions():
def sentry_wrapped_method(*args, **kwargs):
# type: (*Any, **Any) -> Any
middleware_span = _check_middleware_span(old_method)
if middleware_span is None:
return old_method(*args, **kwargs)
with middleware_span:
return old_method(*args, **kwargs)
try:
# fails for __call__ of function on Python 2 (see py2.7-django-1.11)
sentry_wrapped_method = wraps(old_method)(sentry_wrapped_method)
# Necessary for Django 3.1
sentry_wrapped_method.__self__ = old_method.__self__ # type: ignore
except Exception:
pass
return sentry_wrapped_method # type: ignore
return old_method
class SentryWrappingMiddleware(
_asgi_middleware_mixin_factory(_check_middleware_span) # type: ignore
):
sync_capable = getattr(middleware, "sync_capable", True)
async_capable = DJANGO_SUPPORTS_ASYNC_MIDDLEWARE and getattr(
middleware, "async_capable", False
)
def __init__(self, get_response=None, *args, **kwargs):
# type: (Optional[Callable[..., Any]], *Any, **Any) -> None
if get_response:
self._inner = middleware(get_response, *args, **kwargs)
else:
self._inner = middleware(*args, **kwargs)
self.get_response = get_response
self._call_method = None
if self.async_capable:
super().__init__(get_response)
# We need correct behavior for `hasattr()`, which we can only determine
# when we have an instance of the middleware we're wrapping.
def __getattr__(self, method_name):
# type: (str) -> Any
if method_name not in (
"process_request",
"process_view",
"process_template_response",
"process_response",
"process_exception",
):
raise AttributeError()
old_method = getattr(self._inner, method_name)
rv = _get_wrapped_method(old_method)
self.__dict__[method_name] = rv
return rv
def __call__(self, *args, **kwargs):
# type: (*Any, **Any) -> Any
if hasattr(self, "async_route_check") and self.async_route_check():
return self.__acall__(*args, **kwargs)
f = self._call_method
if f is None:
self._call_method = f = self._inner.__call__
middleware_span = _check_middleware_span(old_method=f)
if middleware_span is None:
return f(*args, **kwargs)
with middleware_span:
return f(*args, **kwargs)
for attr in (
"__name__",
"__module__",
"__qualname__",
):
if hasattr(middleware, attr):
setattr(SentryWrappingMiddleware, attr, getattr(middleware, attr))
return SentryWrappingMiddleware
@@ -0,0 +1,91 @@
from functools import wraps
from django.dispatch import Signal
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations.django import DJANGO_VERSION
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Callable
from typing import Any, Union
def _get_receiver_name(receiver):
# type: (Callable[..., Any]) -> str
name = ""
if hasattr(receiver, "__qualname__"):
name = receiver.__qualname__
elif hasattr(receiver, "__name__"): # Python 2.7 has no __qualname__
name = receiver.__name__
elif hasattr(
receiver, "func"
): # certain functions (like partials) dont have a name
if hasattr(receiver, "func") and hasattr(receiver.func, "__name__"):
name = "partial(<function " + receiver.func.__name__ + ">)"
if (
name == ""
): # In case nothing was found, return the string representation (this is the slowest case)
return str(receiver)
if hasattr(receiver, "__module__"): # prepend with module, if there is one
name = receiver.__module__ + "." + name
return name
def patch_signals():
# type: () -> None
"""
Patch django signal receivers to create a span.
This only wraps sync receivers. Django>=5.0 introduced async receivers, but
since we don't create transactions for ASGI Django, we don't wrap them.
"""
from sentry_sdk.integrations.django import DjangoIntegration
old_live_receivers = Signal._live_receivers
def _sentry_live_receivers(self, sender):
# type: (Signal, Any) -> Union[tuple[list[Callable[..., Any]], list[Callable[..., Any]]], list[Callable[..., Any]]]
if DJANGO_VERSION >= (5, 0):
sync_receivers, async_receivers = old_live_receivers(self, sender)
else:
sync_receivers = old_live_receivers(self, sender)
async_receivers = []
def sentry_sync_receiver_wrapper(receiver):
# type: (Callable[..., Any]) -> Callable[..., Any]
@wraps(receiver)
def wrapper(*args, **kwargs):
# type: (Any, Any) -> Any
signal_name = _get_receiver_name(receiver)
with sentry_sdk.start_span(
op=OP.EVENT_DJANGO,
name=signal_name,
origin=DjangoIntegration.origin,
) as span:
span.set_data("signal", signal_name)
return receiver(*args, **kwargs)
return wrapper
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if (
integration
and integration.signals_spans
and self not in integration.signals_denylist
):
for idx, receiver in enumerate(sync_receivers):
sync_receivers[idx] = sentry_sync_receiver_wrapper(receiver)
if DJANGO_VERSION >= (5, 0):
return sync_receivers, async_receivers
else:
return sync_receivers
Signal._live_receivers = _sentry_live_receivers
@@ -0,0 +1,188 @@
import functools
from django.template import TemplateSyntaxError
from django.utils.safestring import mark_safe
from django import VERSION as DJANGO_VERSION
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.utils import ensure_integration_enabled
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Dict
from typing import Optional
from typing import Iterator
from typing import Tuple
try:
# support Django 1.9
from django.template.base import Origin
except ImportError:
# backward compatibility
from django.template.loader import LoaderOrigin as Origin
def get_template_frame_from_exception(exc_value):
# type: (Optional[BaseException]) -> Optional[Dict[str, Any]]
# As of Django 1.9 or so the new template debug thing showed up.
if hasattr(exc_value, "template_debug"):
return _get_template_frame_from_debug(exc_value.template_debug) # type: ignore
# As of r16833 (Django) all exceptions may contain a
# ``django_template_source`` attribute (rather than the legacy
# ``TemplateSyntaxError.source`` check)
if hasattr(exc_value, "django_template_source"):
return _get_template_frame_from_source(
exc_value.django_template_source # type: ignore
)
if isinstance(exc_value, TemplateSyntaxError) and hasattr(exc_value, "source"):
source = exc_value.source
if isinstance(source, (tuple, list)) and isinstance(source[0], Origin):
return _get_template_frame_from_source(source) # type: ignore
return None
def _get_template_name_description(template_name):
# type: (str) -> str
if isinstance(template_name, (list, tuple)):
if template_name:
return "[{}, ...]".format(template_name[0])
else:
return template_name
def patch_templates():
# type: () -> None
from django.template.response import SimpleTemplateResponse
from sentry_sdk.integrations.django import DjangoIntegration
real_rendered_content = SimpleTemplateResponse.rendered_content
@property # type: ignore
@ensure_integration_enabled(DjangoIntegration, real_rendered_content.fget)
def rendered_content(self):
# type: (SimpleTemplateResponse) -> str
with sentry_sdk.start_span(
op=OP.TEMPLATE_RENDER,
name=_get_template_name_description(self.template_name),
origin=DjangoIntegration.origin,
) as span:
span.set_data("context", self.context_data)
return real_rendered_content.fget(self)
SimpleTemplateResponse.rendered_content = rendered_content
if DJANGO_VERSION < (1, 7):
return
import django.shortcuts
real_render = django.shortcuts.render
@functools.wraps(real_render)
@ensure_integration_enabled(DjangoIntegration, real_render)
def render(request, template_name, context=None, *args, **kwargs):
# type: (django.http.HttpRequest, str, Optional[Dict[str, Any]], *Any, **Any) -> django.http.HttpResponse
# Inject trace meta tags into template context
context = context or {}
if "sentry_trace_meta" not in context:
context["sentry_trace_meta"] = mark_safe(
sentry_sdk.get_current_scope().trace_propagation_meta()
)
with sentry_sdk.start_span(
op=OP.TEMPLATE_RENDER,
name=_get_template_name_description(template_name),
origin=DjangoIntegration.origin,
) as span:
span.set_data("context", context)
return real_render(request, template_name, context, *args, **kwargs)
django.shortcuts.render = render
def _get_template_frame_from_debug(debug):
# type: (Dict[str, Any]) -> Dict[str, Any]
if debug is None:
return None
lineno = debug["line"]
filename = debug["name"]
if filename is None:
filename = "<django template>"
pre_context = []
post_context = []
context_line = None
for i, line in debug["source_lines"]:
if i < lineno:
pre_context.append(line)
elif i > lineno:
post_context.append(line)
else:
context_line = line
return {
"filename": filename,
"lineno": lineno,
"pre_context": pre_context[-5:],
"post_context": post_context[:5],
"context_line": context_line,
"in_app": True,
}
def _linebreak_iter(template_source):
# type: (str) -> Iterator[int]
yield 0
p = template_source.find("\n")
while p >= 0:
yield p + 1
p = template_source.find("\n", p + 1)
def _get_template_frame_from_source(source):
# type: (Tuple[Origin, Tuple[int, int]]) -> Optional[Dict[str, Any]]
if not source:
return None
origin, (start, end) = source
filename = getattr(origin, "loadname", None)
if filename is None:
filename = "<django template>"
template_source = origin.reload()
lineno = None
upto = 0
pre_context = []
post_context = []
context_line = None
for num, next in enumerate(_linebreak_iter(template_source)):
line = template_source[upto:next]
if start >= upto and end <= next:
lineno = num
context_line = line
elif lineno is None:
pre_context.append(line)
else:
post_context.append(line)
upto = next
if context_line is None or lineno is None:
return None
return {
"filename": filename,
"lineno": lineno,
"pre_context": pre_context[-5:],
"post_context": post_context[:5],
"context_line": context_line,
}
@@ -0,0 +1,159 @@
"""
Copied from raven-python.
Despite being called "legacy" in some places this resolver is very much still
in use.
"""
import re
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from django.urls.resolvers import URLResolver
from typing import Dict
from typing import List
from typing import Optional
from django.urls.resolvers import URLPattern
from typing import Tuple
from typing import Union
from re import Pattern
from django import VERSION as DJANGO_VERSION
if DJANGO_VERSION >= (2, 0):
from django.urls.resolvers import RoutePattern
else:
RoutePattern = None
try:
from django.urls import get_resolver
except ImportError:
from django.core.urlresolvers import get_resolver
def get_regex(resolver_or_pattern):
# type: (Union[URLPattern, URLResolver]) -> Pattern[str]
"""Utility method for django's deprecated resolver.regex"""
try:
regex = resolver_or_pattern.regex
except AttributeError:
regex = resolver_or_pattern.pattern.regex
return regex
class RavenResolver:
_new_style_group_matcher = re.compile(
r"<(?:([^>:]+):)?([^>]+)>"
) # https://github.com/django/django/blob/21382e2743d06efbf5623e7c9b6dccf2a325669b/django/urls/resolvers.py#L245-L247
_optional_group_matcher = re.compile(r"\(\?\:([^\)]+)\)")
_named_group_matcher = re.compile(r"\(\?P<(\w+)>[^\)]+\)+")
_non_named_group_matcher = re.compile(r"\([^\)]+\)")
# [foo|bar|baz]
_either_option_matcher = re.compile(r"\[([^\]]+)\|([^\]]+)\]")
_camel_re = re.compile(r"([A-Z]+)([a-z])")
_cache = {} # type: Dict[URLPattern, str]
def _simplify(self, pattern):
# type: (Union[URLPattern, URLResolver]) -> str
r"""
Clean up urlpattern regexes into something readable by humans:
From:
> "^(?P<sport_slug>\w+)/athletes/(?P<athlete_slug>\w+)/$"
To:
> "{sport_slug}/athletes/{athlete_slug}/"
"""
# "new-style" path patterns can be parsed directly without turning them
# into regexes first
if (
RoutePattern is not None
and hasattr(pattern, "pattern")
and isinstance(pattern.pattern, RoutePattern)
):
return self._new_style_group_matcher.sub(
lambda m: "{%s}" % m.group(2), str(pattern.pattern._route)
)
result = get_regex(pattern).pattern
# remove optional params
# TODO(dcramer): it'd be nice to change these into [%s] but it currently
# conflicts with the other rules because we're doing regexp matches
# rather than parsing tokens
result = self._optional_group_matcher.sub(lambda m: "%s" % m.group(1), result)
# handle named groups first
result = self._named_group_matcher.sub(lambda m: "{%s}" % m.group(1), result)
# handle non-named groups
result = self._non_named_group_matcher.sub("{var}", result)
# handle optional params
result = self._either_option_matcher.sub(lambda m: m.group(1), result)
# clean up any outstanding regex-y characters.
result = (
result.replace("^", "")
.replace("$", "")
.replace("?", "")
.replace("\\A", "")
.replace("\\Z", "")
.replace("//", "/")
.replace("\\", "")
)
return result
def _resolve(self, resolver, path, parents=None):
# type: (URLResolver, str, Optional[List[URLResolver]]) -> Optional[str]
match = get_regex(resolver).search(path) # Django < 2.0
if not match:
return None
if parents is None:
parents = [resolver]
elif resolver not in parents:
parents = parents + [resolver]
new_path = path[match.end() :]
for pattern in resolver.url_patterns:
# this is an include()
if not pattern.callback:
match_ = self._resolve(pattern, new_path, parents)
if match_:
return match_
continue
elif not get_regex(pattern).search(new_path):
continue
try:
return self._cache[pattern]
except KeyError:
pass
prefix = "".join(self._simplify(p) for p in parents)
result = prefix + self._simplify(pattern)
if not result.startswith("/"):
result = "/" + result
self._cache[pattern] = result
return result
return None
def resolve(
self,
path, # type: str
urlconf=None, # type: Union[None, Tuple[URLPattern, URLPattern, URLResolver], Tuple[URLPattern]]
):
# type: (...) -> Optional[str]
resolver = get_resolver(urlconf)
match = self._resolve(resolver, path)
return match
LEGACY_RESOLVER = RavenResolver()
@@ -0,0 +1,96 @@
import functools
import sentry_sdk
from sentry_sdk.consts import OP
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
try:
from asyncio import iscoroutinefunction
except ImportError:
iscoroutinefunction = None # type: ignore
try:
from sentry_sdk.integrations.django.asgi import wrap_async_view
except (ImportError, SyntaxError):
wrap_async_view = None # type: ignore
def patch_views():
# type: () -> None
from django.core.handlers.base import BaseHandler
from django.template.response import SimpleTemplateResponse
from sentry_sdk.integrations.django import DjangoIntegration
old_make_view_atomic = BaseHandler.make_view_atomic
old_render = SimpleTemplateResponse.render
def sentry_patched_render(self):
# type: (SimpleTemplateResponse) -> Any
with sentry_sdk.start_span(
op=OP.VIEW_RESPONSE_RENDER,
name="serialize response",
origin=DjangoIntegration.origin,
):
return old_render(self)
@functools.wraps(old_make_view_atomic)
def sentry_patched_make_view_atomic(self, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
callback = old_make_view_atomic(self, *args, **kwargs)
# XXX: The wrapper function is created for every request. Find more
# efficient way to wrap views (or build a cache?)
integration = sentry_sdk.get_client().get_integration(DjangoIntegration)
if integration is not None and integration.middleware_spans:
is_async_view = (
iscoroutinefunction is not None
and wrap_async_view is not None
and iscoroutinefunction(callback)
)
if is_async_view:
sentry_wrapped_callback = wrap_async_view(callback)
else:
sentry_wrapped_callback = _wrap_sync_view(callback)
else:
sentry_wrapped_callback = callback
return sentry_wrapped_callback
SimpleTemplateResponse.render = sentry_patched_render
BaseHandler.make_view_atomic = sentry_patched_make_view_atomic
def _wrap_sync_view(callback):
# type: (Any) -> Any
from sentry_sdk.integrations.django import DjangoIntegration
@functools.wraps(callback)
def sentry_wrapped_callback(request, *args, **kwargs):
# type: (Any, *Any, **Any) -> Any
current_scope = sentry_sdk.get_current_scope()
if current_scope.transaction is not None:
current_scope.transaction.update_active_thread()
sentry_scope = sentry_sdk.get_isolation_scope()
# set the active thread id to the handler thread for sync views
# this isn't necessary for async views since that runs on main
if sentry_scope.profile is not None:
sentry_scope.profile.update_active_thread_id()
with sentry_sdk.start_span(
op=OP.VIEW_RENDER,
name=request.resolver_match.view_name,
origin=DjangoIntegration.origin,
):
return callback(request, *args, **kwargs)
return sentry_wrapped_callback
@@ -0,0 +1,168 @@
import json
import sentry_sdk
from sentry_sdk.integrations import Integration
from sentry_sdk.integrations._wsgi_common import request_body_within_bounds
from sentry_sdk.utils import (
AnnotatedValue,
capture_internal_exceptions,
event_from_exception,
)
from dramatiq.broker import Broker # type: ignore
from dramatiq.message import Message # type: ignore
from dramatiq.middleware import Middleware, default_middleware # type: ignore
from dramatiq.errors import Retry # type: ignore
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Dict, Optional, Union
from sentry_sdk._types import Event, Hint
class DramatiqIntegration(Integration):
"""
Dramatiq integration for Sentry
Please make sure that you call `sentry_sdk.init` *before* initializing
your broker, as it monkey patches `Broker.__init__`.
This integration was originally developed and maintained
by https://github.com/jacobsvante and later donated to the Sentry
project.
"""
identifier = "dramatiq"
@staticmethod
def setup_once():
# type: () -> None
_patch_dramatiq_broker()
def _patch_dramatiq_broker():
# type: () -> None
original_broker__init__ = Broker.__init__
def sentry_patched_broker__init__(self, *args, **kw):
# type: (Broker, *Any, **Any) -> None
integration = sentry_sdk.get_client().get_integration(DramatiqIntegration)
try:
middleware = kw.pop("middleware")
except KeyError:
# Unfortunately Broker and StubBroker allows middleware to be
# passed in as positional arguments, whilst RabbitmqBroker and
# RedisBroker does not.
if len(args) == 1:
middleware = args[0]
args = [] # type: ignore
else:
middleware = None
if middleware is None:
middleware = list(m() for m in default_middleware)
else:
middleware = list(middleware)
if integration is not None:
middleware = [m for m in middleware if not isinstance(m, SentryMiddleware)]
middleware.insert(0, SentryMiddleware())
kw["middleware"] = middleware
original_broker__init__(self, *args, **kw)
Broker.__init__ = sentry_patched_broker__init__
class SentryMiddleware(Middleware): # type: ignore[misc]
"""
A Dramatiq middleware that automatically captures and sends
exceptions to Sentry.
This is automatically added to every instantiated broker via the
DramatiqIntegration.
"""
def before_process_message(self, broker, message):
# type: (Broker, Message) -> None
integration = sentry_sdk.get_client().get_integration(DramatiqIntegration)
if integration is None:
return
message._scope_manager = sentry_sdk.new_scope()
message._scope_manager.__enter__()
scope = sentry_sdk.get_current_scope()
scope.transaction = message.actor_name
scope.set_extra("dramatiq_message_id", message.message_id)
scope.add_event_processor(_make_message_event_processor(message, integration))
def after_process_message(self, broker, message, *, result=None, exception=None):
# type: (Broker, Message, Any, Optional[Any], Optional[Exception]) -> None
integration = sentry_sdk.get_client().get_integration(DramatiqIntegration)
if integration is None:
return
actor = broker.get_actor(message.actor_name)
throws = message.options.get("throws") or actor.options.get("throws")
try:
if (
exception is not None
and not (throws and isinstance(exception, throws))
and not isinstance(exception, Retry)
):
event, hint = event_from_exception(
exception,
client_options=sentry_sdk.get_client().options,
mechanism={
"type": DramatiqIntegration.identifier,
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
finally:
message._scope_manager.__exit__(None, None, None)
def _make_message_event_processor(message, integration):
# type: (Message, DramatiqIntegration) -> Callable[[Event, Hint], Optional[Event]]
def inner(event, hint):
# type: (Event, Hint) -> Optional[Event]
with capture_internal_exceptions():
DramatiqMessageExtractor(message).extract_into_event(event)
return event
return inner
class DramatiqMessageExtractor:
def __init__(self, message):
# type: (Message) -> None
self.message_data = dict(message.asdict())
def content_length(self):
# type: () -> int
return len(json.dumps(self.message_data))
def extract_into_event(self, event):
# type: (Event) -> None
client = sentry_sdk.get_client()
if not client.is_active():
return
contexts = event.setdefault("contexts", {})
request_info = contexts.setdefault("dramatiq", {})
request_info["type"] = "dramatiq"
data = None # type: Optional[Union[AnnotatedValue, Dict[str, Any]]]
if not request_body_within_bounds(client, self.content_length()):
data = AnnotatedValue.removed_because_over_size_limit()
else:
data = self.message_data
request_info["data"] = data
@@ -0,0 +1,83 @@
import sys
import sentry_sdk
from sentry_sdk.utils import (
capture_internal_exceptions,
event_from_exception,
)
from sentry_sdk.integrations import Integration
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Callable
from typing import Any
from typing import Type
from typing import Optional
from types import TracebackType
Excepthook = Callable[
[Type[BaseException], BaseException, Optional[TracebackType]],
Any,
]
class ExcepthookIntegration(Integration):
identifier = "excepthook"
always_run = False
def __init__(self, always_run=False):
# type: (bool) -> None
if not isinstance(always_run, bool):
raise ValueError(
"Invalid value for always_run: %s (must be type boolean)"
% (always_run,)
)
self.always_run = always_run
@staticmethod
def setup_once():
# type: () -> None
sys.excepthook = _make_excepthook(sys.excepthook)
def _make_excepthook(old_excepthook):
# type: (Excepthook) -> Excepthook
def sentry_sdk_excepthook(type_, value, traceback):
# type: (Type[BaseException], BaseException, Optional[TracebackType]) -> None
integration = sentry_sdk.get_client().get_integration(ExcepthookIntegration)
# Note: If we replace this with ensure_integration_enabled then
# we break the exceptiongroup backport;
# See: https://github.com/getsentry/sentry-python/issues/3097
if integration is None:
return old_excepthook(type_, value, traceback)
if _should_send(integration.always_run):
with capture_internal_exceptions():
event, hint = event_from_exception(
(type_, value, traceback),
client_options=sentry_sdk.get_client().options,
mechanism={"type": "excepthook", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
return old_excepthook(type_, value, traceback)
return sentry_sdk_excepthook
def _should_send(always_run=False):
# type: (bool) -> bool
if always_run:
return True
if hasattr(sys, "ps1"):
# Disable the excepthook for interactive Python shells, otherwise
# every typo gets sent to Sentry.
return False
return True
@@ -0,0 +1,67 @@
import sentry_sdk
from sentry_sdk.integrations import Integration, DidNotEnable
from sentry_sdk.scope import add_global_event_processor
from sentry_sdk.utils import walk_exception_chain, iter_stacks
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional
from sentry_sdk._types import Event, Hint
try:
import executing
except ImportError:
raise DidNotEnable("executing is not installed")
class ExecutingIntegration(Integration):
identifier = "executing"
@staticmethod
def setup_once():
# type: () -> None
@add_global_event_processor
def add_executing_info(event, hint):
# type: (Event, Optional[Hint]) -> Optional[Event]
if sentry_sdk.get_client().get_integration(ExecutingIntegration) is None:
return event
if hint is None:
return event
exc_info = hint.get("exc_info", None)
if exc_info is None:
return event
exception = event.get("exception", None)
if exception is None:
return event
values = exception.get("values", None)
if values is None:
return event
for exception, (_exc_type, _exc_value, exc_tb) in zip(
reversed(values), walk_exception_chain(exc_info)
):
sentry_frames = [
frame
for frame in exception.get("stacktrace", {}).get("frames", [])
if frame.get("function")
]
tbs = list(iter_stacks(exc_tb))
if len(sentry_frames) != len(tbs):
continue
for sentry_frame, tb in zip(sentry_frames, tbs):
frame = tb.tb_frame
source = executing.Source.for_frame(frame)
sentry_frame["function"] = source.code_qualname(frame.f_code)
return event
@@ -0,0 +1,272 @@
import sentry_sdk
from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
from sentry_sdk.integrations._wsgi_common import RequestExtractor
from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
from sentry_sdk.tracing import SOURCE_FOR_STYLE
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
parse_version,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from typing import Dict
from typing import Optional
from sentry_sdk._types import Event, EventProcessor
# In Falcon 3.0 `falcon.api_helpers` is renamed to `falcon.app_helpers`
# and `falcon.API` to `falcon.App`
try:
import falcon # type: ignore
from falcon import __version__ as FALCON_VERSION
except ImportError:
raise DidNotEnable("Falcon not installed")
try:
import falcon.app_helpers # type: ignore
falcon_helpers = falcon.app_helpers
falcon_app_class = falcon.App
FALCON3 = True
except ImportError:
import falcon.api_helpers # type: ignore
falcon_helpers = falcon.api_helpers
falcon_app_class = falcon.API
FALCON3 = False
_FALCON_UNSET = None # type: Optional[object]
if FALCON3: # falcon.request._UNSET is only available in Falcon 3.0+
with capture_internal_exceptions():
from falcon.request import _UNSET as _FALCON_UNSET # type: ignore[import-not-found, no-redef]
class FalconRequestExtractor(RequestExtractor):
def env(self):
# type: () -> Dict[str, Any]
return self.request.env
def cookies(self):
# type: () -> Dict[str, Any]
return self.request.cookies
def form(self):
# type: () -> None
return None # No such concept in Falcon
def files(self):
# type: () -> None
return None # No such concept in Falcon
def raw_data(self):
# type: () -> Optional[str]
# As request data can only be read once we won't make this available
# to Sentry. Just send back a dummy string in case there was a
# content length.
# TODO(jmagnusson): Figure out if there's a way to support this
content_length = self.content_length()
if content_length > 0:
return "[REQUEST_CONTAINING_RAW_DATA]"
else:
return None
def json(self):
# type: () -> Optional[Dict[str, Any]]
# fallback to cached_media = None if self.request._media is not available
cached_media = None
with capture_internal_exceptions():
# self.request._media is the cached self.request.media
# value. It is only available if self.request.media
# has already been accessed. Therefore, reading
# self.request._media will not exhaust the raw request
# stream (self.request.bounded_stream) because it has
# already been read if self.request._media is set.
cached_media = self.request._media
if cached_media is not _FALCON_UNSET:
return cached_media
return None
class SentryFalconMiddleware:
"""Captures exceptions in Falcon requests and send to Sentry"""
def process_request(self, req, resp, *args, **kwargs):
# type: (Any, Any, *Any, **Any) -> None
integration = sentry_sdk.get_client().get_integration(FalconIntegration)
if integration is None:
return
scope = sentry_sdk.get_isolation_scope()
scope._name = "falcon"
scope.add_event_processor(_make_request_event_processor(req, integration))
TRANSACTION_STYLE_VALUES = ("uri_template", "path")
class FalconIntegration(Integration):
identifier = "falcon"
origin = f"auto.http.{identifier}"
transaction_style = ""
def __init__(self, transaction_style="uri_template"):
# type: (str) -> None
if transaction_style not in TRANSACTION_STYLE_VALUES:
raise ValueError(
"Invalid value for transaction_style: %s (must be in %s)"
% (transaction_style, TRANSACTION_STYLE_VALUES)
)
self.transaction_style = transaction_style
@staticmethod
def setup_once():
# type: () -> None
version = parse_version(FALCON_VERSION)
_check_minimum_version(FalconIntegration, version)
_patch_wsgi_app()
_patch_handle_exception()
_patch_prepare_middleware()
def _patch_wsgi_app():
# type: () -> None
original_wsgi_app = falcon_app_class.__call__
def sentry_patched_wsgi_app(self, env, start_response):
# type: (falcon.API, Any, Any) -> Any
integration = sentry_sdk.get_client().get_integration(FalconIntegration)
if integration is None:
return original_wsgi_app(self, env, start_response)
sentry_wrapped = SentryWsgiMiddleware(
lambda envi, start_resp: original_wsgi_app(self, envi, start_resp),
span_origin=FalconIntegration.origin,
)
return sentry_wrapped(env, start_response)
falcon_app_class.__call__ = sentry_patched_wsgi_app
def _patch_handle_exception():
# type: () -> None
original_handle_exception = falcon_app_class._handle_exception
@ensure_integration_enabled(FalconIntegration, original_handle_exception)
def sentry_patched_handle_exception(self, *args):
# type: (falcon.API, *Any) -> Any
# NOTE(jmagnusson): falcon 2.0 changed falcon.API._handle_exception
# method signature from `(ex, req, resp, params)` to
# `(req, resp, ex, params)`
ex = response = None
with capture_internal_exceptions():
ex = next(argument for argument in args if isinstance(argument, Exception))
response = next(
argument for argument in args if isinstance(argument, falcon.Response)
)
was_handled = original_handle_exception(self, *args)
if ex is None or response is None:
# Both ex and response should have a non-None value at this point; otherwise,
# there is an error with the SDK that will have been captured in the
# capture_internal_exceptions block above.
return was_handled
if _exception_leads_to_http_5xx(ex, response):
event, hint = event_from_exception(
ex,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "falcon", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
return was_handled
falcon_app_class._handle_exception = sentry_patched_handle_exception
def _patch_prepare_middleware():
# type: () -> None
original_prepare_middleware = falcon_helpers.prepare_middleware
def sentry_patched_prepare_middleware(
middleware=None, independent_middleware=False, asgi=False
):
# type: (Any, Any, bool) -> Any
if asgi:
# We don't support ASGI Falcon apps, so we don't patch anything here
return original_prepare_middleware(middleware, independent_middleware, asgi)
integration = sentry_sdk.get_client().get_integration(FalconIntegration)
if integration is not None:
middleware = [SentryFalconMiddleware()] + (middleware or [])
# We intentionally omit the asgi argument here, since the default is False anyways,
# and this way, we remain backwards-compatible with pre-3.0.0 Falcon versions.
return original_prepare_middleware(middleware, independent_middleware)
falcon_helpers.prepare_middleware = sentry_patched_prepare_middleware
def _exception_leads_to_http_5xx(ex, response):
# type: (Exception, falcon.Response) -> bool
is_server_error = isinstance(ex, falcon.HTTPError) and (ex.status or "").startswith(
"5"
)
is_unhandled_error = not isinstance(
ex, (falcon.HTTPError, falcon.http_status.HTTPStatus)
)
# We only check the HTTP status on Falcon 3 because in Falcon 2, the status on the response
# at the stage where we capture it is listed as 200, even though we would expect to see a 500
# status. Since at the time of this change, Falcon 2 is ca. 4 years old, we have decided to
# only perform this check on Falcon 3+, despite the risk that some handled errors might be
# reported to Sentry as unhandled on Falcon 2.
return (is_server_error or is_unhandled_error) and (
not FALCON3 or _has_http_5xx_status(response)
)
def _has_http_5xx_status(response):
# type: (falcon.Response) -> bool
return response.status.startswith("5")
def _set_transaction_name_and_source(event, transaction_style, request):
# type: (Event, str, falcon.Request) -> None
name_for_style = {
"uri_template": request.uri_template,
"path": request.path,
}
event["transaction"] = name_for_style[transaction_style]
event["transaction_info"] = {"source": SOURCE_FOR_STYLE[transaction_style]}
def _make_request_event_processor(req, integration):
# type: (falcon.Request, FalconIntegration) -> EventProcessor
def event_processor(event, hint):
# type: (Event, dict[str, Any]) -> Event
_set_transaction_name_and_source(event, integration.transaction_style, req)
with capture_internal_exceptions():
FalconRequestExtractor(req).extract_into_event(event)
return event
return event_processor
@@ -0,0 +1,147 @@
import asyncio
from copy import deepcopy
from functools import wraps
import sentry_sdk
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_ROUTE
from sentry_sdk.utils import (
transaction_from_function,
logger,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Dict
from sentry_sdk._types import Event
try:
from sentry_sdk.integrations.starlette import (
StarletteIntegration,
StarletteRequestExtractor,
)
except DidNotEnable:
raise DidNotEnable("Starlette is not installed")
try:
import fastapi # type: ignore
except ImportError:
raise DidNotEnable("FastAPI is not installed")
_DEFAULT_TRANSACTION_NAME = "generic FastAPI request"
class FastApiIntegration(StarletteIntegration):
identifier = "fastapi"
@staticmethod
def setup_once():
# type: () -> None
patch_get_request_handler()
def _set_transaction_name_and_source(scope, transaction_style, request):
# type: (sentry_sdk.Scope, str, Any) -> None
name = ""
if transaction_style == "endpoint":
endpoint = request.scope.get("endpoint")
if endpoint:
name = transaction_from_function(endpoint) or ""
elif transaction_style == "url":
route = request.scope.get("route")
if route:
path = getattr(route, "path", None)
if path is not None:
name = path
if not name:
name = _DEFAULT_TRANSACTION_NAME
source = TRANSACTION_SOURCE_ROUTE
else:
source = SOURCE_FOR_STYLE[transaction_style]
scope.set_transaction_name(name, source=source)
logger.debug(
"[FastAPI] Set transaction name and source on scope: %s / %s", name, source
)
def patch_get_request_handler():
# type: () -> None
old_get_request_handler = fastapi.routing.get_request_handler
def _sentry_get_request_handler(*args, **kwargs):
# type: (*Any, **Any) -> Any
dependant = kwargs.get("dependant")
if (
dependant
and dependant.call is not None
and not asyncio.iscoroutinefunction(dependant.call)
):
old_call = dependant.call
@wraps(old_call)
def _sentry_call(*args, **kwargs):
# type: (*Any, **Any) -> Any
current_scope = sentry_sdk.get_current_scope()
if current_scope.transaction is not None:
current_scope.transaction.update_active_thread()
sentry_scope = sentry_sdk.get_isolation_scope()
if sentry_scope.profile is not None:
sentry_scope.profile.update_active_thread_id()
return old_call(*args, **kwargs)
dependant.call = _sentry_call
old_app = old_get_request_handler(*args, **kwargs)
async def _sentry_app(*args, **kwargs):
# type: (*Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(FastApiIntegration)
if integration is None:
return await old_app(*args, **kwargs)
request = args[0]
_set_transaction_name_and_source(
sentry_sdk.get_current_scope(), integration.transaction_style, request
)
sentry_scope = sentry_sdk.get_isolation_scope()
extractor = StarletteRequestExtractor(request)
info = await extractor.extract_request_info()
def _make_request_event_processor(req, integration):
# type: (Any, Any) -> Callable[[Event, Dict[str, Any]], Event]
def event_processor(event, hint):
# type: (Event, Dict[str, Any]) -> Event
# Extract information from request
request_info = event.get("request", {})
if info:
if "cookies" in info and should_send_default_pii():
request_info["cookies"] = info["cookies"]
if "data" in info:
request_info["data"] = info["data"]
event["request"] = deepcopy(request_info)
return event
return event_processor
sentry_scope._name = FastApiIntegration.identifier
sentry_scope.add_event_processor(
_make_request_event_processor(request, integration)
)
return await old_app(*args, **kwargs)
return _sentry_app
fastapi.routing.get_request_handler = _sentry_get_request_handler
@@ -0,0 +1,263 @@
import sentry_sdk
from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
from sentry_sdk.integrations._wsgi_common import (
DEFAULT_HTTP_METHODS_TO_CAPTURE,
RequestExtractor,
)
from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import SOURCE_FOR_STYLE
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
package_version,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Dict, Union
from sentry_sdk._types import Event, EventProcessor
from sentry_sdk.integrations.wsgi import _ScopedResponse
from werkzeug.datastructures import FileStorage, ImmutableMultiDict
try:
import flask_login # type: ignore
except ImportError:
flask_login = None
try:
from flask import Flask, Request # type: ignore
from flask import request as flask_request
from flask.signals import (
before_render_template,
got_request_exception,
request_started,
)
from markupsafe import Markup
except ImportError:
raise DidNotEnable("Flask is not installed")
try:
import blinker # noqa
except ImportError:
raise DidNotEnable("blinker is not installed")
TRANSACTION_STYLE_VALUES = ("endpoint", "url")
class FlaskIntegration(Integration):
identifier = "flask"
origin = f"auto.http.{identifier}"
transaction_style = ""
def __init__(
self,
transaction_style="endpoint", # type: str
http_methods_to_capture=DEFAULT_HTTP_METHODS_TO_CAPTURE, # type: tuple[str, ...]
):
# type: (...) -> None
if transaction_style not in TRANSACTION_STYLE_VALUES:
raise ValueError(
"Invalid value for transaction_style: %s (must be in %s)"
% (transaction_style, TRANSACTION_STYLE_VALUES)
)
self.transaction_style = transaction_style
self.http_methods_to_capture = tuple(map(str.upper, http_methods_to_capture))
@staticmethod
def setup_once():
# type: () -> None
version = package_version("flask")
_check_minimum_version(FlaskIntegration, version)
before_render_template.connect(_add_sentry_trace)
request_started.connect(_request_started)
got_request_exception.connect(_capture_exception)
old_app = Flask.__call__
def sentry_patched_wsgi_app(self, environ, start_response):
# type: (Any, Dict[str, str], Callable[..., Any]) -> _ScopedResponse
if sentry_sdk.get_client().get_integration(FlaskIntegration) is None:
return old_app(self, environ, start_response)
integration = sentry_sdk.get_client().get_integration(FlaskIntegration)
middleware = SentryWsgiMiddleware(
lambda *a, **kw: old_app(self, *a, **kw),
span_origin=FlaskIntegration.origin,
http_methods_to_capture=(
integration.http_methods_to_capture
if integration
else DEFAULT_HTTP_METHODS_TO_CAPTURE
),
)
return middleware(environ, start_response)
Flask.__call__ = sentry_patched_wsgi_app
def _add_sentry_trace(sender, template, context, **extra):
# type: (Flask, Any, Dict[str, Any], **Any) -> None
if "sentry_trace" in context:
return
scope = sentry_sdk.get_current_scope()
trace_meta = Markup(scope.trace_propagation_meta())
context["sentry_trace"] = trace_meta # for backwards compatibility
context["sentry_trace_meta"] = trace_meta
def _set_transaction_name_and_source(scope, transaction_style, request):
# type: (sentry_sdk.Scope, str, Request) -> None
try:
name_for_style = {
"url": request.url_rule.rule,
"endpoint": request.url_rule.endpoint,
}
scope.set_transaction_name(
name_for_style[transaction_style],
source=SOURCE_FOR_STYLE[transaction_style],
)
except Exception:
pass
def _request_started(app, **kwargs):
# type: (Flask, **Any) -> None
integration = sentry_sdk.get_client().get_integration(FlaskIntegration)
if integration is None:
return
request = flask_request._get_current_object()
# Set the transaction name and source here,
# but rely on WSGI middleware to actually start the transaction
_set_transaction_name_and_source(
sentry_sdk.get_current_scope(), integration.transaction_style, request
)
scope = sentry_sdk.get_isolation_scope()
evt_processor = _make_request_event_processor(app, request, integration)
scope.add_event_processor(evt_processor)
class FlaskRequestExtractor(RequestExtractor):
def env(self):
# type: () -> Dict[str, str]
return self.request.environ
def cookies(self):
# type: () -> Dict[Any, Any]
return {
k: v[0] if isinstance(v, list) and len(v) == 1 else v
for k, v in self.request.cookies.items()
}
def raw_data(self):
# type: () -> bytes
return self.request.get_data()
def form(self):
# type: () -> ImmutableMultiDict[str, Any]
return self.request.form
def files(self):
# type: () -> ImmutableMultiDict[str, Any]
return self.request.files
def is_json(self):
# type: () -> bool
return self.request.is_json
def json(self):
# type: () -> Any
return self.request.get_json(silent=True)
def size_of_file(self, file):
# type: (FileStorage) -> int
return file.content_length
def _make_request_event_processor(app, request, integration):
# type: (Flask, Callable[[], Request], FlaskIntegration) -> EventProcessor
def inner(event, hint):
# type: (Event, dict[str, Any]) -> Event
# if the request is gone we are fine not logging the data from
# it. This might happen if the processor is pushed away to
# another thread.
if request is None:
return event
with capture_internal_exceptions():
FlaskRequestExtractor(request).extract_into_event(event)
if should_send_default_pii():
with capture_internal_exceptions():
_add_user_to_event(event)
return event
return inner
@ensure_integration_enabled(FlaskIntegration)
def _capture_exception(sender, exception, **kwargs):
# type: (Flask, Union[ValueError, BaseException], **Any) -> None
event, hint = event_from_exception(
exception,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "flask", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _add_user_to_event(event):
# type: (Event) -> None
if flask_login is None:
return
user = flask_login.current_user
if user is None:
return
with capture_internal_exceptions():
# Access this object as late as possible as accessing the user
# is relatively costly
user_info = event.setdefault("user", {})
try:
user_info.setdefault("id", user.get_id())
# TODO: more configurable user attrs here
except AttributeError:
# might happen if:
# - flask_login could not be imported
# - flask_login is not configured
# - no user is logged in
pass
# The following attribute accesses are ineffective for the general
# Flask-Login case, because the User interface of Flask-Login does not
# care about anything but the ID. However, Flask-User (based on
# Flask-Login) documents a few optional extra attributes.
#
# https://github.com/lingthio/Flask-User/blob/a379fa0a281789618c484b459cb41236779b95b1/docs/source/data_models.rst#fixed-data-model-property-names
try:
user_info.setdefault("email", user.email)
except Exception:
pass
try:
user_info.setdefault("username", user.username)
except Exception:
pass
@@ -0,0 +1,234 @@
import functools
import sys
from copy import deepcopy
from datetime import datetime, timedelta, timezone
from os import environ
import sentry_sdk
from sentry_sdk.api import continue_trace
from sentry_sdk.consts import OP
from sentry_sdk.integrations import Integration
from sentry_sdk.integrations._wsgi_common import _filter_headers
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import TRANSACTION_SOURCE_COMPONENT
from sentry_sdk.utils import (
AnnotatedValue,
capture_internal_exceptions,
event_from_exception,
logger,
TimeoutThread,
reraise,
)
from typing import TYPE_CHECKING
# Constants
TIMEOUT_WARNING_BUFFER = 1.5 # Buffer time required to send timeout warning to Sentry
MILLIS_TO_SECONDS = 1000.0
if TYPE_CHECKING:
from typing import Any
from typing import TypeVar
from typing import Callable
from typing import Optional
from sentry_sdk._types import EventProcessor, Event, Hint
F = TypeVar("F", bound=Callable[..., Any])
def _wrap_func(func):
# type: (F) -> F
@functools.wraps(func)
def sentry_func(functionhandler, gcp_event, *args, **kwargs):
# type: (Any, Any, *Any, **Any) -> Any
client = sentry_sdk.get_client()
integration = client.get_integration(GcpIntegration)
if integration is None:
return func(functionhandler, gcp_event, *args, **kwargs)
configured_time = environ.get("FUNCTION_TIMEOUT_SEC")
if not configured_time:
logger.debug(
"The configured timeout could not be fetched from Cloud Functions configuration."
)
return func(functionhandler, gcp_event, *args, **kwargs)
configured_time = int(configured_time)
initial_time = datetime.now(timezone.utc)
with sentry_sdk.isolation_scope() as scope:
with capture_internal_exceptions():
scope.clear_breadcrumbs()
scope.add_event_processor(
_make_request_event_processor(
gcp_event, configured_time, initial_time
)
)
scope.set_tag("gcp_region", environ.get("FUNCTION_REGION"))
timeout_thread = None
if (
integration.timeout_warning
and configured_time > TIMEOUT_WARNING_BUFFER
):
waiting_time = configured_time - TIMEOUT_WARNING_BUFFER
timeout_thread = TimeoutThread(waiting_time, configured_time)
# Starting the thread to raise timeout warning exception
timeout_thread.start()
headers = {}
if hasattr(gcp_event, "headers"):
headers = gcp_event.headers
transaction = continue_trace(
headers,
op=OP.FUNCTION_GCP,
name=environ.get("FUNCTION_NAME", ""),
source=TRANSACTION_SOURCE_COMPONENT,
origin=GcpIntegration.origin,
)
sampling_context = {
"gcp_env": {
"function_name": environ.get("FUNCTION_NAME"),
"function_entry_point": environ.get("ENTRY_POINT"),
"function_identity": environ.get("FUNCTION_IDENTITY"),
"function_region": environ.get("FUNCTION_REGION"),
"function_project": environ.get("GCP_PROJECT"),
},
"gcp_event": gcp_event,
}
with sentry_sdk.start_transaction(
transaction, custom_sampling_context=sampling_context
):
try:
return func(functionhandler, gcp_event, *args, **kwargs)
except Exception:
exc_info = sys.exc_info()
sentry_event, hint = event_from_exception(
exc_info,
client_options=client.options,
mechanism={"type": "gcp", "handled": False},
)
sentry_sdk.capture_event(sentry_event, hint=hint)
reraise(*exc_info)
finally:
if timeout_thread:
timeout_thread.stop()
# Flush out the event queue
client.flush()
return sentry_func # type: ignore
class GcpIntegration(Integration):
identifier = "gcp"
origin = f"auto.function.{identifier}"
def __init__(self, timeout_warning=False):
# type: (bool) -> None
self.timeout_warning = timeout_warning
@staticmethod
def setup_once():
# type: () -> None
import __main__ as gcp_functions
if not hasattr(gcp_functions, "worker_v1"):
logger.warning(
"GcpIntegration currently supports only Python 3.7 runtime environment."
)
return
worker1 = gcp_functions.worker_v1
worker1.FunctionHandler.invoke_user_function = _wrap_func(
worker1.FunctionHandler.invoke_user_function
)
def _make_request_event_processor(gcp_event, configured_timeout, initial_time):
# type: (Any, Any, Any) -> EventProcessor
def event_processor(event, hint):
# type: (Event, Hint) -> Optional[Event]
final_time = datetime.now(timezone.utc)
time_diff = final_time - initial_time
execution_duration_in_millis = time_diff / timedelta(milliseconds=1)
extra = event.setdefault("extra", {})
extra["google cloud functions"] = {
"function_name": environ.get("FUNCTION_NAME"),
"function_entry_point": environ.get("ENTRY_POINT"),
"function_identity": environ.get("FUNCTION_IDENTITY"),
"function_region": environ.get("FUNCTION_REGION"),
"function_project": environ.get("GCP_PROJECT"),
"execution_duration_in_millis": execution_duration_in_millis,
"configured_timeout_in_seconds": configured_timeout,
}
extra["google cloud logs"] = {
"url": _get_google_cloud_logs_url(final_time),
}
request = event.get("request", {})
request["url"] = "gcp:///{}".format(environ.get("FUNCTION_NAME"))
if hasattr(gcp_event, "method"):
request["method"] = gcp_event.method
if hasattr(gcp_event, "query_string"):
request["query_string"] = gcp_event.query_string.decode("utf-8")
if hasattr(gcp_event, "headers"):
request["headers"] = _filter_headers(gcp_event.headers)
if should_send_default_pii():
if hasattr(gcp_event, "data"):
request["data"] = gcp_event.data
else:
if hasattr(gcp_event, "data"):
# Unfortunately couldn't find a way to get structured body from GCP
# event. Meaning every body is unstructured to us.
request["data"] = AnnotatedValue.removed_because_raw_data()
event["request"] = deepcopy(request)
return event
return event_processor
def _get_google_cloud_logs_url(final_time):
# type: (datetime) -> str
"""
Generates a Google Cloud Logs console URL based on the environment variables
Arguments:
final_time {datetime} -- Final time
Returns:
str -- Google Cloud Logs Console URL to logs.
"""
hour_ago = final_time - timedelta(hours=1)
formatstring = "%Y-%m-%dT%H:%M:%SZ"
url = (
"https://console.cloud.google.com/logs/viewer?project={project}&resource=cloud_function"
"%2Ffunction_name%2F{function_name}%2Fregion%2F{region}&minLogLevel=0&expandAll=false"
"&timestamp={timestamp_end}&customFacets=&limitCustomFacetWidth=true"
"&dateRangeStart={timestamp_start}&dateRangeEnd={timestamp_end}"
"&interval=PT1H&scrollTimestamp={timestamp_end}"
).format(
project=environ.get("GCP_PROJECT"),
function_name=environ.get("FUNCTION_NAME"),
region=environ.get("FUNCTION_REGION"),
timestamp_end=final_time.strftime(formatstring),
timestamp_start=hour_ago.strftime(formatstring),
)
return url
@@ -0,0 +1,107 @@
import re
import sentry_sdk
from sentry_sdk.integrations import Integration
from sentry_sdk.scope import add_global_event_processor
from sentry_sdk.utils import capture_internal_exceptions
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
from sentry_sdk._types import Event
MODULE_RE = r"[a-zA-Z0-9/._:\\-]+"
TYPE_RE = r"[a-zA-Z0-9._:<>,-]+"
HEXVAL_RE = r"[A-Fa-f0-9]+"
FRAME_RE = r"""
^(?P<index>\d+)\.\s
(?P<package>{MODULE_RE})\(
(?P<retval>{TYPE_RE}\ )?
((?P<function>{TYPE_RE})
(?P<args>\(.*\))?
)?
((?P<constoffset>\ const)?\+0x(?P<offset>{HEXVAL_RE}))?
\)\s
\[0x(?P<retaddr>{HEXVAL_RE})\]$
""".format(
MODULE_RE=MODULE_RE, HEXVAL_RE=HEXVAL_RE, TYPE_RE=TYPE_RE
)
FRAME_RE = re.compile(FRAME_RE, re.MULTILINE | re.VERBOSE)
class GnuBacktraceIntegration(Integration):
identifier = "gnu_backtrace"
@staticmethod
def setup_once():
# type: () -> None
@add_global_event_processor
def process_gnu_backtrace(event, hint):
# type: (Event, dict[str, Any]) -> Event
with capture_internal_exceptions():
return _process_gnu_backtrace(event, hint)
def _process_gnu_backtrace(event, hint):
# type: (Event, dict[str, Any]) -> Event
if sentry_sdk.get_client().get_integration(GnuBacktraceIntegration) is None:
return event
exc_info = hint.get("exc_info", None)
if exc_info is None:
return event
exception = event.get("exception", None)
if exception is None:
return event
values = exception.get("values", None)
if values is None:
return event
for exception in values:
frames = exception.get("stacktrace", {}).get("frames", [])
if not frames:
continue
msg = exception.get("value", None)
if not msg:
continue
additional_frames = []
new_msg = []
for line in msg.splitlines():
match = FRAME_RE.match(line)
if match:
additional_frames.append(
(
int(match.group("index")),
{
"package": match.group("package") or None,
"function": match.group("function") or None,
"platform": "native",
},
)
)
else:
# Put garbage lines back into message, not sure what to do with them.
new_msg.append(line)
if additional_frames:
additional_frames.sort(key=lambda x: -x[0])
for _, frame in additional_frames:
frames.append(frame)
new_msg.append("<stacktrace parsed and removed by GnuBacktraceIntegration>")
exception["value"] = "\n".join(new_msg)
return event
@@ -0,0 +1,145 @@
import sentry_sdk
from sentry_sdk.utils import (
event_from_exception,
ensure_integration_enabled,
parse_version,
)
from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
from sentry_sdk.scope import should_send_default_pii
try:
import gql # type: ignore[import-not-found]
from graphql import (
print_ast,
get_operation_ast,
DocumentNode,
VariableDefinitionNode,
)
from gql.transport import Transport, AsyncTransport # type: ignore[import-not-found]
from gql.transport.exceptions import TransportQueryError # type: ignore[import-not-found]
except ImportError:
raise DidNotEnable("gql is not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Dict, Tuple, Union
from sentry_sdk._types import Event, EventProcessor
EventDataType = Dict[str, Union[str, Tuple[VariableDefinitionNode, ...]]]
class GQLIntegration(Integration):
identifier = "gql"
@staticmethod
def setup_once():
# type: () -> None
gql_version = parse_version(gql.__version__)
_check_minimum_version(GQLIntegration, gql_version)
_patch_execute()
def _data_from_document(document):
# type: (DocumentNode) -> EventDataType
try:
operation_ast = get_operation_ast(document)
data = {"query": print_ast(document)} # type: EventDataType
if operation_ast is not None:
data["variables"] = operation_ast.variable_definitions
if operation_ast.name is not None:
data["operationName"] = operation_ast.name.value
return data
except (AttributeError, TypeError):
return dict()
def _transport_method(transport):
# type: (Union[Transport, AsyncTransport]) -> str
"""
The RequestsHTTPTransport allows defining the HTTP method; all
other transports use POST.
"""
try:
return transport.method
except AttributeError:
return "POST"
def _request_info_from_transport(transport):
# type: (Union[Transport, AsyncTransport, None]) -> Dict[str, str]
if transport is None:
return {}
request_info = {
"method": _transport_method(transport),
}
try:
request_info["url"] = transport.url
except AttributeError:
pass
return request_info
def _patch_execute():
# type: () -> None
real_execute = gql.Client.execute
@ensure_integration_enabled(GQLIntegration, real_execute)
def sentry_patched_execute(self, document, *args, **kwargs):
# type: (gql.Client, DocumentNode, Any, Any) -> Any
scope = sentry_sdk.get_isolation_scope()
scope.add_event_processor(_make_gql_event_processor(self, document))
try:
return real_execute(self, document, *args, **kwargs)
except TransportQueryError as e:
event, hint = event_from_exception(
e,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "gql", "handled": False},
)
sentry_sdk.capture_event(event, hint)
raise e
gql.Client.execute = sentry_patched_execute
def _make_gql_event_processor(client, document):
# type: (gql.Client, DocumentNode) -> EventProcessor
def processor(event, hint):
# type: (Event, dict[str, Any]) -> Event
try:
errors = hint["exc_info"][1].errors
except (AttributeError, KeyError):
errors = None
request = event.setdefault("request", {})
request.update(
{
"api_target": "graphql",
**_request_info_from_transport(client.transport),
}
)
if should_send_default_pii():
request["data"] = _data_from_document(document)
contexts = event.setdefault("contexts", {})
response = contexts.setdefault("response", {})
response.update(
{
"data": {"errors": errors},
"type": response,
}
)
return event
return processor
@@ -0,0 +1,151 @@
from contextlib import contextmanager
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
package_version,
)
try:
from graphene.types import schema as graphene_schema # type: ignore
except ImportError:
raise DidNotEnable("graphene is not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Generator
from typing import Any, Dict, Union
from graphene.language.source import Source # type: ignore
from graphql.execution import ExecutionResult
from graphql.type import GraphQLSchema
from sentry_sdk._types import Event
class GrapheneIntegration(Integration):
identifier = "graphene"
@staticmethod
def setup_once():
# type: () -> None
version = package_version("graphene")
_check_minimum_version(GrapheneIntegration, version)
_patch_graphql()
def _patch_graphql():
# type: () -> None
old_graphql_sync = graphene_schema.graphql_sync
old_graphql_async = graphene_schema.graphql
@ensure_integration_enabled(GrapheneIntegration, old_graphql_sync)
def _sentry_patched_graphql_sync(schema, source, *args, **kwargs):
# type: (GraphQLSchema, Union[str, Source], Any, Any) -> ExecutionResult
scope = sentry_sdk.get_isolation_scope()
scope.add_event_processor(_event_processor)
with graphql_span(schema, source, kwargs):
result = old_graphql_sync(schema, source, *args, **kwargs)
with capture_internal_exceptions():
client = sentry_sdk.get_client()
for error in result.errors or []:
event, hint = event_from_exception(
error,
client_options=client.options,
mechanism={
"type": GrapheneIntegration.identifier,
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
return result
async def _sentry_patched_graphql_async(schema, source, *args, **kwargs):
# type: (GraphQLSchema, Union[str, Source], Any, Any) -> ExecutionResult
integration = sentry_sdk.get_client().get_integration(GrapheneIntegration)
if integration is None:
return await old_graphql_async(schema, source, *args, **kwargs)
scope = sentry_sdk.get_isolation_scope()
scope.add_event_processor(_event_processor)
with graphql_span(schema, source, kwargs):
result = await old_graphql_async(schema, source, *args, **kwargs)
with capture_internal_exceptions():
client = sentry_sdk.get_client()
for error in result.errors or []:
event, hint = event_from_exception(
error,
client_options=client.options,
mechanism={
"type": GrapheneIntegration.identifier,
"handled": False,
},
)
sentry_sdk.capture_event(event, hint=hint)
return result
graphene_schema.graphql_sync = _sentry_patched_graphql_sync
graphene_schema.graphql = _sentry_patched_graphql_async
def _event_processor(event, hint):
# type: (Event, Dict[str, Any]) -> Event
if should_send_default_pii():
request_info = event.setdefault("request", {})
request_info["api_target"] = "graphql"
elif event.get("request", {}).get("data"):
del event["request"]["data"]
return event
@contextmanager
def graphql_span(schema, source, kwargs):
# type: (GraphQLSchema, Union[str, Source], Dict[str, Any]) -> Generator[None, None, None]
operation_name = kwargs.get("operation_name")
operation_type = "query"
op = OP.GRAPHQL_QUERY
if source.strip().startswith("mutation"):
operation_type = "mutation"
op = OP.GRAPHQL_MUTATION
elif source.strip().startswith("subscription"):
operation_type = "subscription"
op = OP.GRAPHQL_SUBSCRIPTION
sentry_sdk.add_breadcrumb(
crumb={
"data": {
"operation_name": operation_name,
"operation_type": operation_type,
},
"category": "graphql.operation",
},
)
scope = sentry_sdk.get_current_scope()
if scope.span:
_graphql_span = scope.span.start_child(op=op, name=operation_name)
else:
_graphql_span = sentry_sdk.start_span(op=op, name=operation_name)
_graphql_span.set_data("graphql.document", source)
_graphql_span.set_data("graphql.operation.name", operation_name)
_graphql_span.set_data("graphql.operation.type", operation_type)
try:
yield
finally:
_graphql_span.finish()
@@ -0,0 +1,151 @@
from functools import wraps
import grpc
from grpc import Channel, Server, intercept_channel
from grpc.aio import Channel as AsyncChannel
from grpc.aio import Server as AsyncServer
from sentry_sdk.integrations import Integration
from .client import ClientInterceptor
from .server import ServerInterceptor
from .aio.server import ServerInterceptor as AsyncServerInterceptor
from .aio.client import (
SentryUnaryUnaryClientInterceptor as AsyncUnaryUnaryClientInterceptor,
)
from .aio.client import (
SentryUnaryStreamClientInterceptor as AsyncUnaryStreamClientIntercetor,
)
from typing import TYPE_CHECKING, Any, Optional, Sequence
# Hack to get new Python features working in older versions
# without introducing a hard dependency on `typing_extensions`
# from: https://stackoverflow.com/a/71944042/300572
if TYPE_CHECKING:
from typing import ParamSpec, Callable
else:
# Fake ParamSpec
class ParamSpec:
def __init__(self, _):
self.args = None
self.kwargs = None
# Callable[anything] will return None
class _Callable:
def __getitem__(self, _):
return None
# Make instances
Callable = _Callable()
P = ParamSpec("P")
def _wrap_channel_sync(func: Callable[P, Channel]) -> Callable[P, Channel]:
"Wrapper for synchronous secure and insecure channel."
@wraps(func)
def patched_channel(*args: Any, **kwargs: Any) -> Channel:
channel = func(*args, **kwargs)
if not ClientInterceptor._is_intercepted:
ClientInterceptor._is_intercepted = True
return intercept_channel(channel, ClientInterceptor())
else:
return channel
return patched_channel
def _wrap_intercept_channel(func: Callable[P, Channel]) -> Callable[P, Channel]:
@wraps(func)
def patched_intercept_channel(
channel: Channel, *interceptors: grpc.ServerInterceptor
) -> Channel:
if ClientInterceptor._is_intercepted:
interceptors = tuple(
[
interceptor
for interceptor in interceptors
if not isinstance(interceptor, ClientInterceptor)
]
)
else:
interceptors = interceptors
return intercept_channel(channel, *interceptors)
return patched_intercept_channel # type: ignore
def _wrap_channel_async(func: Callable[P, AsyncChannel]) -> Callable[P, AsyncChannel]:
"Wrapper for asynchronous secure and insecure channel."
@wraps(func)
def patched_channel( # type: ignore
*args: P.args,
interceptors: Optional[Sequence[grpc.aio.ClientInterceptor]] = None,
**kwargs: P.kwargs,
) -> Channel:
sentry_interceptors = [
AsyncUnaryUnaryClientInterceptor(),
AsyncUnaryStreamClientIntercetor(),
]
interceptors = [*sentry_interceptors, *(interceptors or [])]
return func(*args, interceptors=interceptors, **kwargs) # type: ignore
return patched_channel # type: ignore
def _wrap_sync_server(func: Callable[P, Server]) -> Callable[P, Server]:
"""Wrapper for synchronous server."""
@wraps(func)
def patched_server( # type: ignore
*args: P.args,
interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
**kwargs: P.kwargs,
) -> Server:
interceptors = [
interceptor
for interceptor in interceptors or []
if not isinstance(interceptor, ServerInterceptor)
]
server_interceptor = ServerInterceptor()
interceptors = [server_interceptor, *(interceptors or [])]
return func(*args, interceptors=interceptors, **kwargs) # type: ignore
return patched_server # type: ignore
def _wrap_async_server(func: Callable[P, AsyncServer]) -> Callable[P, AsyncServer]:
"""Wrapper for asynchronous server."""
@wraps(func)
def patched_aio_server( # type: ignore
*args: P.args,
interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
**kwargs: P.kwargs,
) -> Server:
server_interceptor = AsyncServerInterceptor()
interceptors = (server_interceptor, *(interceptors or []))
return func(*args, interceptors=interceptors, **kwargs) # type: ignore
return patched_aio_server # type: ignore
class GRPCIntegration(Integration):
identifier = "grpc"
@staticmethod
def setup_once() -> None:
import grpc
grpc.insecure_channel = _wrap_channel_sync(grpc.insecure_channel)
grpc.secure_channel = _wrap_channel_sync(grpc.secure_channel)
grpc.intercept_channel = _wrap_intercept_channel(grpc.intercept_channel)
grpc.aio.insecure_channel = _wrap_channel_async(grpc.aio.insecure_channel)
grpc.aio.secure_channel = _wrap_channel_async(grpc.aio.secure_channel)
grpc.server = _wrap_sync_server(grpc.server)
grpc.aio.server = _wrap_async_server(grpc.aio.server)
@@ -0,0 +1,7 @@
from .server import ServerInterceptor
from .client import ClientInterceptor
__all__ = [
"ClientInterceptor",
"ServerInterceptor",
]
@@ -0,0 +1,94 @@
from typing import Callable, Union, AsyncIterable, Any
from grpc.aio import (
UnaryUnaryClientInterceptor,
UnaryStreamClientInterceptor,
ClientCallDetails,
UnaryUnaryCall,
UnaryStreamCall,
Metadata,
)
from google.protobuf.message import Message
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
class ClientInterceptor:
@staticmethod
def _update_client_call_details_metadata_from_scope(
client_call_details: ClientCallDetails,
) -> ClientCallDetails:
if client_call_details.metadata is None:
client_call_details = client_call_details._replace(metadata=Metadata())
elif not isinstance(client_call_details.metadata, Metadata):
# This is a workaround for a GRPC bug, which was fixed in grpcio v1.60.0
# See https://github.com/grpc/grpc/issues/34298.
client_call_details = client_call_details._replace(
metadata=Metadata.from_tuple(client_call_details.metadata)
)
for (
key,
value,
) in sentry_sdk.get_current_scope().iter_trace_propagation_headers():
client_call_details.metadata.add(key, value)
return client_call_details
class SentryUnaryUnaryClientInterceptor(ClientInterceptor, UnaryUnaryClientInterceptor): # type: ignore
async def intercept_unary_unary(
self,
continuation: Callable[[ClientCallDetails, Message], UnaryUnaryCall],
client_call_details: ClientCallDetails,
request: Message,
) -> Union[UnaryUnaryCall, Message]:
method = client_call_details.method
with sentry_sdk.start_span(
op=OP.GRPC_CLIENT,
name="unary unary call to %s" % method.decode(),
origin=SPAN_ORIGIN,
) as span:
span.set_data("type", "unary unary")
span.set_data("method", method)
client_call_details = self._update_client_call_details_metadata_from_scope(
client_call_details
)
response = await continuation(client_call_details, request)
status_code = await response.code()
span.set_data("code", status_code.name)
return response
class SentryUnaryStreamClientInterceptor(
ClientInterceptor, UnaryStreamClientInterceptor # type: ignore
):
async def intercept_unary_stream(
self,
continuation: Callable[[ClientCallDetails, Message], UnaryStreamCall],
client_call_details: ClientCallDetails,
request: Message,
) -> Union[AsyncIterable[Any], UnaryStreamCall]:
method = client_call_details.method
with sentry_sdk.start_span(
op=OP.GRPC_CLIENT,
name="unary stream call to %s" % method.decode(),
origin=SPAN_ORIGIN,
) as span:
span.set_data("type", "unary stream")
span.set_data("method", method)
client_call_details = self._update_client_call_details_metadata_from_scope(
client_call_details
)
response = await continuation(client_call_details, request)
# status_code = await response.code()
# span.set_data("code", status_code)
return response
@@ -0,0 +1,100 @@
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_CUSTOM
from sentry_sdk.utils import event_from_exception
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from typing import Any, Optional
try:
import grpc
from grpc import HandlerCallDetails, RpcMethodHandler
from grpc.aio import AbortError, ServicerContext
except ImportError:
raise DidNotEnable("grpcio is not installed")
class ServerInterceptor(grpc.aio.ServerInterceptor): # type: ignore
def __init__(self, find_name=None):
# type: (ServerInterceptor, Callable[[ServicerContext], str] | None) -> None
self._find_method_name = find_name or self._find_name
super().__init__()
async def intercept_service(self, continuation, handler_call_details):
# type: (ServerInterceptor, Callable[[HandlerCallDetails], Awaitable[RpcMethodHandler]], HandlerCallDetails) -> Optional[Awaitable[RpcMethodHandler]]
self._handler_call_details = handler_call_details
handler = await continuation(handler_call_details)
if handler is None:
return None
if not handler.request_streaming and not handler.response_streaming:
handler_factory = grpc.unary_unary_rpc_method_handler
async def wrapped(request, context):
# type: (Any, ServicerContext) -> Any
name = self._find_method_name(context)
if not name:
return await handler(request, context)
# What if the headers are empty?
transaction = Transaction.continue_from_headers(
dict(context.invocation_metadata()),
op=OP.GRPC_SERVER,
name=name,
source=TRANSACTION_SOURCE_CUSTOM,
origin=SPAN_ORIGIN,
)
with sentry_sdk.start_transaction(transaction=transaction):
try:
return await handler.unary_unary(request, context)
except AbortError:
raise
except Exception as exc:
event, hint = event_from_exception(
exc,
mechanism={"type": "grpc", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
raise
elif not handler.request_streaming and handler.response_streaming:
handler_factory = grpc.unary_stream_rpc_method_handler
async def wrapped(request, context): # type: ignore
# type: (Any, ServicerContext) -> Any
async for r in handler.unary_stream(request, context):
yield r
elif handler.request_streaming and not handler.response_streaming:
handler_factory = grpc.stream_unary_rpc_method_handler
async def wrapped(request, context):
# type: (Any, ServicerContext) -> Any
response = handler.stream_unary(request, context)
return await response
elif handler.request_streaming and handler.response_streaming:
handler_factory = grpc.stream_stream_rpc_method_handler
async def wrapped(request, context): # type: ignore
# type: (Any, ServicerContext) -> Any
async for r in handler.stream_stream(request, context):
yield r
return handler_factory(
wrapped,
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)
def _find_name(self, context):
# type: (ServicerContext) -> str
return self._handler_call_details.method
@@ -0,0 +1,92 @@
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Iterator, Iterable, Union
try:
import grpc
from grpc import ClientCallDetails, Call
from grpc._interceptor import _UnaryOutcome
from grpc.aio._interceptor import UnaryStreamCall
from google.protobuf.message import Message
except ImportError:
raise DidNotEnable("grpcio is not installed")
class ClientInterceptor(
grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor # type: ignore
):
_is_intercepted = False
def intercept_unary_unary(self, continuation, client_call_details, request):
# type: (ClientInterceptor, Callable[[ClientCallDetails, Message], _UnaryOutcome], ClientCallDetails, Message) -> _UnaryOutcome
method = client_call_details.method
with sentry_sdk.start_span(
op=OP.GRPC_CLIENT,
name="unary unary call to %s" % method,
origin=SPAN_ORIGIN,
) as span:
span.set_data("type", "unary unary")
span.set_data("method", method)
client_call_details = self._update_client_call_details_metadata_from_scope(
client_call_details
)
response = continuation(client_call_details, request)
span.set_data("code", response.code().name)
return response
def intercept_unary_stream(self, continuation, client_call_details, request):
# type: (ClientInterceptor, Callable[[ClientCallDetails, Message], Union[Iterable[Any], UnaryStreamCall]], ClientCallDetails, Message) -> Union[Iterator[Message], Call]
method = client_call_details.method
with sentry_sdk.start_span(
op=OP.GRPC_CLIENT,
name="unary stream call to %s" % method,
origin=SPAN_ORIGIN,
) as span:
span.set_data("type", "unary stream")
span.set_data("method", method)
client_call_details = self._update_client_call_details_metadata_from_scope(
client_call_details
)
response = continuation(
client_call_details, request
) # type: UnaryStreamCall
# Setting code on unary-stream leads to execution getting stuck
# span.set_data("code", response.code().name)
return response
@staticmethod
def _update_client_call_details_metadata_from_scope(client_call_details):
# type: (ClientCallDetails) -> ClientCallDetails
metadata = (
list(client_call_details.metadata) if client_call_details.metadata else []
)
for (
key,
value,
) in sentry_sdk.get_current_scope().iter_trace_propagation_headers():
metadata.append((key, value))
client_call_details = grpc._interceptor._ClientCallDetails(
method=client_call_details.method,
timeout=client_call_details.timeout,
metadata=metadata,
credentials=client_call_details.credentials,
wait_for_ready=client_call_details.wait_for_ready,
compression=client_call_details.compression,
)
return client_call_details
@@ -0,0 +1 @@
SPAN_ORIGIN = "auto.grpc.grpc"
@@ -0,0 +1,66 @@
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.integrations.grpc.consts import SPAN_ORIGIN
from sentry_sdk.tracing import Transaction, TRANSACTION_SOURCE_CUSTOM
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Callable, Optional
from google.protobuf.message import Message
try:
import grpc
from grpc import ServicerContext, HandlerCallDetails, RpcMethodHandler
except ImportError:
raise DidNotEnable("grpcio is not installed")
class ServerInterceptor(grpc.ServerInterceptor): # type: ignore
def __init__(self, find_name=None):
# type: (ServerInterceptor, Optional[Callable[[ServicerContext], str]]) -> None
self._find_method_name = find_name or ServerInterceptor._find_name
super().__init__()
def intercept_service(self, continuation, handler_call_details):
# type: (ServerInterceptor, Callable[[HandlerCallDetails], RpcMethodHandler], HandlerCallDetails) -> RpcMethodHandler
handler = continuation(handler_call_details)
if not handler or not handler.unary_unary:
return handler
def behavior(request, context):
# type: (Message, ServicerContext) -> Message
with sentry_sdk.isolation_scope():
name = self._find_method_name(context)
if name:
metadata = dict(context.invocation_metadata())
transaction = Transaction.continue_from_headers(
metadata,
op=OP.GRPC_SERVER,
name=name,
source=TRANSACTION_SOURCE_CUSTOM,
origin=SPAN_ORIGIN,
)
with sentry_sdk.start_transaction(transaction=transaction):
try:
return handler.unary_unary(request, context)
except BaseException as e:
raise e
else:
return handler.unary_unary(request, context)
return grpc.unary_unary_rpc_method_handler(
behavior,
request_deserializer=handler.request_deserializer,
response_serializer=handler.response_serializer,
)
@staticmethod
def _find_name(context):
# type: (ServicerContext) -> str
return context._rpc_event.call_details.method.decode()
@@ -0,0 +1,167 @@
import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import Integration, DidNotEnable
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
from sentry_sdk.tracing_utils import Baggage, should_propagate_trace
from sentry_sdk.utils import (
SENSITIVE_DATA_SUBSTITUTE,
capture_internal_exceptions,
ensure_integration_enabled,
logger,
parse_url,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import MutableMapping
from typing import Any
try:
from httpx import AsyncClient, Client, Request, Response # type: ignore
except ImportError:
raise DidNotEnable("httpx is not installed")
__all__ = ["HttpxIntegration"]
class HttpxIntegration(Integration):
identifier = "httpx"
origin = f"auto.http.{identifier}"
@staticmethod
def setup_once():
# type: () -> None
"""
httpx has its own transport layer and can be customized when needed,
so patch Client.send and AsyncClient.send to support both synchronous and async interfaces.
"""
_install_httpx_client()
_install_httpx_async_client()
def _install_httpx_client():
# type: () -> None
real_send = Client.send
@ensure_integration_enabled(HttpxIntegration, real_send)
def send(self, request, **kwargs):
# type: (Client, Request, **Any) -> Response
parsed_url = None
with capture_internal_exceptions():
parsed_url = parse_url(str(request.url), sanitize=False)
with sentry_sdk.start_span(
op=OP.HTTP_CLIENT,
name="%s %s"
% (
request.method,
parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE,
),
origin=HttpxIntegration.origin,
) as span:
span.set_data(SPANDATA.HTTP_METHOD, request.method)
if parsed_url is not None:
span.set_data("url", parsed_url.url)
span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
if should_propagate_trace(sentry_sdk.get_client(), str(request.url)):
for (
key,
value,
) in sentry_sdk.get_current_scope().iter_trace_propagation_headers():
logger.debug(
"[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format(
key=key, value=value, url=request.url
)
)
if key == BAGGAGE_HEADER_NAME:
_add_sentry_baggage_to_headers(request.headers, value)
else:
request.headers[key] = value
rv = real_send(self, request, **kwargs)
span.set_http_status(rv.status_code)
span.set_data("reason", rv.reason_phrase)
return rv
Client.send = send
def _install_httpx_async_client():
# type: () -> None
real_send = AsyncClient.send
async def send(self, request, **kwargs):
# type: (AsyncClient, Request, **Any) -> Response
if sentry_sdk.get_client().get_integration(HttpxIntegration) is None:
return await real_send(self, request, **kwargs)
parsed_url = None
with capture_internal_exceptions():
parsed_url = parse_url(str(request.url), sanitize=False)
with sentry_sdk.start_span(
op=OP.HTTP_CLIENT,
name="%s %s"
% (
request.method,
parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE,
),
origin=HttpxIntegration.origin,
) as span:
span.set_data(SPANDATA.HTTP_METHOD, request.method)
if parsed_url is not None:
span.set_data("url", parsed_url.url)
span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
if should_propagate_trace(sentry_sdk.get_client(), str(request.url)):
for (
key,
value,
) in sentry_sdk.get_current_scope().iter_trace_propagation_headers():
logger.debug(
"[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format(
key=key, value=value, url=request.url
)
)
if key == BAGGAGE_HEADER_NAME and request.headers.get(
BAGGAGE_HEADER_NAME
):
# do not overwrite any existing baggage, just append to it
request.headers[key] += "," + value
else:
request.headers[key] = value
rv = await real_send(self, request, **kwargs)
span.set_http_status(rv.status_code)
span.set_data("reason", rv.reason_phrase)
return rv
AsyncClient.send = send
def _add_sentry_baggage_to_headers(headers, sentry_baggage):
# type: (MutableMapping[str, str], str) -> None
"""Add the Sentry baggage to the headers.
This function directly mutates the provided headers. The provided sentry_baggage
is appended to the existing baggage. If the baggage already contains Sentry items,
they are stripped out first.
"""
existing_baggage = headers.get(BAGGAGE_HEADER_NAME, "")
stripped_existing_baggage = Baggage.strip_sentry_baggage(existing_baggage)
separator = "," if len(stripped_existing_baggage) > 0 else ""
headers[BAGGAGE_HEADER_NAME] = (
stripped_existing_baggage + separator + sentry_baggage
)
@@ -0,0 +1,174 @@
import sys
from datetime import datetime
import sentry_sdk
from sentry_sdk.api import continue_trace, get_baggage, get_traceparent
from sentry_sdk.consts import OP, SPANSTATUS
from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import (
BAGGAGE_HEADER_NAME,
SENTRY_TRACE_HEADER_NAME,
TRANSACTION_SOURCE_TASK,
)
from sentry_sdk.utils import (
capture_internal_exceptions,
ensure_integration_enabled,
event_from_exception,
SENSITIVE_DATA_SUBSTITUTE,
reraise,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Callable, Optional, Union, TypeVar
from sentry_sdk._types import EventProcessor, Event, Hint
from sentry_sdk.utils import ExcInfo
F = TypeVar("F", bound=Callable[..., Any])
try:
from huey.api import Huey, Result, ResultGroup, Task, PeriodicTask
from huey.exceptions import CancelExecution, RetryTask, TaskLockedException
except ImportError:
raise DidNotEnable("Huey is not installed")
HUEY_CONTROL_FLOW_EXCEPTIONS = (CancelExecution, RetryTask, TaskLockedException)
class HueyIntegration(Integration):
identifier = "huey"
origin = f"auto.queue.{identifier}"
@staticmethod
def setup_once():
# type: () -> None
patch_enqueue()
patch_execute()
def patch_enqueue():
# type: () -> None
old_enqueue = Huey.enqueue
@ensure_integration_enabled(HueyIntegration, old_enqueue)
def _sentry_enqueue(self, task):
# type: (Huey, Task) -> Optional[Union[Result, ResultGroup]]
with sentry_sdk.start_span(
op=OP.QUEUE_SUBMIT_HUEY,
name=task.name,
origin=HueyIntegration.origin,
):
if not isinstance(task, PeriodicTask):
# Attach trace propagation data to task kwargs. We do
# not do this for periodic tasks, as these don't
# really have an originating transaction.
task.kwargs["sentry_headers"] = {
BAGGAGE_HEADER_NAME: get_baggage(),
SENTRY_TRACE_HEADER_NAME: get_traceparent(),
}
return old_enqueue(self, task)
Huey.enqueue = _sentry_enqueue
def _make_event_processor(task):
# type: (Any) -> EventProcessor
def event_processor(event, hint):
# type: (Event, Hint) -> Optional[Event]
with capture_internal_exceptions():
tags = event.setdefault("tags", {})
tags["huey_task_id"] = task.id
tags["huey_task_retry"] = task.default_retries > task.retries
extra = event.setdefault("extra", {})
extra["huey-job"] = {
"task": task.name,
"args": (
task.args
if should_send_default_pii()
else SENSITIVE_DATA_SUBSTITUTE
),
"kwargs": (
task.kwargs
if should_send_default_pii()
else SENSITIVE_DATA_SUBSTITUTE
),
"retry": (task.default_retries or 0) - task.retries,
}
return event
return event_processor
def _capture_exception(exc_info):
# type: (ExcInfo) -> None
scope = sentry_sdk.get_current_scope()
if exc_info[0] in HUEY_CONTROL_FLOW_EXCEPTIONS:
scope.transaction.set_status(SPANSTATUS.ABORTED)
return
scope.transaction.set_status(SPANSTATUS.INTERNAL_ERROR)
event, hint = event_from_exception(
exc_info,
client_options=sentry_sdk.get_client().options,
mechanism={"type": HueyIntegration.identifier, "handled": False},
)
scope.capture_event(event, hint=hint)
def _wrap_task_execute(func):
# type: (F) -> F
@ensure_integration_enabled(HueyIntegration, func)
def _sentry_execute(*args, **kwargs):
# type: (*Any, **Any) -> Any
try:
result = func(*args, **kwargs)
except Exception:
exc_info = sys.exc_info()
_capture_exception(exc_info)
reraise(*exc_info)
return result
return _sentry_execute # type: ignore
def patch_execute():
# type: () -> None
old_execute = Huey._execute
@ensure_integration_enabled(HueyIntegration, old_execute)
def _sentry_execute(self, task, timestamp=None):
# type: (Huey, Task, Optional[datetime]) -> Any
with sentry_sdk.isolation_scope() as scope:
with capture_internal_exceptions():
scope._name = "huey"
scope.clear_breadcrumbs()
scope.add_event_processor(_make_event_processor(task))
sentry_headers = task.kwargs.pop("sentry_headers", None)
transaction = continue_trace(
sentry_headers or {},
name=task.name,
op=OP.QUEUE_TASK_HUEY,
source=TRANSACTION_SOURCE_TASK,
origin=HueyIntegration.origin,
)
transaction.set_status(SPANSTATUS.OK)
if not getattr(task, "_sentry_is_patched", False):
task.execute = _wrap_task_execute(task.execute)
task._sentry_is_patched = True
with sentry_sdk.start_transaction(transaction):
return old_execute(self, task, timestamp)
Huey._execute = _sentry_execute
@@ -0,0 +1,175 @@
from functools import wraps
from sentry_sdk import consts
from sentry_sdk.ai.monitoring import record_token_usage
from sentry_sdk.ai.utils import set_data_normalized
from sentry_sdk.consts import SPANDATA
from typing import Any, Iterable, Callable
import sentry_sdk
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.utils import (
capture_internal_exceptions,
event_from_exception,
)
try:
import huggingface_hub.inference._client
from huggingface_hub import ChatCompletionStreamOutput, TextGenerationOutput
except ImportError:
raise DidNotEnable("Huggingface not installed")
class HuggingfaceHubIntegration(Integration):
identifier = "huggingface_hub"
origin = f"auto.ai.{identifier}"
def __init__(self, include_prompts=True):
# type: (HuggingfaceHubIntegration, bool) -> None
self.include_prompts = include_prompts
@staticmethod
def setup_once():
# type: () -> None
huggingface_hub.inference._client.InferenceClient.text_generation = (
_wrap_text_generation(
huggingface_hub.inference._client.InferenceClient.text_generation
)
)
def _capture_exception(exc):
# type: (Any) -> None
event, hint = event_from_exception(
exc,
client_options=sentry_sdk.get_client().options,
mechanism={"type": "huggingface_hub", "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
def _wrap_text_generation(f):
# type: (Callable[..., Any]) -> Callable[..., Any]
@wraps(f)
def new_text_generation(*args, **kwargs):
# type: (*Any, **Any) -> Any
integration = sentry_sdk.get_client().get_integration(HuggingfaceHubIntegration)
if integration is None:
return f(*args, **kwargs)
if "prompt" in kwargs:
prompt = kwargs["prompt"]
elif len(args) >= 2:
kwargs["prompt"] = args[1]
prompt = kwargs["prompt"]
args = (args[0],) + args[2:]
else:
# invalid call, let it return error
return f(*args, **kwargs)
model = kwargs.get("model")
streaming = kwargs.get("stream")
span = sentry_sdk.start_span(
op=consts.OP.HUGGINGFACE_HUB_CHAT_COMPLETIONS_CREATE,
name="Text Generation",
origin=HuggingfaceHubIntegration.origin,
)
span.__enter__()
try:
res = f(*args, **kwargs)
except Exception as e:
_capture_exception(e)
span.__exit__(None, None, None)
raise e from None
with capture_internal_exceptions():
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, prompt)
set_data_normalized(span, SPANDATA.AI_MODEL_ID, model)
set_data_normalized(span, SPANDATA.AI_STREAMING, streaming)
if isinstance(res, str):
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(
span,
"ai.responses",
[res],
)
span.__exit__(None, None, None)
return res
if isinstance(res, TextGenerationOutput):
if should_send_default_pii() and integration.include_prompts:
set_data_normalized(
span,
"ai.responses",
[res.generated_text],
)
if res.details is not None and res.details.generated_tokens > 0:
record_token_usage(span, total_tokens=res.details.generated_tokens)
span.__exit__(None, None, None)
return res
if not isinstance(res, Iterable):
# we only know how to deal with strings and iterables, ignore
set_data_normalized(span, "unknown_response", True)
span.__exit__(None, None, None)
return res
if kwargs.get("details", False):
# res is Iterable[TextGenerationStreamOutput]
def new_details_iterator():
# type: () -> Iterable[ChatCompletionStreamOutput]
with capture_internal_exceptions():
tokens_used = 0
data_buf: list[str] = []
for x in res:
if hasattr(x, "token") and hasattr(x.token, "text"):
data_buf.append(x.token.text)
if hasattr(x, "details") and hasattr(
x.details, "generated_tokens"
):
tokens_used = x.details.generated_tokens
yield x
if (
len(data_buf) > 0
and should_send_default_pii()
and integration.include_prompts
):
set_data_normalized(
span, SPANDATA.AI_RESPONSES, "".join(data_buf)
)
if tokens_used > 0:
record_token_usage(span, total_tokens=tokens_used)
span.__exit__(None, None, None)
return new_details_iterator()
else:
# res is Iterable[str]
def new_iterator():
# type: () -> Iterable[str]
data_buf: list[str] = []
with capture_internal_exceptions():
for s in res:
if isinstance(s, str):
data_buf.append(s)
yield s
if (
len(data_buf) > 0
and should_send_default_pii()
and integration.include_prompts
):
set_data_normalized(
span, SPANDATA.AI_RESPONSES, "".join(data_buf)
)
span.__exit__(None, None, None)
return new_iterator()
return new_text_generation
@@ -0,0 +1,465 @@
from collections import OrderedDict
from functools import wraps
import sentry_sdk
from sentry_sdk.ai.monitoring import set_ai_pipeline_name, record_token_usage
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.ai.utils import set_data_normalized
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import Span
from sentry_sdk.integrations import DidNotEnable, Integration
from sentry_sdk.utils import logger, capture_internal_exceptions
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, List, Callable, Dict, Union, Optional
from uuid import UUID
try:
from langchain_core.messages import BaseMessage
from langchain_core.outputs import LLMResult
from langchain_core.callbacks import (
manager,
BaseCallbackHandler,
)
from langchain_core.agents import AgentAction, AgentFinish
except ImportError:
raise DidNotEnable("langchain not installed")
DATA_FIELDS = {
"temperature": SPANDATA.AI_TEMPERATURE,
"top_p": SPANDATA.AI_TOP_P,
"top_k": SPANDATA.AI_TOP_K,
"function_call": SPANDATA.AI_FUNCTION_CALL,
"tool_calls": SPANDATA.AI_TOOL_CALLS,
"tools": SPANDATA.AI_TOOLS,
"response_format": SPANDATA.AI_RESPONSE_FORMAT,
"logit_bias": SPANDATA.AI_LOGIT_BIAS,
"tags": SPANDATA.AI_TAGS,
}
# To avoid double collecting tokens, we do *not* measure
# token counts for models for which we have an explicit integration
NO_COLLECT_TOKEN_MODELS = [
"openai-chat",
"anthropic-chat",
"cohere-chat",
"huggingface_endpoint",
]
class LangchainIntegration(Integration):
identifier = "langchain"
origin = f"auto.ai.{identifier}"
# The most number of spans (e.g., LLM calls) that can be processed at the same time.
max_spans = 1024
def __init__(
self, include_prompts=True, max_spans=1024, tiktoken_encoding_name=None
):
# type: (LangchainIntegration, bool, int, Optional[str]) -> None
self.include_prompts = include_prompts
self.max_spans = max_spans
self.tiktoken_encoding_name = tiktoken_encoding_name
@staticmethod
def setup_once():
# type: () -> None
manager._configure = _wrap_configure(manager._configure)
class WatchedSpan:
span = None # type: Span
num_completion_tokens = 0 # type: int
num_prompt_tokens = 0 # type: int
no_collect_tokens = False # type: bool
children = [] # type: List[WatchedSpan]
is_pipeline = False # type: bool
def __init__(self, span):
# type: (Span) -> None
self.span = span
class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
"""Base callback handler that can be used to handle callbacks from langchain."""
span_map = OrderedDict() # type: OrderedDict[UUID, WatchedSpan]
max_span_map_size = 0
def __init__(self, max_span_map_size, include_prompts, tiktoken_encoding_name=None):
# type: (int, bool, Optional[str]) -> None
self.max_span_map_size = max_span_map_size
self.include_prompts = include_prompts
self.tiktoken_encoding = None
if tiktoken_encoding_name is not None:
import tiktoken # type: ignore
self.tiktoken_encoding = tiktoken.get_encoding(tiktoken_encoding_name)
def count_tokens(self, s):
# type: (str) -> int
if self.tiktoken_encoding is not None:
return len(self.tiktoken_encoding.encode_ordinary(s))
return 0
def gc_span_map(self):
# type: () -> None
while len(self.span_map) > self.max_span_map_size:
run_id, watched_span = self.span_map.popitem(last=False)
self._exit_span(watched_span, run_id)
def _handle_error(self, run_id, error):
# type: (UUID, Any) -> None
if not run_id or run_id not in self.span_map:
return
span_data = self.span_map[run_id]
if not span_data:
return
sentry_sdk.capture_exception(error, span_data.span.scope)
span_data.span.__exit__(None, None, None)
del self.span_map[run_id]
def _normalize_langchain_message(self, message):
# type: (BaseMessage) -> Any
parsed = {"content": message.content, "role": message.type}
parsed.update(message.additional_kwargs)
return parsed
def _create_span(self, run_id, parent_id, **kwargs):
# type: (SentryLangchainCallback, UUID, Optional[Any], Any) -> WatchedSpan
watched_span = None # type: Optional[WatchedSpan]
if parent_id:
parent_span = self.span_map.get(parent_id) # type: Optional[WatchedSpan]
if parent_span:
watched_span = WatchedSpan(parent_span.span.start_child(**kwargs))
parent_span.children.append(watched_span)
if watched_span is None:
watched_span = WatchedSpan(sentry_sdk.start_span(**kwargs))
if kwargs.get("op", "").startswith("ai.pipeline."):
if kwargs.get("name"):
set_ai_pipeline_name(kwargs.get("name"))
watched_span.is_pipeline = True
watched_span.span.__enter__()
self.span_map[run_id] = watched_span
self.gc_span_map()
return watched_span
def _exit_span(self, span_data, run_id):
# type: (SentryLangchainCallback, WatchedSpan, UUID) -> None
if span_data.is_pipeline:
set_ai_pipeline_name(None)
span_data.span.__exit__(None, None, None)
del self.span_map[run_id]
def on_llm_start(
self,
serialized,
prompts,
*,
run_id,
tags=None,
parent_run_id=None,
metadata=None,
**kwargs,
):
# type: (SentryLangchainCallback, Dict[str, Any], List[str], UUID, Optional[List[str]], Optional[UUID], Optional[Dict[str, Any]], Any) -> Any
"""Run when LLM starts running."""
with capture_internal_exceptions():
if not run_id:
return
all_params = kwargs.get("invocation_params", {})
all_params.update(serialized.get("kwargs", {}))
watched_span = self._create_span(
run_id,
kwargs.get("parent_run_id"),
op=OP.LANGCHAIN_RUN,
name=kwargs.get("name") or "Langchain LLM call",
origin=LangchainIntegration.origin,
)
span = watched_span.span
if should_send_default_pii() and self.include_prompts:
set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, prompts)
for k, v in DATA_FIELDS.items():
if k in all_params:
set_data_normalized(span, v, all_params[k])
def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
# type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
"""Run when Chat Model starts running."""
with capture_internal_exceptions():
if not run_id:
return
all_params = kwargs.get("invocation_params", {})
all_params.update(serialized.get("kwargs", {}))
watched_span = self._create_span(
run_id,
kwargs.get("parent_run_id"),
op=OP.LANGCHAIN_CHAT_COMPLETIONS_CREATE,
name=kwargs.get("name") or "Langchain Chat Model",
origin=LangchainIntegration.origin,
)
span = watched_span.span
model = all_params.get(
"model", all_params.get("model_name", all_params.get("model_id"))
)
watched_span.no_collect_tokens = any(
x in all_params.get("_type", "") for x in NO_COLLECT_TOKEN_MODELS
)
if not model and "anthropic" in all_params.get("_type"):
model = "claude-2"
if model:
span.set_data(SPANDATA.AI_MODEL_ID, model)
if should_send_default_pii() and self.include_prompts:
set_data_normalized(
span,
SPANDATA.AI_INPUT_MESSAGES,
[
[self._normalize_langchain_message(x) for x in list_]
for list_ in messages
],
)
for k, v in DATA_FIELDS.items():
if k in all_params:
set_data_normalized(span, v, all_params[k])
if not watched_span.no_collect_tokens:
for list_ in messages:
for message in list_:
self.span_map[run_id].num_prompt_tokens += self.count_tokens(
message.content
) + self.count_tokens(message.type)
def on_llm_new_token(self, token, *, run_id, **kwargs):
# type: (SentryLangchainCallback, str, UUID, Any) -> Any
"""Run on new LLM token. Only available when streaming is enabled."""
with capture_internal_exceptions():
if not run_id or run_id not in self.span_map:
return
span_data = self.span_map[run_id]
if not span_data or span_data.no_collect_tokens:
return
span_data.num_completion_tokens += self.count_tokens(token)
def on_llm_end(self, response, *, run_id, **kwargs):
# type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
"""Run when LLM ends running."""
with capture_internal_exceptions():
if not run_id:
return
token_usage = (
response.llm_output.get("token_usage") if response.llm_output else None
)
span_data = self.span_map[run_id]
if not span_data:
return
if should_send_default_pii() and self.include_prompts:
set_data_normalized(
span_data.span,
SPANDATA.AI_RESPONSES,
[[x.text for x in list_] for list_ in response.generations],
)
if not span_data.no_collect_tokens:
if token_usage:
record_token_usage(
span_data.span,
token_usage.get("prompt_tokens"),
token_usage.get("completion_tokens"),
token_usage.get("total_tokens"),
)
else:
record_token_usage(
span_data.span,
span_data.num_prompt_tokens,
span_data.num_completion_tokens,
)
self._exit_span(span_data, run_id)
def on_llm_error(self, error, *, run_id, **kwargs):
# type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
"""Run when LLM errors."""
with capture_internal_exceptions():
self._handle_error(run_id, error)
def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
# type: (SentryLangchainCallback, Dict[str, Any], Dict[str, Any], UUID, Any) -> Any
"""Run when chain starts running."""
with capture_internal_exceptions():
if not run_id:
return
watched_span = self._create_span(
run_id,
kwargs.get("parent_run_id"),
op=(
OP.LANGCHAIN_RUN
if kwargs.get("parent_run_id") is not None
else OP.LANGCHAIN_PIPELINE
),
name=kwargs.get("name") or "Chain execution",
origin=LangchainIntegration.origin,
)
metadata = kwargs.get("metadata")
if metadata:
set_data_normalized(watched_span.span, SPANDATA.AI_METADATA, metadata)
def on_chain_end(self, outputs, *, run_id, **kwargs):
# type: (SentryLangchainCallback, Dict[str, Any], UUID, Any) -> Any
"""Run when chain ends running."""
with capture_internal_exceptions():
if not run_id or run_id not in self.span_map:
return
span_data = self.span_map[run_id]
if not span_data:
return
self._exit_span(span_data, run_id)
def on_chain_error(self, error, *, run_id, **kwargs):
# type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
"""Run when chain errors."""
self._handle_error(run_id, error)
def on_agent_action(self, action, *, run_id, **kwargs):
# type: (SentryLangchainCallback, AgentAction, UUID, Any) -> Any
with capture_internal_exceptions():
if not run_id:
return
watched_span = self._create_span(
run_id,
kwargs.get("parent_run_id"),
op=OP.LANGCHAIN_AGENT,
name=action.tool or "AI tool usage",
origin=LangchainIntegration.origin,
)
if action.tool_input and should_send_default_pii() and self.include_prompts:
set_data_normalized(
watched_span.span, SPANDATA.AI_INPUT_MESSAGES, action.tool_input
)
def on_agent_finish(self, finish, *, run_id, **kwargs):
# type: (SentryLangchainCallback, AgentFinish, UUID, Any) -> Any
with capture_internal_exceptions():
if not run_id:
return
span_data = self.span_map[run_id]
if not span_data:
return
if should_send_default_pii() and self.include_prompts:
set_data_normalized(
span_data.span, SPANDATA.AI_RESPONSES, finish.return_values.items()
)
self._exit_span(span_data, run_id)
def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
# type: (SentryLangchainCallback, Dict[str, Any], str, UUID, Any) -> Any
"""Run when tool starts running."""
with capture_internal_exceptions():
if not run_id:
return
watched_span = self._create_span(
run_id,
kwargs.get("parent_run_id"),
op=OP.LANGCHAIN_TOOL,
name=serialized.get("name") or kwargs.get("name") or "AI tool usage",
origin=LangchainIntegration.origin,
)
if should_send_default_pii() and self.include_prompts:
set_data_normalized(
watched_span.span,
SPANDATA.AI_INPUT_MESSAGES,
kwargs.get("inputs", [input_str]),
)
if kwargs.get("metadata"):
set_data_normalized(
watched_span.span, SPANDATA.AI_METADATA, kwargs.get("metadata")
)
def on_tool_end(self, output, *, run_id, **kwargs):
# type: (SentryLangchainCallback, str, UUID, Any) -> Any
"""Run when tool ends running."""
with capture_internal_exceptions():
if not run_id or run_id not in self.span_map:
return
span_data = self.span_map[run_id]
if not span_data:
return
if should_send_default_pii() and self.include_prompts:
set_data_normalized(span_data.span, SPANDATA.AI_RESPONSES, output)
self._exit_span(span_data, run_id)
def on_tool_error(self, error, *args, run_id, **kwargs):
# type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
"""Run when tool errors."""
self._handle_error(run_id, error)
def _wrap_configure(f):
# type: (Callable[..., Any]) -> Callable[..., Any]
@wraps(f)
def new_configure(*args, **kwargs):
# type: (Any, Any) -> Any
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
if integration is None:
return f(*args, **kwargs)
with capture_internal_exceptions():
new_callbacks = [] # type: List[BaseCallbackHandler]
if "local_callbacks" in kwargs:
existing_callbacks = kwargs["local_callbacks"]
kwargs["local_callbacks"] = new_callbacks
elif len(args) > 2:
existing_callbacks = args[2]
args = (
args[0],
args[1],
new_callbacks,
) + args[3:]
else:
existing_callbacks = []
if existing_callbacks:
if isinstance(existing_callbacks, list):
for cb in existing_callbacks:
new_callbacks.append(cb)
elif isinstance(existing_callbacks, BaseCallbackHandler):
new_callbacks.append(existing_callbacks)
else:
logger.debug("Unknown callback type: %s", existing_callbacks)
already_added = False
for callback in new_callbacks:
if isinstance(callback, SentryLangchainCallback):
already_added = True
if not already_added:
new_callbacks.append(
SentryLangchainCallback(
integration.max_spans,
integration.include_prompts,
integration.tiktoken_encoding_name,
)
)
return f(*args, **kwargs)
return new_configure
@@ -0,0 +1,62 @@
from typing import TYPE_CHECKING
import sentry_sdk
from sentry_sdk.integrations import DidNotEnable, Integration
try:
import ldclient
from ldclient.hook import Hook, Metadata
if TYPE_CHECKING:
from ldclient import LDClient
from ldclient.hook import EvaluationSeriesContext
from ldclient.evaluation import EvaluationDetail
from typing import Any
except ImportError:
raise DidNotEnable("LaunchDarkly is not installed")
class LaunchDarklyIntegration(Integration):
identifier = "launchdarkly"
def __init__(self, ld_client=None):
# type: (LDClient | None) -> None
"""
:param client: An initialized LDClient instance. If a client is not provided, this
integration will attempt to use the shared global instance.
"""
try:
client = ld_client or ldclient.get()
except Exception as exc:
raise DidNotEnable("Error getting LaunchDarkly client. " + repr(exc))
if not client.is_initialized():
raise DidNotEnable("LaunchDarkly client is not initialized.")
# Register the flag collection hook with the LD client.
client.add_hook(LaunchDarklyHook())
@staticmethod
def setup_once():
# type: () -> None
pass
class LaunchDarklyHook(Hook):
@property
def metadata(self):
# type: () -> Metadata
return Metadata(name="sentry-flag-auditor")
def after_evaluation(self, series_context, data, detail):
# type: (EvaluationSeriesContext, dict[Any, Any], EvaluationDetail) -> dict[Any, Any]
if isinstance(detail.value, bool):
flags = sentry_sdk.get_current_scope().flags
flags.set(series_context.key, detail.value)
return data
def before_evaluation(self, series_context, data):
# type: (EvaluationSeriesContext, dict[Any, Any]) -> dict[Any, Any]
return data # No-op.
@@ -0,0 +1,306 @@
from collections.abc import Set
import sentry_sdk
from sentry_sdk.consts import OP
from sentry_sdk.integrations import (
_DEFAULT_FAILED_REQUEST_STATUS_CODES,
DidNotEnable,
Integration,
)
from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
from sentry_sdk.integrations.logging import ignore_logger
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.tracing import SOURCE_FOR_STYLE, TRANSACTION_SOURCE_ROUTE
from sentry_sdk.utils import (
ensure_integration_enabled,
event_from_exception,
transaction_from_function,
)
try:
from litestar import Request, Litestar # type: ignore
from litestar.handlers.base import BaseRouteHandler # type: ignore
from litestar.middleware import DefineMiddleware # type: ignore
from litestar.routes.http import HTTPRoute # type: ignore
from litestar.data_extractors import ConnectionDataExtractor # type: ignore
from litestar.exceptions import HTTPException # type: ignore
except ImportError:
raise DidNotEnable("Litestar is not installed")
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any, Optional, Union
from litestar.types.asgi_types import ASGIApp # type: ignore
from litestar.types import ( # type: ignore
HTTPReceiveMessage,
HTTPScope,
Message,
Middleware,
Receive,
Scope as LitestarScope,
Send,
WebSocketReceiveMessage,
)
from litestar.middleware import MiddlewareProtocol
from sentry_sdk._types import Event, Hint
_DEFAULT_TRANSACTION_NAME = "generic Litestar request"
class LitestarIntegration(Integration):
identifier = "litestar"
origin = f"auto.http.{identifier}"
def __init__(
self,
failed_request_status_codes=_DEFAULT_FAILED_REQUEST_STATUS_CODES, # type: Set[int]
) -> None:
self.failed_request_status_codes = failed_request_status_codes
@staticmethod
def setup_once():
# type: () -> None
patch_app_init()
patch_middlewares()
patch_http_route_handle()
# The following line follows the pattern found in other integrations such as `DjangoIntegration.setup_once`.
# The Litestar `ExceptionHandlerMiddleware.__call__` catches exceptions and does the following
# (among other things):
# 1. Logs them, some at least (such as 500s) as errors
# 2. Calls after_exception hooks
# The `LitestarIntegration`` provides an after_exception hook (see `patch_app_init` below) to create a Sentry event
# from an exception, which ends up being called during step 2 above. However, the Sentry `LoggingIntegration` will
# by default create a Sentry event from error logs made in step 1 if we do not prevent it from doing so.
ignore_logger("litestar")
class SentryLitestarASGIMiddleware(SentryAsgiMiddleware):
def __init__(self, app, span_origin=LitestarIntegration.origin):
# type: (ASGIApp, str) -> None
super().__init__(
app=app,
unsafe_context_data=False,
transaction_style="endpoint",
mechanism_type="asgi",
span_origin=span_origin,
)
def patch_app_init():
# type: () -> None
"""
Replaces the Litestar class's `__init__` function in order to inject `after_exception` handlers and set the
`SentryLitestarASGIMiddleware` as the outmost middleware in the stack.
See:
- https://docs.litestar.dev/2/usage/applications.html#after-exception
- https://docs.litestar.dev/2/usage/middleware/using-middleware.html
"""
old__init__ = Litestar.__init__
@ensure_integration_enabled(LitestarIntegration, old__init__)
def injection_wrapper(self, *args, **kwargs):
# type: (Litestar, *Any, **Any) -> None
kwargs["after_exception"] = [
exception_handler,
*(kwargs.get("after_exception") or []),
]
SentryLitestarASGIMiddleware.__call__ = SentryLitestarASGIMiddleware._run_asgi3 # type: ignore
middleware = kwargs.get("middleware") or []
kwargs["middleware"] = [SentryLitestarASGIMiddleware, *middleware]
old__init__(self, *args, **kwargs)
Litestar.__init__ = injection_wrapper
def patch_middlewares():
# type: () -> None
old_resolve_middleware_stack = BaseRouteHandler.resolve_middleware
@ensure_integration_enabled(LitestarIntegration, old_resolve_middleware_stack)
def resolve_middleware_wrapper(self):
# type: (BaseRouteHandler) -> list[Middleware]
return [
enable_span_for_middleware(middleware)
for middleware in old_resolve_middleware_stack(self)
]
BaseRouteHandler.resolve_middleware = resolve_middleware_wrapper
def enable_span_for_middleware(middleware):
# type: (Middleware) -> Middleware
if (
not hasattr(middleware, "__call__") # noqa: B004
or middleware is SentryLitestarASGIMiddleware
):
return middleware
if isinstance(middleware, DefineMiddleware):
old_call = middleware.middleware.__call__ # type: ASGIApp
else:
old_call = middleware.__call__
async def _create_span_call(self, scope, receive, send):
# type: (MiddlewareProtocol, LitestarScope, Receive, Send) -> None
if sentry_sdk.get_client().get_integration(LitestarIntegration) is None:
return await old_call(self, scope, receive, send)
middleware_name = self.__class__.__name__
with sentry_sdk.start_span(
op=OP.MIDDLEWARE_LITESTAR,
name=middleware_name,
origin=LitestarIntegration.origin,
) as middleware_span:
middleware_span.set_tag("litestar.middleware_name", middleware_name)
# Creating spans for the "receive" callback
async def _sentry_receive(*args, **kwargs):
# type: (*Any, **Any) -> Union[HTTPReceiveMessage, WebSocketReceiveMessage]
if sentry_sdk.get_client().get_integration(LitestarIntegration) is None:
return await receive(*args, **kwargs)
with sentry_sdk.start_span(
op=OP.MIDDLEWARE_LITESTAR_RECEIVE,
name=getattr(receive, "__qualname__", str(receive)),
origin=LitestarIntegration.origin,
) as span:
span.set_tag("litestar.middleware_name", middleware_name)
return await receive(*args, **kwargs)
receive_name = getattr(receive, "__name__", str(receive))
receive_patched = receive_name == "_sentry_receive"
new_receive = _sentry_receive if not receive_patched else receive
# Creating spans for the "send" callback
async def _sentry_send(message):
# type: (Message) -> None
if sentry_sdk.get_client().get_integration(LitestarIntegration) is None:
return await send(message)
with sentry_sdk.start_span(
op=OP.MIDDLEWARE_LITESTAR_SEND,
name=getattr(send, "__qualname__", str(send)),
origin=LitestarIntegration.origin,
) as span:
span.set_tag("litestar.middleware_name", middleware_name)
return await send(message)
send_name = getattr(send, "__name__", str(send))
send_patched = send_name == "_sentry_send"
new_send = _sentry_send if not send_patched else send
return await old_call(self, scope, new_receive, new_send)
not_yet_patched = old_call.__name__ not in ["_create_span_call"]
if not_yet_patched:
if isinstance(middleware, DefineMiddleware):
middleware.middleware.__call__ = _create_span_call
else:
middleware.__call__ = _create_span_call
return middleware
def patch_http_route_handle():
# type: () -> None
old_handle = HTTPRoute.handle
async def handle_wrapper(self, scope, receive, send):
# type: (HTTPRoute, HTTPScope, Receive, Send) -> None
if sentry_sdk.get_client().get_integration(LitestarIntegration) is None:
return await old_handle(self, scope, receive, send)
sentry_scope = sentry_sdk.get_isolation_scope()
request = scope["app"].request_class(
scope=scope, receive=receive, send=send
) # type: Request[Any, Any]
extracted_request_data = ConnectionDataExtractor(
parse_body=True, parse_query=True
)(request)
body = extracted_request_data.pop("body")
request_data = await body
def event_processor(event, _):
# type: (Event, Hint) -> Event
route_handler = scope.get("route_handler")
request_info = event.get("request", {})
request_info["content_length"] = len(scope.get("_body", b""))
if should_send_default_pii():
request_info["cookies"] = extracted_request_data["cookies"]
if request_data is not None:
request_info["data"] = request_data
func = None
if route_handler.name is not None:
tx_name = route_handler.name
# Accounts for use of type `Ref` in earlier versions of litestar without the need to reference it as a type
elif hasattr(route_handler.fn, "value"):
func = route_handler.fn.value
else:
func = route_handler.fn
if func is not None:
tx_name = transaction_from_function(func)
tx_info = {"source": SOURCE_FOR_STYLE["endpoint"]}
if not tx_name:
tx_name = _DEFAULT_TRANSACTION_NAME
tx_info = {"source": TRANSACTION_SOURCE_ROUTE}
event.update(
{
"request": request_info,
"transaction": tx_name,
"transaction_info": tx_info,
}
)
return event
sentry_scope._name = LitestarIntegration.identifier
sentry_scope.add_event_processor(event_processor)
return await old_handle(self, scope, receive, send)
HTTPRoute.handle = handle_wrapper
def retrieve_user_from_scope(scope):
# type: (LitestarScope) -> Optional[dict[str, Any]]
scope_user = scope.get("user")
if isinstance(scope_user, dict):
return scope_user
if hasattr(scope_user, "asdict"): # dataclasses
return scope_user.asdict()
return None
@ensure_integration_enabled(LitestarIntegration)
def exception_handler(exc, scope):
# type: (Exception, LitestarScope) -> None
user_info = None # type: Optional[dict[str, Any]]
if should_send_default_pii():
user_info = retrieve_user_from_scope(scope)
if user_info and isinstance(user_info, dict):
sentry_scope = sentry_sdk.get_isolation_scope()
sentry_scope.set_user(user_info)
if isinstance(exc, HTTPException):
integration = sentry_sdk.get_client().get_integration(LitestarIntegration)
if (
integration is not None
and exc.status_code not in integration.failed_request_status_codes
):
return
event, hint = event_from_exception(
exc,
client_options=sentry_sdk.get_client().options,
mechanism={"type": LitestarIntegration.identifier, "handled": False},
)
sentry_sdk.capture_event(event, hint=hint)
@@ -0,0 +1,294 @@
import logging
from datetime import datetime, timezone
from fnmatch import fnmatch
import sentry_sdk
from sentry_sdk.utils import (
to_string,
event_from_exception,
current_stacktrace,
capture_internal_exceptions,
)
from sentry_sdk.integrations import Integration
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import MutableMapping
from logging import LogRecord
from typing import Any
from typing import Dict
from typing import Optional
DEFAULT_LEVEL = logging.INFO
DEFAULT_EVENT_LEVEL = logging.ERROR
LOGGING_TO_EVENT_LEVEL = {
logging.NOTSET: "notset",
logging.DEBUG: "debug",
logging.INFO: "info",
logging.WARN: "warning", # WARN is same a WARNING
logging.WARNING: "warning",
logging.ERROR: "error",
logging.FATAL: "fatal",
logging.CRITICAL: "fatal", # CRITICAL is same as FATAL
}
# Capturing events from those loggers causes recursion errors. We cannot allow
# the user to unconditionally create events from those loggers under any
# circumstances.
#
# Note: Ignoring by logger name here is better than mucking with thread-locals.
# We do not necessarily know whether thread-locals work 100% correctly in the user's environment.
_IGNORED_LOGGERS = set(
["sentry_sdk.errors", "urllib3.connectionpool", "urllib3.connection"]
)
def ignore_logger(
name, # type: str
):
# type: (...) -> None
"""This disables recording (both in breadcrumbs and as events) calls to
a logger of a specific name. Among other uses, many of our integrations
use this to prevent their actions being recorded as breadcrumbs. Exposed
to users as a way to quiet spammy loggers.
:param name: The name of the logger to ignore (same string you would pass to ``logging.getLogger``).
"""
_IGNORED_LOGGERS.add(name)
class LoggingIntegration(Integration):
identifier = "logging"
def __init__(self, level=DEFAULT_LEVEL, event_level=DEFAULT_EVENT_LEVEL):
# type: (Optional[int], Optional[int]) -> None
self._handler = None
self._breadcrumb_handler = None
if level is not None:
self._breadcrumb_handler = BreadcrumbHandler(level=level)
if event_level is not None:
self._handler = EventHandler(level=event_level)
def _handle_record(self, record):
# type: (LogRecord) -> None
if self._handler is not None and record.levelno >= self._handler.level:
self._handler.handle(record)
if (
self._breadcrumb_handler is not None
and record.levelno >= self._breadcrumb_handler.level
):
self._breadcrumb_handler.handle(record)
@staticmethod
def setup_once():
# type: () -> None
old_callhandlers = logging.Logger.callHandlers
def sentry_patched_callhandlers(self, record):
# type: (Any, LogRecord) -> Any
# keeping a local reference because the
# global might be discarded on shutdown
ignored_loggers = _IGNORED_LOGGERS
try:
return old_callhandlers(self, record)
finally:
# This check is done twice, once also here before we even get
# the integration. Otherwise we have a high chance of getting
# into a recursion error when the integration is resolved
# (this also is slower).
if ignored_loggers is not None and record.name not in ignored_loggers:
integration = sentry_sdk.get_client().get_integration(
LoggingIntegration
)
if integration is not None:
integration._handle_record(record)
logging.Logger.callHandlers = sentry_patched_callhandlers # type: ignore
class _BaseHandler(logging.Handler):
COMMON_RECORD_ATTRS = frozenset(
(
"args",
"created",
"exc_info",
"exc_text",
"filename",
"funcName",
"levelname",
"levelno",
"linenno",
"lineno",
"message",
"module",
"msecs",
"msg",
"name",
"pathname",
"process",
"processName",
"relativeCreated",
"stack",
"tags",
"taskName",
"thread",
"threadName",
"stack_info",
)
)
def _can_record(self, record):
# type: (LogRecord) -> bool
"""Prevents ignored loggers from recording"""
for logger in _IGNORED_LOGGERS:
if fnmatch(record.name, logger):
return False
return True
def _logging_to_event_level(self, record):
# type: (LogRecord) -> str
return LOGGING_TO_EVENT_LEVEL.get(
record.levelno, record.levelname.lower() if record.levelname else ""
)
def _extra_from_record(self, record):
# type: (LogRecord) -> MutableMapping[str, object]
return {
k: v
for k, v in vars(record).items()
if k not in self.COMMON_RECORD_ATTRS
and (not isinstance(k, str) or not k.startswith("_"))
}
class EventHandler(_BaseHandler):
"""
A logging handler that emits Sentry events for each log record
Note that you do not have to use this class if the logging integration is enabled, which it is by default.
"""
def emit(self, record):
# type: (LogRecord) -> Any
with capture_internal_exceptions():
self.format(record)
return self._emit(record)
def _emit(self, record):
# type: (LogRecord) -> None
if not self._can_record(record):
return
client = sentry_sdk.get_client()
if not client.is_active():
return
client_options = client.options
# exc_info might be None or (None, None, None)
#
# exc_info may also be any falsy value due to Python stdlib being
# liberal with what it receives and Celery's billiard being "liberal"
# with what it sends. See
# https://github.com/getsentry/sentry-python/issues/904
if record.exc_info and record.exc_info[0] is not None:
event, hint = event_from_exception(
record.exc_info,
client_options=client_options,
mechanism={"type": "logging", "handled": True},
)
elif (record.exc_info and record.exc_info[0] is None) or record.stack_info:
event = {}
hint = {}
with capture_internal_exceptions():
event["threads"] = {
"values": [
{
"stacktrace": current_stacktrace(
include_local_variables=client_options[
"include_local_variables"
],
max_value_length=client_options["max_value_length"],
),
"crashed": False,
"current": True,
}
]
}
else:
event = {}
hint = {}
hint["log_record"] = record
level = self._logging_to_event_level(record)
if level in {"debug", "info", "warning", "error", "critical", "fatal"}:
event["level"] = level # type: ignore[typeddict-item]
event["logger"] = record.name
# Log records from `warnings` module as separate issues
record_caputured_from_warnings_module = (
record.name == "py.warnings" and record.msg == "%s"
)
if record_caputured_from_warnings_module:
# use the actual message and not "%s" as the message
# this prevents grouping all warnings under one "%s" issue
msg = record.args[0] # type: ignore
event["logentry"] = {
"message": msg,
"params": (),
}
else:
event["logentry"] = {
"message": to_string(record.msg),
"params": record.args,
}
event["extra"] = self._extra_from_record(record)
sentry_sdk.capture_event(event, hint=hint)
# Legacy name
SentryHandler = EventHandler
class BreadcrumbHandler(_BaseHandler):
"""
A logging handler that records breadcrumbs for each log record.
Note that you do not have to use this class if the logging integration is enabled, which it is by default.
"""
def emit(self, record):
# type: (LogRecord) -> Any
with capture_internal_exceptions():
self.format(record)
return self._emit(record)
def _emit(self, record):
# type: (LogRecord) -> None
if not self._can_record(record):
return
sentry_sdk.add_breadcrumb(
self._breadcrumb_from_record(record), hint={"log_record": record}
)
def _breadcrumb_from_record(self, record):
# type: (LogRecord) -> Dict[str, Any]
return {
"type": "log",
"level": self._logging_to_event_level(record),
"category": record.name,
"message": record.message,
"timestamp": datetime.fromtimestamp(record.created, timezone.utc),
"data": self._extra_from_record(record),
}

Some files were not shown because too many files have changed in this diff Show More