diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..d544de7 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,15 @@ +root = true + +[*] +end_of_line = lf +charset = utf-8 +insert_final_newline = true +trim_trailing_whitespace = true + +[*.{md}] +indent_size = 2 +indent_style = space + +[*.{cfg,py,sh}] +indent_size = 4 +indent_style = space diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..096d742 --- /dev/null +++ b/.gitignore @@ -0,0 +1,269 @@ + +# Created by https://www.toptal.com/developers/gitignore/api/pycharm+all,python,vim +# Edit at https://www.toptal.com/developers/gitignore?templates=pycharm+all,python,vim + +### PyCharm+all ### +# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider +# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 + +# User-specific stuff +.idea/**/workspace.xml +.idea/**/tasks.xml +.idea/**/usage.statistics.xml +.idea/**/dictionaries +.idea/**/shelf + +# AWS User-specific +.idea/**/aws.xml + +# Generated files +.idea/**/contentModel.xml + +# Sensitive or high-churn files +.idea/**/dataSources/ +.idea/**/dataSources.ids +.idea/**/dataSources.local.xml +.idea/**/sqlDataSources.xml +.idea/**/dynamic.xml +.idea/**/uiDesigner.xml +.idea/**/dbnavigator.xml + +# Gradle +.idea/**/gradle.xml +.idea/**/libraries + +# Gradle and Maven with auto-import +# When using Gradle or Maven with auto-import, you should exclude module files, +# since they will be recreated, and may cause churn. Uncomment if using +# auto-import. +# .idea/artifacts +# .idea/compiler.xml +# .idea/jarRepositories.xml +# .idea/modules.xml +# .idea/*.iml +# .idea/modules +# *.iml +# *.ipr + +# CMake +cmake-build-*/ + +# Mongo Explorer plugin +.idea/**/mongoSettings.xml + +# File-based project format +*.iws + +# IntelliJ +out/ + +# mpeltonen/sbt-idea plugin +.idea_modules/ + +# JIRA plugin +atlassian-ide-plugin.xml + +# Cursive Clojure plugin +.idea/replstate.xml + +# SonarLint plugin +.idea/sonarlint/ + +# Crashlytics plugin (for Android Studio and IntelliJ) +com_crashlytics_export_strings.xml +crashlytics.properties +crashlytics-build.properties +fabric.properties + +# Editor-based Rest Client +.idea/httpRequests + +# Android studio 3.1+ serialized cache file +.idea/caches/build_file_checksums.ser + +### PyCharm+all Patch ### +# Ignore everything but code style settings and run configurations +# that are supposed to be shared within teams. + +.idea/* + +!.idea/codeStyles +!.idea/runConfigurations + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +### Vim ### +# Swap +[._]*.s[a-v][a-z] +!*.svg # comment out if you don't need vector files +[._]*.sw[a-p] +[._]s[a-rt-v][a-z] +[._]ss[a-gi-z] +[._]sw[a-p] + +# Session +Session.vim +Sessionx.vim + +# Temporary +.netrwhist +*~ +# Auto-generated tag files +tags +# Persistent undo +[._]*.un~ + +# End of https://www.toptal.com/developers/gitignore/api/pycharm+all,python,vim + diff --git a/celery_aws_xray_sdk_extension/__init__.py b/celery_aws_xray_sdk_extension/__init__.py new file mode 100644 index 0000000..b794fd4 --- /dev/null +++ b/celery_aws_xray_sdk_extension/__init__.py @@ -0,0 +1 @@ +__version__ = '0.1.0' diff --git a/celery_aws_xray_sdk_extension/handlers.py b/celery_aws_xray_sdk_extension/handlers.py new file mode 100644 index 0000000..eaf57c6 --- /dev/null +++ b/celery_aws_xray_sdk_extension/handlers.py @@ -0,0 +1,49 @@ +from aws_xray_sdk.core import xray_recorder +from aws_xray_sdk.core.utils import stacktrace +from aws_xray_sdk.ext.util import construct_xray_header, inject_trace_header + +__all__ = ( + 'xray_after_task_publish', + 'xray_before_task_publish', + 'xray_task_failure', + 'xray_task_postrun', + 'xray_task_prerun', +) + +CELERY_NAMESPACE = 'celery' + + +def xray_before_task_publish(sender=None, headers=None, **kwargs): + headers = headers if headers else {} + task_id = headers.get('id') + + subsegment = xray_recorder.begin_subsegment(name=sender, namespace='remote') + if not subsegment: + return + + subsegment.put_metadata('task_id', task_id, namespace=CELERY_NAMESPACE) + inject_trace_header(headers, subsegment) + + +def xray_after_task_publish(**kwargs): + xray_recorder.end_subsegment() + + +def xray_task_prerun(task_id=None, task=None, **kwargs): + xray_header = construct_xray_header(task.request) + segment = xray_recorder.begin_segment(name=task.name, traceid=xray_header.root, parent_id=xray_header.parent) + segment.save_origin_trace_header(xray_header) + segment.put_annotation('routing_key', task.request.properties["delivery_info"]["routing_key"]) + segment.put_annotation('task_name', task.name) + segment.put_metadata('task_id', task_id, namespace=CELERY_NAMESPACE) + + +def xray_task_postrun(**kwargs): + xray_recorder.end_segment() + + +def xray_task_failure(exception=None, **kwargs): + segment = xray_recorder.current_segment() + if exception: + stack = stacktrace.get_stacktrace(limit=xray_recorder._max_trace_back) + segment.add_exception(exception, stack) diff --git a/example/readme.md b/example/readme.md new file mode 100644 index 0000000..82ed10c --- /dev/null +++ b/example/readme.md @@ -0,0 +1,37 @@ +# Example for celery-aws-xray-sdk-extension + +## Installation + +[Docker](https://www.docker.com/) is required for this example setup. + +1. Create an empty [virtualenv](https://virtualenv.pypa.io/en/latest/) in this directory and activate it. +2. Run `pip install -e .`in this directory. +3. Set environment variables to your testing AWS account (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` and `AWS_SESSION_TOKEN`). +4. Run `./run_xray_daemon.sh`. This starts in background a Docker container with AWS X-Ray daemon using latest version on port 2000. +5. Run `./run.sh`. This starts the Celery worker with AWS X-Ray SDK patch for `httplib`. + +## Running demo Celery tasks + +This demo project contains 3 Celery tasks: `add`, `power_2` and `complex`. + +### Task `add` + +`celery -A tasks call -a '[2, 2]' tasks.add` + +Simply adds together two numbers. + +### Task `power_2` + +`celery -A tasks call -a '[4]' tasks.power_2` + +Returns second power of a number. + +### Task `complex` + +`celery -A tasks call -a '[2, 2]' tasks.complex` + +Sends a HTTP request to [example.com](https://example.com/) and returns a second power of sum of two numbers. + +`celery -A tasks call -a '[2, 2]' -k '{"fail": true}' tasks.complex` + +You can optionally make it fail by providing `fail` kwarg set to `True`. diff --git a/example/requirements.txt b/example/requirements.txt new file mode 100644 index 0000000..2c28e6e --- /dev/null +++ b/example/requirements.txt @@ -0,0 +1,4 @@ +-e .. +aws-xray-sdk>=2.9.0,<3 +celery[redis]>=5,<6 +requests<3 diff --git a/example/run.sh b/example/run.sh new file mode 100755 index 0000000..62dc66c --- /dev/null +++ b/example/run.sh @@ -0,0 +1,4 @@ +#!/bin/sh +docker run --rm -d -p 6379:6379 --name celery-aws-xray-sdk-extension redis:alpine +celery -A tasks worker -n celery-aws-xray-sdk-extension-example -E -Q xray-demo -l INFO +docker stop celery-aws-xray-sdk-extension diff --git a/example/run_xray_daemon.sh b/example/run_xray_daemon.sh new file mode 100755 index 0000000..fe2e5d8 --- /dev/null +++ b/example/run_xray_daemon.sh @@ -0,0 +1,2 @@ +#!/bin/sh +docker run --name aws-xray-daemon --rm -d -p 2000:2000 -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} --network host public.ecr.aws/xray/aws-xray-daemon:latest -o -n eu-central-1 diff --git a/example/tasks.py b/example/tasks.py new file mode 100644 index 0000000..1a0d106 --- /dev/null +++ b/example/tasks.py @@ -0,0 +1,38 @@ +import requests +from aws_xray_sdk.core.patcher import patch +from celery import Celery, signals + +from celery_aws_xray_sdk_extension.handlers import (xray_after_task_publish, + xray_before_task_publish, + xray_task_failure, + xray_task_postrun, + xray_task_prerun) + +patch(('httplib',)) + +app = Celery(backend='redis://localhost:6379/1', broker='redis://localhost:6379/0') +app.conf.task_default_queue = 'xray-demo' + +signals.after_task_publish.connect(xray_after_task_publish) +signals.before_task_publish.connect(xray_before_task_publish) +signals.task_failure.connect(xray_task_failure) +signals.task_postrun.connect(xray_task_postrun) +signals.task_prerun.connect(xray_task_prerun) + + +@app.task +def add(x, y): + return x + y + + +@app.task +def power_2(self, x): + return x ** 2 + + +@app.task +def complex(x, y, fail=False): + requests.get('https://example.com') + add.apply_async((2, 2), link=power_2.s()) + if fail: + raise Exception('A really serious bug') diff --git a/license b/license new file mode 100644 index 0000000..0764868 --- /dev/null +++ b/license @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Radim Sückr + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..105fc92 --- /dev/null +++ b/readme.md @@ -0,0 +1,29 @@ +# celery-aws-xray-sdk-extension + +[![PyPI version](https://badge.fury.io/py/celery-aws-xray-sdk-extension.svg)](https://pypi.org/project/celery-aws-xray-sdk-extension/) + +Celery signal handlers that integrate Celery task lifecycle with AWS X-Ray tracing. + +There's a tiny example in the directory `example`. + +## Installation + +You can install it easily with `pip`: `pip install celery-aws-xray-sdk-extension`. For latest version visit [PyPI](https://pypi.org/project/celery-aws-xray-sdk-extension/). + +## Setup + +This guide doesn't cover setting up AWS X-Ray SDK for Python or AWS X-Ray daemon. It's expected that you've already got some experience with it. If you don't have any experience with AWS X-Ray, please visit [Amazon documentation](https://docs.aws.amazon.com/xray/latest/devguide/aws-xray.html). + +1. You have to have [Celery signals](https://docs.celeryq.dev/en/stable/userguide/signals.html) enabled. +2. Connect Celery signals to signal handlers from `celery_aws_xray_sdk_extension` module in your Celery setup. Example code is below this list. +3. You're good to go! + +### Connecting handlers to Celery signals + +```python +signals.after_task_publish.connect(xray_after_task_publish) +signals.before_task_publish.connect(xray_before_task_publish) +signals.task_failure.connect(xray_task_failure) +signals.task_postrun.connect(xray_task_postrun) +signals.task_prerun.connect(xray_task_prerun) +``` diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..e00fa8b --- /dev/null +++ b/setup.cfg @@ -0,0 +1,5 @@ +[flake8] +max-line-length = 120 + +[metadata] +version = attr: celery_aws_xray_sdk_extension.__version__ diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..1794033 --- /dev/null +++ b/setup.py @@ -0,0 +1,33 @@ +from pathlib import Path + +from setuptools import find_packages, setup + + +def read_root_file(filename): + root = Path('.') + with open(root / filename, 'r') as f: + return f.read().strip() + + +setup( + name='celery-aws-xray-sdk-extension', + # version=read_root_file('version'), + author='Radim Sückr', + author_email='kontakt@radimsuckr.cz', + description='Extension for AWS X-Ray SDK which enables tracing of Celery tasks', + long_description=read_root_file('readme.md'), + long_description_content_type='text/markdown', + license='MIT', + url='https://github.com/druids/celery-aws-xray-sdk-extension', + packages=find_packages(include=['celery_aws_xray_sdk_extension', 'tests']), + install_requires=[ + 'aws-xray-sdk>=2.9,<3', + 'celery>=5,<6', + ], + python_requires='>=3.6', + classifiers=[ + 'Development Status :: 3 - Alpha', + 'Framework :: Celery', + 'License :: OSI Approved :: MIT License', + ], +)