import opentelemetry.instrumentation.dbapi
from opentelemetry.instrumentation.dbapi import DatabaseApiIntegration

original_DatabaseApiIntegration = None


class _ApproziumDatabaseApiIntegration(DatabaseApiIntegration):
    def get_connection_attributes(self, connection):
        super().get_connection_attributes(connection)
        # check if this is an Approzium connection
        if hasattr(connection, "authenticator"):
            self.add_approzium_span_attributes(connection)

    def add_approzium_span_attributes(self, connection):
        for key, value in connection.authenticator.attribution_info.items():
            self.span_attributes["approzium." + key] = value


def instrument():
    """Instruments Opentelemetry ``opentelemetry.instrumentation.dbapi`` to automatically
    include attribution tags in traces generated by Approzium database connections.
    Call this method before calling the ``instrument`` method provided by Opentelemetry
    database extensions.

    Refer to :py:obj:`approzium.AuthClient.attribution_info` for generated tags.

    Example (using Opentelemetry Psycopg2 integration):

    .. code-block:: python

        >>> import approzium.opentelemetry
        >>> from approzium.psycopg2 import connect
        >>> from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
        >>> approzium.opentelemetry.instrument()
        >>> Psycopg2Instrumentor().instrument()
        >>> cnx = connect("host=dbmd5 dbname=db user=bob")
        >>> cursor.execute("SELECT 1")
    """
    global original_DatabaseApiIntegration
    original_DatabaseApiIntegration = DatabaseApiIntegration
    opentelemetry.instrumentation.dbapi.DatabaseApiIntegration = (
        _ApproziumDatabaseApiIntegration
    )


def uninstrument():
    """Removes Approzium instrumentation from Opentelemetry
    ``opentelemetry.instrumentation.dbapi``

    :raises: Exception, if ``opentelemetry.instrumentation.dbapi`` was not instrumented
    by Approzium
    """
    if original_DatabaseApiIntegration is None:
        raise Exception("Opentelemetry is not currently instrumented by Approzium")
    opentelemetry.instrumentation.dbapi.DatabaseApiIntegration = (
        original_DatabaseApiIntegration
    )
