def test_instrumentor(self, mock_connect):
Psycopg2Instrumentor().instrument()
cnx = psycopg2.connect(database="test")
cursor = cnx.cursor()
query = "SELECT * FROM test"
cursor.execute(query)
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
# Check version and name in span's instrumentation info
self.check_span_instrumentation_info(
span, opentelemetry.instrumentation.psycopg2)
# check that no spans are generated after uninstrument
Psycopg2Instrumentor().uninstrument()
cnx = psycopg2.connect(database="test")
cursor = cnx.cursor()
query = "SELECT * FROM test"
cursor.execute(query)
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
def tearDown(self):
super().tearDown()
self.memory_exporter.clear()
self.cursor_mock.stop()
self.connection_mock.stop()
with self.disable_logging():
Psycopg2Instrumentor().uninstrument()
def test_uninstrument_connection(self, mock_connect):
Psycopg2Instrumentor().instrument()
cnx = psycopg2.connect(database="test")
query = "SELECT * FROM test"
cursor = cnx.cursor()
cursor.execute(query)
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
cnx = Psycopg2Instrumentor().uninstrument_connection(cnx)
cursor = cnx.cursor()
cursor.execute(query)
spans_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
def test_sqlcommenter_disabled(self, event_mocked):
cnx = psycopg2.connect(database="test")
Psycopg2Instrumentor().instrument()
query = "SELECT * FROM test"
cursor = cnx.cursor()
cursor.execute(query)
kwargs = event_mocked.call_args[1]
self.assertEqual(kwargs["enable_commenter"], False)
def test_not_recording(self):
mock_tracer = mock.Mock()
mock_span = mock.Mock()
mock_span.is_recording.return_value = False
mock_tracer.start_span.return_value = mock_span
Psycopg2Instrumentor().instrument()
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
tracer.return_value = mock_tracer
cnx = psycopg2.connect(database="test")
cursor = cnx.cursor()
query = "SELECT * FROM test"
cursor.execute(query)
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)
Psycopg2Instrumentor().uninstrument()
def instrument_app(self) -> None:
logger.info("Activating Opentelemetry tracing to app", app=self.title)
trace.set_tracer_provider(tracer_provider)
FastAPIInstrumentor.instrument_app(self)
RequestsInstrumentor().instrument()
HTTPXClientInstrumentor().instrument()
RedisInstrumentor().instrument()
Psycopg2Instrumentor().instrument()
SQLAlchemyInstrumentor().instrument(engine=db.engine,
tracer_provider=tracer_provider)
def setUp(self):
super().setUp()
self._tracer = self.tracer_provider.get_tracer(__name__)
Psycopg2Instrumentor().instrument(tracer_provider=self.tracer_provider)
self._connection = psycopg2.connect(
dbname=POSTGRES_DB_NAME,
user=POSTGRES_USER,
password=POSTGRES_PASSWORD,
host=POSTGRES_HOST,
port=POSTGRES_PORT,
)
self._connection.set_session(autocommit=True)
self._cursor = self._connection.cursor()
def test_commenter_enabled(self):
stmt = "CREATE TABLE IF NOT EXISTS users (id integer, name varchar)"
with self._tracer.start_as_current_span("rootSpan"):
self._cursor.execute(stmt)
self.validate_spans("CREATE")
Psycopg2Instrumentor().uninstrument()
Psycopg2Instrumentor().instrument(enable_commenter=True)
self._cursor.execute(
sql.SQL("SELECT FROM {table} where {field}='{value}'").format(
table=sql.Identifier("users"),
field=sql.Identifier("name"),
value=sql.Identifier("abc"),
)
)
spans = self.memory_exporter.get_finished_spans()
span = spans[2]
self.assertEqual(span.name, "SELECT")
self.assertEqual(
span.attributes[SpanAttributes.DB_STATEMENT],
'SELECT FROM "users" where "name"=\'"abc"\'',
)
def test_custom_tracer_provider(self, mock_connect):
resource = resources.Resource.create({})
result = self.create_tracer_provider(resource=resource)
tracer_provider, exporter = result
Psycopg2Instrumentor().instrument(tracer_provider=tracer_provider)
cnx = psycopg2.connect(database="test")
cursor = cnx.cursor()
query = "SELECT * FROM test"
cursor.execute(query)
spans_list = exporter.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
self.assertIs(span.resource, resource)
from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import (BatchSpanProcessor,
ConsoleSpanExporter,
SimpleSpanProcessor)
# Set up OpenTelemetry tracing
trace.set_tracer_provider(TracerProvider())
trace.get_tracer_provider().add_span_processor(
SimpleSpanProcessor(ConsoleSpanExporter()))
trace.get_tracer_provider().add_span_processor(
BatchSpanProcessor(CloudTraceSpanExporter(), schedule_delay_millis=5000))
# Trace postgres queries as well
Psycopg2Instrumentor().instrument()
import psycopg2
from google.cloud.sqlcommenter.psycopg2.extension import CommenterCursorFactory
tracer = trace.get_tracer(__name__)
def main():
cursor_factory = CommenterCursorFactory(
with_db_driver=True,
with_dbapi_level=True,
with_dbapi_threadsafety=True,
with_driver_paramstyle=True,
with_libpq_version=True,
with_opentelemetry=True,
trace.set_tracer_provider(TracerProvider())
tracer = trace.get_tracer(__name__)
jaeger_exporter = jaeger.JaegerSpanExporter(
service_name="point_service",
agent_host_name="jaeger",
agent_port=6831,
)
span_processor = BatchExportSpanProcessor(jaeger_exporter)
trace.get_tracer_provider().add_span_processor(span_processor)
app = Flask("point_app")
FlaskInstrumentor().instrument_app(app)
Psycopg2Instrumentor().instrument(tracer_provider=trace.get_tracer_provider())
postgres_con = psycopg2.connect(
database="postgres",
user="******",
host="db",
port="5432",
)
@app.route("/", methods=["POST"])
def add_point():
data, errors = AddPointValidator.validate_or_error(request.json)
if errors:
return jsonify(dict(errors)), 400
with postgres_con:
cur = postgres_con.cursor()
def tearDownClass(cls):
if cls._cursor:
cls._cursor.close()
if cls._connection:
cls._connection.close()
Psycopg2Instrumentor().uninstrument()
def tearDown(self):
super().tearDown()
with self.disable_logging():
Psycopg2Instrumentor().uninstrument()
def tearDown(self):
self._cursor.close()
self._connection.close()
Psycopg2Instrumentor().uninstrument()
super().tearDown()
def instrument_psycopg2():
from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
Psycopg2Instrumentor().instrument()
