2025-07-01
This commit is contained in:
@@ -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)
|
||||
+1343
File diff suppressed because it is too large
Load Diff
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;$" => 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}")
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -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("You’re out of downloads")
|
||||
text = _t("You’ll 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
|
||||
+108
@@ -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
|
||||
+271
@@ -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)
|
||||
+528
@@ -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
|
||||
+157
@@ -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)
|
||||
+258
@@ -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
|
||||
+747
@@ -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
|
||||
+191
@@ -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
|
||||
+187
@@ -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
|
||||
+91
@@ -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
|
||||
+188
@@ -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,
|
||||
}
|
||||
+159
@@ -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"
|
||||
"×tamp={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
|
||||
+107
@@ -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()
|
||||
+151
@@ -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)
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
from .server import ServerInterceptor
|
||||
from .client import ClientInterceptor
|
||||
|
||||
__all__ = [
|
||||
"ClientInterceptor",
|
||||
"ServerInterceptor",
|
||||
]
|
||||
+94
@@ -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
|
||||
+100
@@ -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
|
||||
+175
@@ -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
Reference in New Issue
Block a user