예제 #1
0
def instrument_connection(
    name: str,
    connection,
    database_component: str,
    database_type: str = "",
    connection_attributes: typing.Dict = None,
    version: str = "",
    tracer_provider: typing.Optional[TracerProvider] = None,
):
    """Enable instrumentation in a database connection.

    Args:
        name: Name of opentelemetry extension for aiopg.
        connection: The connection to instrument.
        database_component: Database driver name or database name "postgreSQL".
        database_type: The Database type. For any SQL database, "sql".
        connection_attributes: Attribute names for database, port, host and
            user in a connection object.
        version: Version of opentelemetry extension for aiopg.
        tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
            use. If ommited the current configured one is used.

    Returns:
        An instrumented connection.
    """
    db_integration = AiopgIntegration(
        name,
        database_component,
        database_type,
        connection_attributes=connection_attributes,
        version=version,
        tracer_provider=tracer_provider,
    )
    db_integration.get_connection_attributes(connection)
    return get_traced_connection_proxy(connection, db_integration)
예제 #2
0
 def test_span_not_recording(self):
     connection_props = {
         "database": "testdatabase",
         "server_host": "testhost",
         "server_port": 123,
         "user": "******",
     }
     connection_attributes = {
         "database": "database",
         "port": "server_port",
         "host": "server_host",
         "user": "******",
     }
     mock_tracer = mock.Mock()
     mock_span = mock.Mock()
     mock_span.is_recording.return_value = False
     mock_tracer.start_span.return_value = mock_span
     mock_tracer.use_span.return_value.__enter__ = mock_span
     mock_tracer.use_span.return_value.__exit__ = True
     db_integration = AiopgIntegration(
         mock_tracer, "testcomponent", "testtype", connection_attributes
     )
     mock_connection = async_call(
         db_integration.wrapped_connection(
             mock_connect, {}, connection_props
         )
     )
     cursor = async_call(mock_connection.cursor())
     async_call(cursor.execute("Test query", ("param1Value", False)))
     self.assertFalse(mock_span.is_recording())
     self.assertTrue(mock_span.is_recording.called)
     self.assertFalse(mock_span.set_attribute.called)
     self.assertFalse(mock_span.set_status.called)
예제 #3
0
    def test_executemany(self):
        db_integration = AiopgIntegration(self.tracer, "testcomponent")
        mock_connection = async_call(
            db_integration.wrapped_connection(mock_connect, {}, {}))
        cursor = async_call(mock_connection.cursor())
        async_call(cursor.executemany("Test query"))
        spans_list = self.memory_exporter.get_finished_spans()

        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]
        self.assertEqual(span.attributes["db.statement"], "Test query")
    def test_span_succeeded(self):
        connection_props = {
            "database": "testdatabase",
            "server_host": "testhost",
            "server_port": 123,
            "user": "******",
        }
        connection_attributes = {
            "database": "database",
            "port": "server_port",
            "host": "server_host",
            "user": "******",
        }
        db_integration = AiopgIntegration(
            self.tracer,
            "testcomponent",
            connection_attributes,
            capture_parameters=True,
        )
        mock_connection = async_call(
            db_integration.wrapped_connection(
                mock_connect, {}, connection_props
            )
        )
        cursor = async_call(mock_connection.cursor())
        async_call(cursor.execute("Test query", ("param1Value", False)))
        spans_list = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]
        self.assertEqual(span.name, "Test")
        self.assertIs(span.kind, trace_api.SpanKind.CLIENT)

        self.assertEqual(
            span.attributes[SpanAttributes.DB_SYSTEM], "testcomponent"
        )
        self.assertEqual(
            span.attributes[SpanAttributes.DB_NAME], "testdatabase"
        )
        self.assertEqual(
            span.attributes[SpanAttributes.DB_STATEMENT], "Test query"
        )
        self.assertEqual(
            span.attributes["db.statement.parameters"],
            "('param1Value', False)",
        )
        self.assertEqual(span.attributes[SpanAttributes.DB_USER], "testuser")
        self.assertEqual(
            span.attributes[SpanAttributes.NET_PEER_NAME], "testhost"
        )
        self.assertEqual(span.attributes[SpanAttributes.NET_PEER_PORT], 123)
        self.assertIs(span.status.status_code, trace_api.StatusCode.UNSET)
    def test_span_failed(self):
        db_integration = AiopgIntegration(self.tracer, "testcomponent")
        mock_connection = async_call(
            db_integration.wrapped_connection(mock_connect, {}, {}))
        cursor = async_call(mock_connection.cursor())
        with self.assertRaises(Exception):
            async_call(cursor.execute("Test query", throw_exception=True))

        spans_list = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]
        self.assertEqual(span.attributes["db.statement"], "Test query")
        self.assertIs(span.status.status_code, trace_api.StatusCode.ERROR)
        self.assertEqual(span.status.description, "Exception: Test Exception")
    def test_callproc(self):
        db_integration = AiopgIntegration(self.tracer, "testcomponent")
        mock_connection = async_call(
            db_integration.wrapped_connection(mock_connect, {}, {}))
        cursor = async_call(mock_connection.cursor())
        async_call(cursor.callproc("Test stored procedure"))
        spans_list = self.memory_exporter.get_finished_spans()

        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]
        self.assertEqual(
            span.attributes[SpanAttributes.DB_STATEMENT],
            "Test stored procedure",
        )
 def wrap_connect_(
     wrapped: typing.Callable[..., typing.Any],
     instance: typing.Any,
     args: typing.Tuple[typing.Any, typing.Any],
     kwargs: typing.Dict[typing.Any, typing.Any],
 ):
     db_integration = AiopgIntegration(
         name,
         database_system,
         connection_attributes=connection_attributes,
         version=version,
         tracer_provider=tracer_provider,
     )
     return _ContextManager(  # pylint: disable=no-value-for-parameter
         db_integration.wrapped_connection(wrapped, args, kwargs))
 def wrap_create_pool_(
     wrapped: typing.Callable[..., typing.Any],
     instance: typing.Any,
     args: typing.Tuple[typing.Any, typing.Any],
     kwargs: typing.Dict[typing.Any, typing.Any],
 ):
     db_integration = AiopgIntegration(
         name,
         database_system,
         connection_attributes=connection_attributes,
         version=version,
         tracer_provider=tracer_provider,
     )
     return _PoolContextManager(
         db_integration.wrapped_pool(wrapped, args, kwargs))
 def wrap_connect_(
     wrapped: typing.Callable[..., typing.Any],
     instance: typing.Any,
     args: typing.Tuple[typing.Any, typing.Any],
     kwargs: typing.Dict[typing.Any, typing.Any],
 ):
     db_integration = AiopgIntegration(
         name,
         database_component,
         database_type=database_type,
         connection_attributes=connection_attributes,
         version=version,
         tracer_provider=tracer_provider,
     )
     return _ContextManager(
         db_integration.wrapped_connection(wrapped, args, kwargs))
    def test_span_succeeded(self):
        connection_props = {
            "database": "testdatabase",
            "server_host": "testhost",
            "server_port": 123,
            "user": "******",
        }
        connection_attributes = {
            "database": "database",
            "port": "server_port",
            "host": "server_host",
            "user": "******",
        }
        db_integration = AiopgIntegration(
            self.tracer, "testcomponent", "testtype", connection_attributes
        )
        mock_connection = async_call(
            db_integration.wrapped_connection(
                mock_connect, {}, connection_props
            )
        )
        cursor = async_call(mock_connection.cursor())
        async_call(cursor.execute("Test query", ("param1Value", False)))
        spans_list = self.memory_exporter.get_finished_spans()
        self.assertEqual(len(spans_list), 1)
        span = spans_list[0]
        self.assertEqual(span.name, "testcomponent.testdatabase")
        self.assertIs(span.kind, trace_api.SpanKind.CLIENT)

        self.assertEqual(span.attributes["component"], "testcomponent")
        self.assertEqual(span.attributes["db.type"], "testtype")
        self.assertEqual(span.attributes["db.instance"], "testdatabase")
        self.assertEqual(span.attributes["db.statement"], "Test query")
        self.assertEqual(
            span.attributes["db.statement.parameters"],
            "('param1Value', False)",
        )
        self.assertEqual(span.attributes["db.user"], "testuser")
        self.assertEqual(span.attributes["net.peer.name"], "testhost")
        self.assertEqual(span.attributes["net.peer.port"], 123)
        self.assertIs(
            span.status.canonical_code,
            trace_api.status.StatusCanonicalCode.OK,
        )
def instrument_connection(
    name: str,
    connection,
    database_system: str,
    connection_attributes: typing.Dict = None,
    version: str = "",
    tracer_provider: typing.Optional[TracerProvider] = None,
):
    """Enable instrumentation in a database connection.

    Args:
        name: Name of opentelemetry extension for aiopg.
        connection: The connection to instrument.
        database_system: An identifier for the database management system (DBMS)
            product being used.
        connection_attributes: Attribute names for database, port, host and
            user in a connection object.
        version: Version of opentelemetry extension for aiopg.
        tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
            use. If omitted the current configured one is used.

    Returns:
        An instrumented connection.
    """
    if isinstance(connection, AsyncProxyObject):
        logger.warning("Connection already instrumented")
        return connection

    db_integration = AiopgIntegration(
        name,
        database_system,
        connection_attributes=connection_attributes,
        version=version,
        tracer_provider=tracer_provider,
    )
    db_integration.get_connection_attributes(connection)
    return get_traced_connection_proxy(connection, db_integration)
예제 #12
0
 async def wrap_create_pool_(
     wrapped: typing.Callable[..., typing.Any],
     instance: typing.Any,
     args: typing.Tuple[typing.Any, typing.Any],
     kwargs: typing.Dict[typing.Any, typing.Any],
 ):
     db_integration = AiopgIntegration(
         name,
         database_component,
         database_type,
         connection_attributes=connection_attributes,
         version=version,
         tracer_provider=tracer_provider,
     )
     return await db_integration.wrapped_pool(wrapped, args, kwargs)