skywalking/plugins/sw_asyncpg.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from skywalking import Layer, Component, config
from skywalking.trace.context import get_context
from skywalking.trace.tags import TagDbType, TagDbInstance, TagDbStatement, TagDbSqlParameters

link_vector = ['https://github.com/MagicStack/asyncpg']
support_matrix = {
    'asyncpg': {
        '>=3.7': ['0.25.0'],
    }
}
note = """"""


def install():
    from asyncpg import Connection
    from asyncpg.protocol import Protocol

    def _sw_init(self, *args, **kwargs):
        _init(self, *args, **kwargs)
        self._protocol._addr = f'{self._addr[0]}:{self._addr[1]}'
        self._protocol._database = self._params.database

    async def __bind(proto, query, params, future, is_many=False):
        peer = getattr(proto, '_addr', '<unavailable>')  # just in case

        with get_context().new_exit_span(op='PostgreSQL/AsyncPG/bind', peer=peer,
                                         component=Component.AsyncPG) as span:
            span.layer = Layer.Database

            span.tag(TagDbType('PostgreSQL'))
            span.tag(TagDbInstance(getattr(proto, '_database', '<unavailable>')))
            span.tag(TagDbStatement(query))

            if config.plugin_sql_parameters_max_length and params is not None:
                if not is_many:
                    text = ','.join(str(v) for v in params)

                    if len(text) > config.plugin_sql_parameters_max_length:
                        text = f'{text[:config.plugin_sql_parameters_max_length]}...'

                    span.tag(TagDbSqlParameters(f'[{text}]'))

                else:
                    max_len = config.plugin_sql_parameters_max_length
                    total_len = 0
                    text_list = []

                    for _params in params:
                        text = f"[{','.join(str(v) for v in _params)}]"
                        total_len += len(text)

                        if total_len > max_len:
                            text_list.append(f'{text[:max_len - total_len]}...')

                            break

                        text_list.append(text)

                    span.tag(TagDbSqlParameters(f"[{','.join(text_list)}]"))

            return await future

    async def _sw_bind(proto, stmt, params, *args, **kwargs):
        return await __bind(proto, stmt.query, params, _bind(proto, stmt, params, *args, **kwargs))

    async def _sw_bind_execute(proto, stmt, params, *args, **kwargs):
        return await __bind(proto, stmt.query, params, _bind_execute(proto, stmt, params, *args, **kwargs))

    async def _sw_bind_execute_many(proto, stmt, params, *args, **kwargs):
        return await __bind(proto, stmt.query, params, _bind_execute_many(proto, stmt, params, *args, **kwargs), True)

    async def _sw_query(proto, query, *args, **kwargs):
        return await __bind(proto, query, (), _query(proto, query, *args, **kwargs))

    # async def _sw_execute(proto, stmt, *args, **kwargs):  # these may be useful in the future, left here for documentation purposes
    # async def _sw_prepare(*args, **kwargs):

    _init = Connection.__init__
    _bind = Protocol.bind
    _bind_execute = Protocol.bind_execute
    _bind_execute_many = Protocol.bind_execute_many
    _query = Protocol.query
    # _execute = Protocol.execute
    # _prepare = Protocol.prepare

    Connection.__init__ = _sw_init
    Protocol.bind = _sw_bind
    Protocol.bind_execute = _sw_bind_execute
    Protocol.bind_execute_many = _sw_bind_execute_many
    Protocol.query = _sw_query
    # Protocol.execute = _sw_execute
    # Protocol.prepare = _sw_prepare


# Example code for someone who might want to make tests:
#
# async def conncetAsycPg():
#     power=3
#     con = await asyncpg.connect(user='user',password='cv9493a32',database="db" , host='localhost')
#     await con.fetchval('SELECT 2 ^ $1', power)
#     types = await con.fetch('SELECT * FROM pg_type')
#     async with con.transaction():
#         async for record in con.cursor('SELECT generate_series(0, 100)'):
#             print(record)
#     await con.close()