Skip to content

Commit 9e4ad27

Browse files
ag-ramachandranag-ramachandran
committed
* Refactor to classes
1 parent ca981ba commit 9e4ad27

File tree

8 files changed

+499
-295
lines changed

8 files changed

+499
-295
lines changed

‎lib/logstash/outputs/kusto.rb‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
require 'logstash/outputs/kusto/ingestor'
88
require 'logstash/outputs/kusto/interval'
99
require 'logstash/outputs/kusto/custom_size_based_buffer'
10+
require 'logstash/outputs/kusto/kustoLogstashConfiguration'
1011

1112
##
1213
# This plugin sends messages to Azure Kusto in batches.
@@ -77,8 +78,8 @@ class LogStash::Outputs::Kusto < LogStash::Outputs::Base
7778
default :codec, 'json_lines'
7879

7980
def register
80-
# Initialize the custom buffer with size, interval, and buffer file
81-
@buffer = LogStash::Outputs::CustomSizeBasedBuffer.new(@max_size, @max_interval) do |events|
81+
# Initialize the custom buffer with size and interval
82+
@buffer = LogStash::Outputs::KustoInternal::CustomSizeBasedBuffer.new(@max_size, @max_interval) do |events|
8283
flush_buffer(events)
8384
end
8485

@@ -91,13 +92,13 @@ def register
9192
max_threads: upload_concurrent_count,
9293
max_queue: upload_queue_size,
9394
fallback_policy: :caller_runs)
94-
95-
@ingestor = Ingestor.new(ingest_url, app_id, app_key, app_tenant, managed_identity, cli_auth, database, table, final_mapping, proxy_host, proxy_port, proxy_protocol, @logger, executor)
96-
97-
# Deprecation warning for path
98-
if @path
99-
@logger.warn("The 'path' configuration option is deprecated and will be removed in a future release.")
100-
end
95+
96+
kusto_ingest_base = LogStash::Outputs::KustoInternal::KustoIngestConfiguration.new(ingest_url, database, table, final_mapping)
97+
kusto_auth_base = LogStash::Outputs::KustoInternal::KustoAuthConfiguration.new(app_id, app_key, app_tenant, managed_identity, cli_auth)
98+
kusto_proxy_base = LogStash::Outputs::KustoInternal::KustoProxyConfiguration.new(proxy_host , proxy_port , proxy_protocol, false)
99+
@kusto_logstash_configuration = LogStash::Outputs::KustoInternal::KustoLogstashConfiguration.new(kusto_ingest_base, kusto_auth_base , kusto_proxy_base, logger)
100+
@ingestor = Ingestor.new(@kusto_logstash_configuration, @logger, executor)
101+
101102
end
102103

103104

@@ -114,8 +115,7 @@ def multi_receive_encoded(events_and_encoded)
114115

115116
def close
116117
@logger.info("Closing Kusto output plugin")
117-
118-
begin
118+
begin
119119
@buffer.shutdown unless @buffer.nil?
120120
@logger.info("Buffer shutdown") unless @buffer.nil?
121121
rescue => e

‎lib/logstash/outputs/kusto/custom_size_based_buffer.rb‎

Lines changed: 74 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -2,69 +2,70 @@
22
require 'thread'
33

44
module LogStash
5-
module Outputs
6-
class CustomSizeBasedBuffer
7-
def initialize(max_size_mb, max_interval, &flush_callback)
8-
@buffer_config = {
9-
max_size: max_size_mb * 1024 * 1024, # Convert MB to bytes
10-
max_interval: max_interval,
11-
logger: Logger.new(STDOUT)
12-
}
13-
@buffer_state = {
14-
pending_items: [],
15-
pending_size: 0,
16-
last_flush: Time.now.to_i,
17-
timer: Thread.new do
18-
loop do
19-
sleep(@buffer_config[:max_interval])
20-
buffer_flush(force: true)
21-
end
22-
end
23-
}
24-
@flush_callback = flush_callback
25-
@shutdown = false
26-
@pending_mutex = Mutex.new
27-
@flush_mutex = Mutex.new
28-
@buffer_config[:logger].info("CustomSizeBasedBuffer initialized with max_size: #{max_size_mb} MB, max_interval: #{max_interval} seconds")
29-
end
30-
31-
def <<(event)
32-
while buffer_full? do
33-
sleep 0.1
34-
end
35-
36-
@pending_mutex.synchronize do
37-
@buffer_state[:pending_items] << event
38-
@buffer_state[:pending_size] += event.bytesize
39-
end
40-
41-
buffer_flush
42-
end
43-
44-
def shutdown
45-
@buffer_config[:logger].info("Shutting down buffer")
46-
@shutdown = true
47-
@buffer_state[:timer].kill
48-
buffer_flush(final: true)
49-
end
50-
51-
private
52-
53-
def buffer_full?
54-
@pending_mutex.synchronize do
55-
@buffer_state[:pending_size] >= @buffer_config[:max_size]
56-
end
57-
end
58-
59-
def buffer_flush(options = {})
60-
force = options[:force] || options[:final]
61-
final = options[:final]
62-
63-
if final
64-
@flush_mutex.lock
65-
elsif !@flush_mutex.try_lock
66-
return 0
67-
end
5+
module Outputs
6+
module KustoInternal
7+
class CustomSizeBasedBuffer
8+
def initialize(max_size_mb, max_interval, &flush_callback)
9+
@buffer_config = {
10+
max_size: max_size_mb * 1024 * 1024, # Convert MB to bytes
11+
max_interval: max_interval,
12+
logger: Logger.new(STDOUT)
13+
}
14+
@buffer_state = {
15+
pending_items: [],
16+
pending_size: 0,
17+
last_flush: Time.now.to_i,
18+
timer: Thread.new do
19+
loop do
20+
sleep(@buffer_config[:max_interval])
21+
buffer_flush(force: true)
22+
end
23+
end
24+
}
25+
@flush_callback = flush_callback
26+
@shutdown = false
27+
@pending_mutex = Mutex.new
28+
@flush_mutex = Mutex.new
29+
@buffer_config[:logger].info("CustomSizeBasedBuffer initialized with max_size: #{max_size_mb} MB, max_interval: #{max_interval} seconds")
30+
end
31+
32+
def <<(event)
33+
while buffer_full? do
34+
sleep 0.1
35+
end
36+
37+
@pending_mutex.synchronize do
38+
@buffer_state[:pending_items] << event
39+
@buffer_state[:pending_size] += event.bytesize
40+
end
41+
42+
buffer_flush
43+
end
44+
45+
def shutdown
46+
@buffer_config[:logger].info("Shutting down buffer")
47+
@shutdown = true
48+
@buffer_state[:timer].kill
49+
buffer_flush(final: true)
50+
end
51+
52+
private
53+
54+
def buffer_full?
55+
@pending_mutex.synchronize do
56+
@buffer_state[:pending_size] >= @buffer_config[:max_size]
57+
end
58+
end
59+
60+
def buffer_flush(options = {})
61+
force = options[:force] || options[:final]
62+
final = options[:final]
63+
64+
if final
65+
@flush_mutex.lock
66+
elsif !@flush_mutex.try_lock
67+
return 0
68+
end
6869

6970
items_flushed = 0
7071
max_retries = 5
@@ -77,7 +78,7 @@ def buffer_flush(options = {})
7778
@pending_mutex.synchronize do
7879
return 0 if @buffer_state[:pending_size] == 0
7980

80-
time_since_last_flush = Time.now.to_i - @buffer_state[:last_flush]
81+
time_since_last_flush = Time.now.to_i - @buffer_state[:last_flush]
8182

8283
if !force && @buffer_state[:pending_size] < @buffer_config[:max_size] && time_since_last_flush < @buffer_config[:max_interval]
8384
return 0
@@ -121,13 +122,14 @@ def buffer_flush(options = {})
121122
@flush_mutex.unlock
122123
end
123124

124-
items_flushed
125-
end
125+
items_flushed
126+
end
126127

127-
def buffer_initialize
128-
@buffer_state[:pending_items] = []
129-
@buffer_state[:pending_size] = 0
130-
end
131-
end
132-
end
128+
def buffer_initialize
129+
@buffer_state[:pending_items] = []
130+
@buffer_state[:pending_size] = 0
131+
end
132+
end
133+
end
134+
end
133135
end

‎lib/logstash/outputs/kusto/ingestor.rb‎

Lines changed: 21 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -20,38 +20,39 @@ class Ingestor
2020
LOW_QUEUE_LENGTH = 3
2121
FIELD_REF = /%\{[^}]+\}/
2222

23-
def initialize(ingest_url, app_id, app_key, app_tenant, managed_identity_id, cli_auth, database, table, json_mapping, proxy_host , proxy_port , proxy_protocol,logger, threadpool = DEFAULT_THREADPOOL)
23+
def initialize(kusto_logstash_configuration, logger, threadpool = DEFAULT_THREADPOOL)
2424
@workers_pool = threadpool
2525
@logger = logger
26-
validate_config(database, table, json_mapping,proxy_protocol,app_id, app_key, managed_identity_id,cli_auth)
26+
#Validate and assign
27+
kusto_logstash_configuration.validate_config()
28+
@kusto_logstash_configuration = kusto_logstash_configuration
29+
2730
@logger.info('Preparing Kusto resources.')
2831

2932
kusto_java = Java::com.microsoft.azure.kusto
3033
apache_http = Java::org.apache.http
31-
# kusto_connection_string = kusto_java.data.auth.ConnectionStringBuilder.createWithAadApplicationCredentials(ingest_url, app_id, app_key.value, app_tenant)
32-
# If there is managed identity, use it. This means the AppId and AppKey are empty/nil
33-
# If there is CLI Auth, use that instead of managed identity
34-
is_managed_identity = (app_id.nil? && app_key.nil? && !cli_auth)
34+
35+
is_managed_identity = @kusto_logstash_configuration.kusto_auth.is_managed_identity
3536
# If it is system managed identity, propagate the system identity
36-
is_system_assigned_managed_identity = is_managed_identity && 0 == "system".casecmp(managed_identity_id)
37+
is_system_assigned_managed_identity = @kusto_logstash_configuration.kusto_auth.is_system_assigned_managed_identity
3738
# Is it direct connection
38-
is_direct_conn = (proxy_host.nil? || proxy_host.empty?)
39+
is_direct_conn = @kusto_logstash_configuration.kusto_proxy.is_direct_conn
3940
# Create a connection string
4041
kusto_connection_string = if is_managed_identity
4142
if is_system_assigned_managed_identity
4243
@logger.info('Using system managed identity.')
43-
kusto_java.data.auth.ConnectionStringBuilder.createWithAadManagedIdentity(ingest_url)
44+
kusto_java.data.auth.ConnectionStringBuilder.createWithAadManagedIdentity(@kusto_logstash_configuration.kusto_ingest.ingest_url)
4445
else
4546
@logger.info('Using user managed identity.')
46-
kusto_java.data.auth.ConnectionStringBuilder.createWithAadManagedIdentity(ingest_url, managed_identity_id)
47+
kusto_java.data.auth.ConnectionStringBuilder.createWithAadManagedIdentity(@kusto_logstash_configuration.kusto_ingest.ingest_url, @kusto_logstash_configuration.kusto_ingest.managed_identity_id)
4748
end
4849
else
49-
if cli_auth
50+
if @kusto_logstash_configuration.kusto_auth.cli_auth
5051
@logger.warn('*Use of CLI Auth is only for dev-test scenarios. This is ***NOT RECOMMENDED*** for production*')
51-
kusto_java.data.auth.ConnectionStringBuilder.createWithAzureCli(ingest_url)
52+
kusto_java.data.auth.ConnectionStringBuilder.createWithAzureCli(@kusto_logstash_configuration.kusto_ingest.ingest_url)
5253
else
5354
@logger.info('Using app id and app key.')
54-
kusto_java.data.auth.ConnectionStringBuilder.createWithAadApplicationCredentials(ingest_url, app_id, app_key.value, app_tenant)
55+
kusto_java.data.auth.ConnectionStringBuilder.createWithAadApplicationCredentials(@kusto_logstash_configuration.kusto_ingest.ingest_url, @kusto_logstash_configuration.kusto_auth.app_id, @kusto_logstash_configuration.kusto_auth.app_key.value, @kusto_logstash_configuration.kusto_auth.app_tenant)
5556
end
5657
end
5758
#
@@ -63,22 +64,22 @@ def initialize(ingest_url, app_id, app_key, app_tenant, managed_identity_id, cli
6364
tuple_utils = Java::org.apache.commons.lang3.tuple
6465
# kusto_connection_string.setClientVersionForTracing(name_for_tracing)
6566
version_for_tracing=Gem.loaded_specs['logstash-output-kusto']&.version || "unknown"
66-
kusto_connection_string.setConnectorDetails("Logstash",version_for_tracing.to_s,"","",false,"", tuple_utils.Pair.emptyArray());
67+
kusto_connection_string.setConnectorDetails("Logstash",version_for_tracing.to_s,name_for_tracing.to_s,version_for_tracing.to_s,false,"", tuple_utils.Pair.emptyArray());
6768

6869
@kusto_client = begin
6970
if is_direct_conn
7071
kusto_java.ingest.IngestClientFactory.createClient(kusto_connection_string)
7172
else
72-
kusto_http_client_properties = kusto_java.data.HttpClientProperties.builder().proxy(apache_http.HttpHost.new(proxy_host,proxy_port,proxy_protocol)).build()
73+
kusto_http_client_properties = kusto_java.data.HttpClientProperties.builder().proxy(apache_http.HttpHost.new(@kusto_logstash_configuration.kusto_proxy.proxy_host,@kusto_logstash_configuration.kusto_proxy.proxy_port,@kusto_logstash_configuration.kusto_proxy.proxy_protocol)).build()
7374
kusto_java.ingest.IngestClientFactory.createClient(kusto_connection_string, kusto_http_client_properties)
7475
end
7576
end
7677

77-
@ingestion_properties = kusto_java.ingest.IngestionProperties.new(database, table)
78-
is_mapping_ref_provided = !(json_mapping.nil? || json_mapping.empty?)
79-
if is_mapping_ref_provided
80-
@logger.debug('Using mapping reference.', json_mapping)
81-
@ingestion_properties.setIngestionMapping(json_mapping, kusto_java.ingest.IngestionMapping::IngestionMappingKind::JSON)
78+
@ingestion_properties = kusto_java.ingest.IngestionProperties.new(@kusto_logstash_configuration.kusto_ingest.database, @kusto_logstash_configuration.kusto_ingest.table)
79+
80+
if @kusto_logstash_configuration.kusto_ingest.is_mapping_ref_provided
81+
@logger.debug('Using mapping reference.', @kusto_logstash_configuration.kusto_ingest.json_mapping)
82+
@ingestion_properties.setIngestionMapping(@kusto_logstash_configuration.kusto_ingest.json_mapping, kusto_java.ingest.IngestionMapping::IngestionMappingKind::JSON)
8283
@ingestion_properties.setDataFormat(kusto_java.ingest.IngestionProperties::DataFormat::JSON)
8384
else
8485
@logger.debug('No mapping reference provided. Columns will be mapped by names in the logstash output')
@@ -87,38 +88,6 @@ def initialize(ingest_url, app_id, app_key, app_tenant, managed_identity_id, cli
8788
@logger.debug('Kusto resources are ready.')
8889
end
8990

90-
def validate_config(database, table, json_mapping, proxy_protocol, app_id, app_key, managed_identity_id,cli_auth)
91-
# Add an additional validation and fail this upfront
92-
if app_id.nil? && app_key.nil? && managed_identity_id.nil?
93-
if cli_auth
94-
@logger.info('Using CLI Auth, this is only for dev-test scenarios. This is ***NOT RECOMMENDED*** for production')
95-
else
96-
@logger.error('managed_identity_id is not provided and app_id/app_key is empty.')
97-
raise LogStash::ConfigurationError.new('managed_identity_id is not provided and app_id/app_key is empty.')
98-
end
99-
end
100-
if database =~ FIELD_REF
101-
@logger.error('database config value should not be dynamic.', database)
102-
raise LogStash::ConfigurationError.new('database config value should not be dynamic.')
103-
end
104-
105-
if table =~ FIELD_REF
106-
@logger.error('table config value should not be dynamic.', table)
107-
raise LogStash::ConfigurationError.new('table config value should not be dynamic.')
108-
end
109-
110-
if json_mapping =~ FIELD_REF
111-
@logger.error('json_mapping config value should not be dynamic.', json_mapping)
112-
raise LogStash::ConfigurationError.new('json_mapping config value should not be dynamic.')
113-
end
114-
115-
if not(["https", "http"].include? proxy_protocol)
116-
@logger.error('proxy_protocol has to be http or https.', proxy_protocol)
117-
raise LogStash::ConfigurationError.new('proxy_protocol has to be http or https.')
118-
end
119-
120-
end
121-
12291
def upload_async(data)
12392
if @workers_pool.remaining_capacity <= LOW_QUEUE_LENGTH
12493
@logger.warn("Ingestor queue capacity is running low with #{@workers_pool.remaining_capacity} free slots.")

0 commit comments

Comments
 (0)