elasticapm/instrumentation/packages/psycopg2.py (87 lines of code) (raw):
# BSD 3-Clause License
#
# Copyright (c) 2019, Elasticsearch BV
# All 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 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, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from __future__ import absolute_import
from typing import Optional, Union
from elasticapm.instrumentation.packages.dbapi2 import (
ConnectionProxy,
CursorProxy,
DbApi2Instrumentation,
extract_signature,
)
from elasticapm.traces import capture_span
from elasticapm.utils import default_ports
class PGCursorProxy(CursorProxy):
provider_name = "postgresql"
def _bake_sql(self, sql):
from psycopg2 import extensions as psycopg2_extensions
# if this is a Composable object, use its `as_string` method
# see http://initd.org/psycopg/docs/sql.html
if hasattr(sql, "as_string"):
sql = sql.as_string(self.__wrapped__)
# if the sql string is already a byte string, we need to decode it using the connection encoding
if isinstance(sql, bytes):
sql = sql.decode(psycopg2_extensions.encodings[self.__wrapped__.connection.encoding])
return sql
def extract_signature(self, sql):
return extract_signature(sql)
def __enter__(self):
return PGCursorProxy(self.__wrapped__.__enter__(), destination_info=self._self_destination_info)
@property
def _self_database(self):
return self.connection.info.dbname or ""
class PGConnectionProxy(ConnectionProxy):
cursor_proxy = PGCursorProxy
def __enter__(self):
return PGConnectionProxy(self.__wrapped__.__enter__(), destination_info=self._self_destination_info)
class Psycopg2Instrumentation(DbApi2Instrumentation):
name = "psycopg2"
instrument_list = [("psycopg2", "connect")]
def call(self, module, method, wrapped, instance, args, kwargs):
signature = "psycopg2.connect"
host, port = get_destination_info(kwargs.get("host"), kwargs.get("port"))
database = kwargs.get("database")
signature = f"{signature} {host}:{port}"
destination_info = {
"address": host,
"port": port,
}
with capture_span(
signature,
span_type="db",
span_subtype="postgresql",
span_action="connect",
leaf=True,
extra={"destination": destination_info, "db": {"type": "sql", "instance": database}},
):
return PGConnectionProxy(wrapped(*args, **kwargs), destination_info=destination_info)
class Psycopg2ExtensionsInstrumentation(DbApi2Instrumentation):
"""
Some extensions do a type check on the Connection/Cursor in C-code, which our
proxy fails. For these extensions, we need to ensure that the unwrapped
Connection/Cursor is passed.
"""
name = "psycopg2"
instrument_list = [
("psycopg2.extensions", "register_type"),
# specifically instrument `register_json` as it bypasses `register_type`
("psycopg2._json", "register_json"),
("psycopg2.extensions", "quote_ident"),
("psycopg2.extensions", "encrypt_password"),
]
def call(self, module, method, wrapped, instance, args, kwargs):
if "conn_or_curs" in kwargs and hasattr(kwargs["conn_or_curs"], "__wrapped__"):
kwargs["conn_or_curs"] = kwargs["conn_or_curs"].__wrapped__
# register_type takes the connection as second argument
elif len(args) == 2 and hasattr(args[1], "__wrapped__"):
args = (args[0], args[1].__wrapped__)
# register_json takes the connection as first argument, and can have
# several more arguments
elif method == "register_json":
if args and hasattr(args[0], "__wrapped__"):
args = (args[0].__wrapped__,) + args[1:]
elif method == "encrypt_password":
# connection/cursor is either 3rd argument, or "scope" keyword argument
if len(args) >= 3 and hasattr(args[2], "__wrapped__"):
args = args[:2] + (args[2].__wrapped__,) + args[3:]
elif "scope" in kwargs and hasattr(kwargs["scope"], "__wrapped__"):
kwargs["scope"] = kwargs["scope"].__wrapped__
return wrapped(*args, **kwargs)
def get_destination_info(host: Optional[str], port: Union[None, str, int]) -> tuple:
if host:
if "," in host: # multiple hosts defined, take first
host = host.split(",")[0]
else:
host = "localhost"
if port:
port = str(port)
if "," in port: # multiple ports defined, take first
port = port.split(",")[0]
port = int(port)
else:
port = default_ports.get("postgresql")
return host, port