Use is_recording flag in aiopg, asyncpg, dbapi, psycopg2, pymemcache, pymongo, redis, sqlalchemy instrumentations (#1212)

This commit is contained in:
Leighton Chen
2020-10-08 15:25:20 -04:00
committed by GitHub
parent 59ee8b7f4d
commit dbc5f3b166
2 changed files with 52 additions and 20 deletions

View File

@ -79,15 +79,16 @@ class EngineTracer:
def _before_cur_exec(self, conn, cursor, statement, *args):
self.current_span = self.tracer.start_span(self.name)
with self.tracer.use_span(self.current_span, end_on_exit=False):
self.current_span.set_attribute("service", self.vendor)
self.current_span.set_attribute(_STMT, statement)
if self.current_span.is_recording():
self.current_span.set_attribute("service", self.vendor)
self.current_span.set_attribute(_STMT, statement)
if not _set_attributes_from_url(
self.current_span, conn.engine.url
):
_set_attributes_from_cursor(
self.current_span, self.vendor, cursor
)
if not _set_attributes_from_url(
self.current_span, conn.engine.url
):
_set_attributes_from_cursor(
self.current_span, self.vendor, cursor
)
# pylint: disable=unused-argument
def _after_cur_exec(self, conn, cursor, statement, *args):
@ -95,7 +96,11 @@ class EngineTracer:
return
try:
if cursor and cursor.rowcount >= 0:
if (
cursor
and cursor.rowcount >= 0
and self.current_span.is_recording()
):
self.current_span.set_attribute(_ROWS, cursor.rowcount)
finally:
self.current_span.end()
@ -105,30 +110,34 @@ class EngineTracer:
return
try:
self.current_span.set_status(
Status(
StatusCanonicalCode.UNKNOWN,
str(context.original_exception),
if self.current_span.is_recording():
self.current_span.set_status(
Status(
StatusCanonicalCode.UNKNOWN,
str(context.original_exception),
)
)
)
finally:
self.current_span.end()
def _set_attributes_from_url(span: trace.Span, url):
"""Set connection tags from the url. return true if successful."""
if url.host:
span.set_attribute(_HOST, url.host)
if url.port:
span.set_attribute(_PORT, url.port)
if url.database:
span.set_attribute(_DB, url.database)
if span.is_recording():
if url.host:
span.set_attribute(_HOST, url.host)
if url.port:
span.set_attribute(_PORT, url.port)
if url.database:
span.set_attribute(_DB, url.database)
return bool(url.host)
def _set_attributes_from_cursor(span: trace.Span, vendor, cursor):
"""Attempt to set db connection attributes by introspecting the cursor."""
if not span.is_recording():
return
if vendor == "postgres":
# pylint: disable=import-outside-toplevel
from psycopg2.extensions import parse_dsn

View File

@ -11,6 +11,7 @@
# 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 unittest import mock
from sqlalchemy import create_engine
@ -37,6 +38,28 @@ class TestSqlalchemyInstrumentation(TestBase):
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0].name, "sqlite.query")
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
mock_tracer.use_span.return_value.__enter__ = mock_span
mock_tracer.use_span.return_value.__exit__ = True
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
tracer.return_value = mock_tracer
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
service="my-database",
)
cnx = engine.connect()
cnx.execute("SELECT 1 + 1;").fetchall()
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)
def test_create_engine_wrapper(self):
SQLAlchemyInstrumentor().instrument()
from sqlalchemy import create_engine # pylint: disable-all