# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "opentelemetry-api",
#     "opentelemetry-instrumentation-sqlalchemy",
#     "opentelemetry-sdk",
#     "sqlalchemy",
# ]
# ///
from opentelemetry import trace
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.instrumentation.utils import suppress_instrumentation
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from sqlalchemy import create_engine

exporter = ConsoleSpanExporter()
provider = TracerProvider()
provider.add_span_processor(BatchSpanProcessor(exporter))
trace.set_tracer_provider(provider)

engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(engine=engine)

with engine.connect():  # produces spans
    pass

with suppress_instrumentation(), engine.connect():  # should not produce spans but does
    pass

def _wrap_connect(tracer):
    # pylint: disable=unused-argument
    def _wrap_connect_internal(func, module, args, kwargs):
        if not is_instrumentation_enabled(): # this is new
            return func(*args, **kwargs)

        with tracer.start_as_current_span(
            "connect", kind=trace.SpanKind.CLIENT
        ) as span:
            if span.is_recording():
                attrs, _ = _get_attributes_from_url(module.url)
                span.set_attributes(attrs)
                span.set_attribute(
                    SpanAttributes.DB_SYSTEM, _normalize_vendor(module.name)
                )
            return func(*args, **kwargs)

    return _wrap_connect_internal