Skip to main content

Python API

The DAGZ pytest plugin exposes a small Python API for two integrations:

  • DB rerouting for parallel database tests (dagz.integ.psycopg, dagz.integ.pymysql, etc.).
  • OpenTelemetry integration that turns OTel spans into DAGZ sub-spans (dagz.integ.otel).

For the conceptual walkthrough and setup patterns of DB rerouting, see Parallelizing DB Tests. The page below documents the public API surface.

DB Rerouting

Each integration exposes a singleton config object: PG_CONFIG, MYSQL_CONFIG, CASSANDRA_CONFIG, etc. All share the same surface, described here in terms of a generic CONFIG.

CONFIG.configure(...)

Activates rerouting for this database family. Call once at conftest module load.

ParameterRequiredPurpose
rewrite_db_name(db_name, session) -> stryesCompute the per-worker name. Use default_rewrite_db_name for the _jw{N} suffix.
should_reroute(db_name) -> boolyesDecide whether to reroute a given DB. Use default_should_reroute to skip system DBs.
prepare(db_name) -> NoneyesPer-test preparation. Pass None for no preparation.
worker_init(worker_num) -> Noneno, default NoneOne-time per-worker setup.
single_worker_disable: boolno, default FalseUnhook entirely when running with one worker.
main_worker_bypass: boolno, default FalseSkip rerouting on the main worker. Used together with the dagz_main_worker mark.

The first three are required positional contract: pass default_rewrite_db_name / default_should_reroute for the standard scheme, and None for prepare if you do not need per-test preparation.

Manual init helpers

Use these to opt out of the global prepare callback for a single test, typically inside a fixture that pre-loads its own data.

CONFIG.get_manual_init_db_name(name) -> str Resolves name to its per-worker form through the configured rewrite_db_name policy (e.g. myappmyapp_jw03), and marks the resolved name as already initialized for the current test. The auto-prepare callback then skips it on the first connection of this test. The marker is per-test: cleared at the start of every test via a test_prerun hook.

CONFIG.get_manual_init_db_url(url) -> str Same as above for URL-based connection strings. Parses the URL, rewrites the path component, and reassembles.

CONFIG.bypass()

Context manager that disables rerouting on the current thread / async task. Backed by contextvars, so it is thread-safe and async-task-local.

with PG_CONFIG.bypass():
conn = psycopg.connect(dbname="postgres") # admin DB, not rerouted

Driver-level helpers

For wiring a client library that DAGZ does not ship a hook for. Call from your own override of the library's connect function.

Client call shapeHelper
connect(database="myapp", ...)CONFIG.maybe_reroute(name)
connect(dbname="myapp", ...)CONFIG.maybe_reroute(name)
connect("postgresql://user:pw@host/myapp")CONFIG.maybe_reroute_uri(uri)

maybe_reroute(name) -> str returns the rerouted name if the config wants to reroute it, the original name otherwise. maybe_reroute_uri(uri) -> str does the same for a URI string.

For libraries with a custom shape (multi-statement keyspace strings, byte-encoded auth blobs), call CONFIG.should_reroute(name) -> bool and CONFIG.reroute_db(name) -> str directly.

CONFIG.truncate(db_name, **connect_kwargs)

Connects to the named DB and empties every user table. Suitable as the body of a prepare callback for the common case of clearing all user tables between tests.

def _prepare(db_name):
PG_CONFIG.truncate(db_name, host="localhost", port=5432, user="postgres", password="postgres")

OpenTelemetry

dagz.integ.otel registers a span processor that mirrors OpenTelemetry spans into DAGZ as sub-spans. Each OTel span becomes a DAGZ sub-span tagged with an otel_type attribute derived from the OTel instrumentation scope and operation, and OK / ERROR status maps to DAGZ's SpanResult.

dagz.otel.configure(...)

Call once from conftest.py. No-op if the DAGZ pytest plugin is not enabled.

ParameterRequiredDefaultPurpose
dd: boolyes(no default)Must be False. True is not yet supported and raises if ddtrace is importable.
start_provider: boolnoTrueCreate and set a TracerProvider if one is not already configured.
auto_instrument: boolnoTrueRun OpenTelemetry's auto_instrumentation.initialize() to install all available instrumentations.
instrument_logging: boolnoFalseIf False, adds logging to OTEL_PYTHON_DISABLED_INSTRUMENTATIONS so OTel does not replace the standard logging module.
# conftest.py
import dagz.otel

dagz.otel.configure(dd=False)

After configure() runs, any TracerProvider set later (e.g. by application code) is also instrumented: trace.set_tracer_provider is patched to attach the DAGZ span processor to subsequent providers.