Skip to content

Commit

Permalink
OpenTelemetry support #230
Browse files Browse the repository at this point in the history
  • Loading branch information
xzkostyan committed Aug 6, 2021
1 parent f1b56bb commit 9b88188
Show file tree
Hide file tree
Showing 7 changed files with 188 additions and 6 deletions.
17 changes: 15 additions & 2 deletions clickhouse_driver/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,21 @@ class Client(object):
* ``strings_encoding`` -- specifies string encoding. UTF-8 by default.
* ``use_numpy`` -- Use numpy for columns reading. New in version
*0.2.0*.
* ``opentelemetry_traceparent`` -- OpenTelemetry traceparent header as
described by W3C Trace Context recommendation.
New in version *0.2.2*.
* ``opentelemetry_tracestate`` -- OpenTelemetry tracestate header as
described by W3C Trace Context recommendation.
New in version *0.2.2*.
"""

available_client_settings = (
'insert_block_size', # TODO: rename to max_insert_block_size
'strings_as_bytes',
'strings_encoding',
'use_numpy'
'use_numpy',
'opentelemetry_traceparent',
'opentelemetry_tracestate'
)

def __init__(self, *args, **kwargs):
Expand All @@ -65,6 +72,12 @@ def __init__(self, *args, **kwargs):
),
'use_numpy': self.settings.pop(
'use_numpy', False
),
'opentelemetry_traceparent': self.settings.pop(
'opentelemetry_traceparent', None
),
'opentelemetry_tracestate': self.settings.pop(
'opentelemetry_tracestate', ''
)
}

Expand Down
23 changes: 21 additions & 2 deletions clickhouse_driver/clientinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@

from . import defines
from . import errors
from .opentelemetry import OpenTelemetryTraceContext
from .varint import write_varint
from .writer import write_binary_str, write_binary_uint8
from .writer import write_binary_str, write_binary_uint8, \
write_binary_uint64, write_binary_uint128


class ClientInfo(object):
Expand Down Expand Up @@ -34,7 +36,7 @@ class QueryKind(object):

quota_key = ''

def __init__(self, client_name):
def __init__(self, client_name, context):
self.query_kind = ClientInfo.QueryKind.NO_QUERY

try:
Expand All @@ -44,6 +46,11 @@ def __init__(self, client_name):
self.client_hostname = socket.gethostname()
self.client_name = client_name

self.client_trace_context = OpenTelemetryTraceContext(
context.client_settings['opentelemetry_traceparent'],
context.client_settings['opentelemetry_tracestate']
)

super(ClientInfo, self).__init__()

@property
Expand Down Expand Up @@ -78,3 +85,15 @@ def write(self, server_revision, fout):

if revision >= defines.DBMS_MIN_REVISION_WITH_VERSION_PATCH:
write_varint(self.client_version_patch, fout)

if revision >= defines.DBMS_MIN_REVISION_WITH_OPENTELEMETRY:
if self.client_trace_context.trace_id is not None:
# Have OpenTelemetry header.
write_binary_uint8(1, fout)
write_binary_uint128(self.client_trace_context.trace_id, fout)
write_binary_uint64(self.client_trace_context.span_id, fout)
write_binary_str(self.client_trace_context.tracestate, fout)
write_binary_uint8(self.client_trace_context.trace_flags, fout)
else:
# Don't have OpenTelemetry header.
write_binary_uint8(0, fout)
4 changes: 3 additions & 1 deletion clickhouse_driver/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,7 @@ def __init__(

self.connected = False

self.client_trace_context = None
self.server_info = None
self.context = Context()

Expand Down Expand Up @@ -347,6 +348,7 @@ def reset_state(self):

self.connected = False

self.client_trace_context = None
self.server_info = None

self.block_in = None
Expand Down Expand Up @@ -576,7 +578,7 @@ def send_query(self, query, query_id=None):

revision = self.server_info.revision
if revision >= defines.DBMS_MIN_REVISION_WITH_CLIENT_INFO:
client_info = ClientInfo(self.client_name)
client_info = ClientInfo(self.client_name, self.context)
client_info.query_kind = ClientInfo.QueryKind.INITIAL_QUERY

client_info.write(revision, self.fout)
Expand Down
3 changes: 2 additions & 1 deletion clickhouse_driver/defines.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO = 54420
DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS = 54429
DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET = 54441
DBMS_MIN_REVISION_WITH_OPENTELEMETRY = 54442

# Timeouts
DBMS_DEFAULT_CONNECT_TIMEOUT_SEC = 10
Expand All @@ -35,7 +36,7 @@
CLIENT_VERSION_MAJOR = 20
CLIENT_VERSION_MINOR = 10
CLIENT_VERSION_PATCH = 2
CLIENT_REVISION = 54441
CLIENT_REVISION = 54442

BUFFER_SIZE = 1048576

Expand Down
43 changes: 43 additions & 0 deletions clickhouse_driver/opentelemetry.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@

class OpenTelemetryTraceContext(object):
traceparent_tpl = 'xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxxx-xx'
translation = str.maketrans('1234567890abcdef', 'xxxxxxxxxxxxxxxx')

def __init__(self, traceparent, tracestate):
# xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxxx-xx
# ^ ^ ^ ^
# version trace_id span_id flags

self.trace_id = None # UUID
self.span_id = None # UInt64
self.tracestate = tracestate # String
self.trace_flags = None # UInt8

if traceparent is not None:
self.parse_traceparent(traceparent)

super(OpenTelemetryTraceContext, self).__init__()

def parse_traceparent(self, traceparent):
traceparent = traceparent.lower()

if len(traceparent) != len(self.traceparent_tpl):
raise ValueError('unexpected length {}, expected {}'.format(
len(traceparent), len(self.traceparent_tpl)
))

if traceparent.translate(self.translation) != self.traceparent_tpl:
raise ValueError(
'Malformed traceparant header: {}'.format(traceparent)
)

parts = traceparent.split('-')
version = int(parts[0], 16)
if version != 0:
raise ValueError(
'unexpected version {}, expected 00'.format(parts[0])
)

self.trace_id = (int(parts[1][16:], 16) << 64) + int(parts[1][:16], 16)
self.span_id = int(parts[2], 16)
self.trace_flags = int(parts[3], 16)
28 changes: 28 additions & 0 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,3 +227,31 @@ def test_settings_is_important(self):
def test_use_numpy(self):
c = Client.from_url('clickhouse://host?use_numpy=true')
self.assertTrue(c.connection.context.client_settings['use_numpy'])

def test_opentelemetry(self):
c = Client.from_url(
'clickhouse://host?opentelemetry_traceparent='
'00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00'
)
self.assertEqual(
c.connection.context.client_settings['opentelemetry_traceparent'],
'00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00'
)
self.assertEqual(
c.connection.context.client_settings['opentelemetry_tracestate'],
''
)

c = Client.from_url(
'clickhouse://host?opentelemetry_traceparent='
'00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00&'
'opentelemetry_tracestate=state'
)
self.assertEqual(
c.connection.context.client_settings['opentelemetry_traceparent'],
'00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00'
)
self.assertEqual(
c.connection.context.client_settings['opentelemetry_tracestate'],
'state'
)
76 changes: 76 additions & 0 deletions tests/test_opentelemetry.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
from tests.testcase import BaseTestCase
from tests.util import capture_logging


class OpenTelemetryTestCase(BaseTestCase):
required_server_version = (20, 11, 2)

def test_server_logs(self):
tracestate = 'tracestate'
traceparent = '00-1af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01'

settings = {
'opentelemetry_tracestate': tracestate,
'opentelemetry_traceparent': traceparent

}
with self.created_client(settings=settings) as client:
with capture_logging('clickhouse_driver.log', 'INFO') as buffer:
settings = {'send_logs_level': 'trace'}
query = 'SELECT 1'
client.execute(query, settings=settings)
value = buffer.getvalue()
self.assertIn('OpenTelemetry', value)
self.assertIn('1af7651916cd43dd8448eb211c80319c', value)

def test_no_tracestate(self):
traceparent = '00-1af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01'

settings = {
'opentelemetry_traceparent': traceparent

}
with self.created_client(settings=settings) as client:
with capture_logging('clickhouse_driver.log', 'INFO') as buffer:
settings = {'send_logs_level': 'trace'}
query = 'SELECT 1'
client.execute(query, settings=settings)
value = buffer.getvalue()
self.assertIn('OpenTelemetry', value)
self.assertIn('1af7651916cd43dd8448eb211c80319c', value)

def test_bad_traceparent(self):
settings = {'opentelemetry_traceparent': 'bad'}
with self.created_client(settings=settings) as client:
with self.assertRaises(ValueError) as e:
client.execute('SELECT 1')

self.assertEqual(
str(e.exception),
'unexpected length 3, expected 55'
)

traceparent = '00-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-yyyyyyyyyyyyyyyy-01'
settings = {'opentelemetry_traceparent': traceparent}
with self.created_client(settings=settings) as client:
with self.assertRaises(ValueError) as e:
client.execute('SELECT 1')

self.assertEqual(
str(e.exception),
'Malformed traceparant header: {}'.format(traceparent)
)

def test_bad_traceparent_version(self):
settings = {
'opentelemetry_traceparent':
'01-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01'
}
with self.created_client(settings=settings) as client:
with self.assertRaises(ValueError) as e:
client.execute('SELECT 1')

self.assertEqual(
str(e.exception),
'unexpected version 01, expected 00'
)

0 comments on commit 9b88188

Please sign in to comment.