-
Notifications
You must be signed in to change notification settings - Fork 369
/
instrumentation.rb
110 lines (95 loc) · 3.7 KB
/
instrumentation.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
require 'ddtrace/ext/net'
require 'ddtrace/ext/sql'
require 'ddtrace/ext/app_types'
require 'ddtrace/contrib/presto/ext'
module Datadog
module Contrib
module Presto
# Instrumentation for Presto integration
module Instrumentation
# Instrumentation for Presto::Client::Client
module Client
def self.included(base)
base.send(:prepend, InstanceMethods)
end
# Instance methods for Presto::Client
module InstanceMethods
def run(query)
tracer.trace(Ext::SPAN_QUERY, span_options) do |span|
begin
decorate!(span)
span.resource = query
span.span_type = Datadog::Ext::SQL::TYPE
span.set_tag(Ext::TAG_QUERY_ASYNC, false)
rescue StandardError => e
Datadog.logger.debug("error preparing span for presto: #{e}")
end
super(query)
end
end
def query(query, &blk)
tracer.trace(Ext::SPAN_QUERY, span_options) do |span|
begin
decorate!(span)
span.resource = query
span.span_type = Datadog::Ext::SQL::TYPE
span.set_tag(Ext::TAG_QUERY_ASYNC, !blk.nil?)
rescue StandardError => e
Datadog.logger.debug("error preparing span for presto: #{e}")
end
super(query, &blk)
end
end
def kill(query_id)
tracer.trace(Ext::SPAN_KILL, span_options) do |span|
begin
decorate!(span)
span.resource = Ext::SPAN_KILL
span.span_type = Datadog::Ext::AppTypes::DB
# ^ not an SQL type span, since there's no SQL query
span.set_tag(Ext::TAG_QUERY_ID, query_id)
rescue StandardError => e
Datadog.logger.debug("error preparing span for presto: #{e}")
end
super(query_id)
end
end
private
def datadog_configuration
Datadog.configuration[:presto]
end
def span_options
{
service: datadog_configuration[:service_name],
app: Ext::APP,
app_type: Datadog::Ext::AppTypes::DB
}
end
def tracer
datadog_configuration.tracer
end
def decorate!(span)
set_nilable_tag!(span, :server, Datadog::Ext::NET::TARGET_HOST)
set_nilable_tag!(span, :user, Ext::TAG_USER_NAME)
set_nilable_tag!(span, :schema, Ext::TAG_SCHEMA_NAME)
set_nilable_tag!(span, :catalog, Ext::TAG_CATALOG_NAME)
set_nilable_tag!(span, :time_zone, Ext::TAG_TIME_ZONE)
set_nilable_tag!(span, :language, Ext::TAG_LANGUAGE)
set_nilable_tag!(span, :http_proxy, Ext::TAG_PROXY)
set_nilable_tag!(span, :model_version, Ext::TAG_MODEL_VERSION)
# Tag as an external peer service
span.set_tag(Datadog::Ext::Integration::TAG_PEER_SERVICE, span.service)
# Set analytics sample rate
if Contrib::Analytics.enabled?(datadog_configuration[:analytics_enabled])
Contrib::Analytics.set_sample_rate(span, datadog_configuration[:analytics_sample_rate])
end
end
def set_nilable_tag!(span, key, tag_name)
@options[key].tap { |val| span.set_tag(tag_name, val) if val }
end
end
end
end
end
end
end