fluent-plugin-kusto 0.0.1.beta

data/lib/fluent/plugin/client.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+# Client handles authentication and resource fetching for Azure Data Explorer (Kusto) ingestion.
+# It supports both Managed Identity and AAD Client Credentials authentication methods.
+#
+# The Client class is responsible for:
+# - Authenticating using Managed Identity or AAD Client Credentials
+# - Fetching and caching Kusto ingestion resources
+# - Providing access to blob SAS URI, queue SAS URI, and identity token
+require_relative 'auth/aad_tokenprovider'
+require_relative 'auth/mi_tokenprovider'
+require_relative 'auth/azcli_tokenprovider'
+require_relative 'auth/wif_tokenprovider'
+require_relative 'kusto_query'
+require 'logger'
+
+class Client
+  def initialize(outconfiguration)
+    # Set up queries for resource fetching
+    @user_query_blob_container = '.get ingestion resources'
+    @user_query_aad_token = '.get kusto identity token'
+    # Use provided logger or default to stdout
+    @logger = initialize_logger(outconfiguration)
+    @data_endpoint = outconfiguration.kusto_endpoint
+    @cached_resources = nil
+    @resources_expiry_time = nil
+    @outconfiguration = outconfiguration
+    @token_provider = create_token_provider(outconfiguration)
+  end
+
+  def resources
+    # Return cached resources if valid, otherwise fetch and cache
+    return @cached_resources if resources_cached?
+
+    fetch_and_cache_resources
+    @cached_resources
+  end
+
+  attr_reader :blob_sas_uri, :queue_sas_uri, :identity_token, :logger, :blob_rows, :data_endpoint, :token_provider
+
+  private
+
+  def initialize_logger(outconfiguration)
+    # Prefer logger from configuration, fallback to stdout
+    outconfiguration.logger
+  rescue StandardError
+    Logger.new($stdout)
+  end
+
+  def resources_cached?
+    # Check if resources are cached and not expired
+    @cached_resources && @resources_expiry_time && @resources_expiry_time > Time.now
+  end
+
+  def fetch_and_cache_resources
+    # Fetch resources from Kusto and cache them
+    @logger.info('Fetching resources from Kusto...')
+    blob_rows, aad_token_rows = fetch_kusto_resources
+    return unless blob_rows && aad_token_rows
+
+    blob_sas_uri, queue_sas_uri, identity_token = extract_resource_uris(blob_rows, aad_token_rows)
+    return unless validate_resource_uris(blob_sas_uri, queue_sas_uri, identity_token)
+
+    assign_and_cache_resources(blob_sas_uri, queue_sas_uri, identity_token)
+  end
+
+  def fetch_kusto_resources
+    # Fetch resource rows and validate them
+    blob_rows, aad_token_rows = fetch_kusto_rows_with_error_handling
+    validate_kusto_resource_rows(blob_rows, aad_token_rows)
+  end
+
+  def fetch_kusto_rows_with_error_handling
+    # Fetch blob and AAD token rows with error handling
+    blob_rows = fetch_blob_rows
+    aad_token_rows = fetch_aad_token_rows
+    [blob_rows, aad_token_rows]
+  end
+
+  def fetch_blob_rows
+    # Run Kusto query for blob resources
+    run_kusto_api_query(@user_query_blob_container, @data_endpoint, @token_provider,
+                        use_ingest_endpoint: true)
+  rescue StandardError => e
+    @logger.error("Failed to fetch blob resources from Kusto: #{e.message}")
+    nil
+  end
+
+  def fetch_aad_token_rows
+    # Run Kusto query for AAD token resources
+    run_kusto_api_query(@user_query_aad_token, @data_endpoint, @token_provider,
+                        use_ingest_endpoint: true, database_name: nil)
+  rescue StandardError => e
+    @logger.error("Failed to fetch AAD token resources from Kusto: #{e.message}")
+    nil
+  end
+
+  def create_token_provider(outconfiguration)
+    case outconfiguration.auth_type&.downcase
+    when 'aad'
+      AadTokenProvider.new(outconfiguration)
+    when 'azcli'
+      AzCliTokenProvider.new(outconfiguration)
+    when 'workload_identity'
+      WorkloadIdentity.new(outconfiguration)
+    when 'user_managed_identity', 'system_managed_identity'
+      ManagedIdentityTokenProvider.new(outconfiguration)
+    else
+      raise "Unknown auth_type: #{outconfiguration.auth_type}"
+    end
+  end
+
+  def extract_resource_uris(blob_rows, aad_token_rows)
+    # Extract URIs from resource rows
+    blob_sas_uri = blob_rows.find { |row| row[0] == 'TempStorage' }&.[](1)
+    queue_sas_uri = blob_rows.find { |row| row[0] == 'SecuredReadyForAggregationQueue' }&.[](1)
+    identity_token = aad_token_rows[0][0] if aad_token_rows.any?
+    [blob_sas_uri, queue_sas_uri, identity_token]
+  end
+
+  def validate_resource_uris(blob_sas_uri, queue_sas_uri, identity_token)
+    # Ensure all required URIs are present
+    if blob_sas_uri.nil? || queue_sas_uri.nil? || identity_token.nil?
+      @logger.error('Failed to retrieve all required resources: blob_sas_uri, queue_s_uri, or identity_token is nil.')
+      return false
+    end
+    true
+  end
+
+  def assign_and_cache_resources(blob_sas_uri, queue_sas_uri, identity_token)
+    # Assign and cache resource URIs
+    @blob_sas_uri = blob_sas_uri
+    @queue_sas_uri = queue_sas_uri
+    @identity_token = identity_token
+    @cached_resources = {
+      blob_sas_uri: blob_sas_uri,
+      queue_sas_uri: queue_sas_uri,
+      identity_token: identity_token
+    }
+    @resources_expiry_time = Time.now + 21_600 # Cache for 6 hours
+  end
+
+  def validate_kusto_resource_rows(blob_rows, aad_token_rows)
+    # Validate resource rows are present
+    if blob_rows.nil? || blob_rows.empty?
+      @logger.error('No blob rows found in the response.')
+      return [nil, nil]
+    end
+    if aad_token_rows.nil? || aad_token_rows.empty?
+      @logger.error('No AAD token rows found in the response.')
+      return [nil, nil]
+    end
+    [blob_rows, aad_token_rows]
+  end
+end

data/lib/fluent/plugin/conffile.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+# OutputConfiguration holds and validates configuration for the Kusto output plugin.
+# It supports both Managed Identity and AAD Client Credentials authentication methods.
+#
+# Responsibilities:
+# - Store configuration options
+# - Validate configuration for Managed Identity or AAD
+# - Provide Azure AD endpoint based on cloud
+
+require 'logger'
+
+AZURE_CLOUDS = {
+  'AzureCloud' => { 'aad' => 'https://login.microsoftonline.com' },
+  'AzureChinaCloud' => { 'aad' => 'https://login.chinacloudapi.cn' },
+  'AzureUSGovernment' => { 'aad' => 'https://login.microsoftonline.us' }
+}.freeze
+
+class OutputConfiguration
+  def initialize(opts = {})
+    # Initialize configuration options and logger
+    @logger_path = opts[:logger_path]
+    @logger = initialize_logger
+    @client_app_id = opts[:client_app_id]
+    @client_app_secret = opts[:client_app_secret]
+    @tenant_id = opts[:tenant_id]
+    @kusto_endpoint = opts[:kusto_endpoint]
+    @database_name = opts[:database_name]
+    @table_name = opts[:table_name]
+    @azure_cloud = opts[:azure_cloud] || 'AzureCloud'
+    @managed_identity_client_id = opts[:managed_identity_client_id]
+    @azure_clouds = AZURE_CLOUDS
+    @auth_type = opts[:auth_type] || 'aad'
+    @workload_identity_client_id = opts[:workload_identity_client_id]
+    @workload_identity_tenant_id = opts[:workload_identity_tenant_id]
+    @workload_identity_token_file_path = opts[:workload_identity_token_file_path]
+    validate_configuration
+  end
+
+  def validate_configuration
+    # Validate configuration based on authentication method
+    case @auth_type&.downcase
+    when 'aad'
+      validate_aad_config
+    when 'user_managed_identity', 'system_managed_identity', 'azcli'
+      validate_base_config
+    when 'workload_identity'
+      validate_workload_identity_config
+    else
+      raise ArgumentError, "Unknown auth_type: #{@auth_type}"
+    end
+    validate_azure_cloud
+    true
+  end
+
+  def print_missing_parameter_message_and_raise(param_name)
+    # Print error and raise if a required parameter is missing
+    @logger.error(
+      "Missing a required setting for the Kusto output plugin configuration:\n" \
+      "output {\nkusto {\n#{param_name} => # SETTING MISSING\n ...\n}\n}\n"
+    )
+    raise ArgumentError, "The setting #{param_name} is required for Kusto configuration."
+  end
+
+  attr_reader :logger, :client_app_id, :client_app_secret, :tenant_id, :kusto_endpoint, :database_name, :table_name,
+              :managed_identity_client_id, :azure_cloud, :auth_type, :workload_identity_client_id,
+              :workload_identity_tenant_id, :workload_identity_token_file_path
+
+  def aad_endpoint
+    # Return Azure AD endpoint for selected cloud
+    @azure_clouds[@azure_cloud]['aad']
+  end
+
+  private
+
+  def initialize_logger
+    # Use logger_path if provided, otherwise log to stdout
+    logger = if @logger_path && !@logger_path.strip.empty?
+               Logger.new(@logger_path, 'daily')
+             else
+               Logger.new($stdout)
+             end
+    logger.level = Logger::DEBUG
+    logger
+  end
+
+  def using_managed_identity?
+    val = @managed_identity_client_id.to_s.strip
+    !val.empty?
+  end
+
+  def validate_base_config
+    # Validate required configs for Managed Identity
+    required = {
+      'kusto_endpoint' => @kusto_endpoint,
+      'database_name' => @database_name,
+      'table_name' => @table_name
+    }
+    check_required_configs(required, %w[kusto_endpoint database_name table_name])
+    # No further validation needed for SYSTEM or GUID
+  end
+
+  def validate_aad_config
+    # Validate required configs for AAD
+    required = aad_required_hash
+    check_required_configs(
+      required,
+      %w[client_app_id client_app_secret tenant_id kusto_endpoint database_name table_name]
+    )
+  end
+
+  def validate_workload_identity_config
+    # Validate required configs for Workload Identity
+    required = {
+      'workload_identity_client_id' => @workload_identity_client_id,
+      'workload_identity_tenant_id' => @workload_identity_tenant_id,
+      'kusto_endpoint' => @kusto_endpoint,
+      'database_name' => @database_name,
+      'table_name' => @table_name
+    }
+    check_required_configs(required, %w[client_app_id tenant_id kusto_endpoint database_name table_name])
+  end
+
+  def aad_required_hash
+    # Return required config hash for AAD
+    {
+      'client_app_id' => @client_app_id,
+      'client_app_secret' => @client_app_secret,
+      'tenant_id' => @tenant_id,
+      'kusto_endpoint' => @kusto_endpoint,
+      'database_name' => @database_name,
+      'table_name' => @table_name
+    }
+  end
+
+  def check_required_configs(required_configs, names)
+    # Check for missing or empty required configs
+    required_configs.each do |name, conf|
+      print_missing_parameter_message_and_raise(name) if conf.nil?
+    end
+    return unless required_configs.values.any? { |conf| conf.to_s.strip.empty? }
+
+    raise ArgumentError,
+          "Malformed configuration, the following arguments can not be null or empty. [#{names.join(', ')}]"
+  end
+
+  def validate_azure_cloud
+    # Validate that the selected Azure cloud is supported
+    return if @azure_clouds.key?(@azure_cloud)
+
+    raise ArgumentError,
+          "The specified Azure cloud #{@azure_cloud} is not supported. Supported clouds are: " \
+          "#{@azure_clouds.keys.join(', ')}."
+  end
+end

data/lib/fluent/plugin/ingester.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+# Ingester handles uploading data to Azure Blob Storage and sending ingestion messages to Azure Queue for Kusto ingestion.
+#
+# Responsibilities:
+# - Upload data to blob storage
+# - Prepare and send ingestion messages to queue
+# - Handle errors during upload and ingestion
+
+require 'uri'
+require 'json'
+require 'securerandom'
+require 'base64'
+require_relative 'client'
+require_relative 'kusto_error_handler'
+require 'logger'
+require 'net/http'
+class Ingester
+  # Use a class instance variable instead of a class variable for client cache
+  @client_cache = nil
+  class << self
+    attr_accessor :client_cache
+  end
+
+  def initialize(outconfiguration)
+    # Initialize Ingester with configuration and resources
+    @client = self.class.client(outconfiguration)
+    @resources = @client.resources
+    @logger = begin
+      outconfiguration.logger
+    rescue StandardError
+      Logger.new($stdout)
+    end
+  end
+
+  def self.client(outconfiguration)
+    # Cache and return a Client instance
+    self.client_cache ||= Client.new(outconfiguration)
+  end
+
+  def build_uri(container_sas_uri, name)
+    # Build a blob URI with SAS token
+    base_uri, sas_token = container_sas_uri.split('?', 2)
+    base_uri = base_uri.chomp('/')
+    "#{base_uri}/#{name}?#{sas_token}"
+  end
+
+  # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+  def upload_to_blob(blob_uri, raw_data, blob_name)
+    # Upload raw data to Azure Blob Storage
+    uri_str = build_uri(blob_uri, blob_name)
+    uri = URI.parse(uri_str)
+    blob_size = raw_data.bytesize
+    request = Net::HTTP::Put.new(uri)
+    request.body = raw_data
+    request['x-ms-blob-type'] = 'BlockBlob'
+    request['Content-Length'] = blob_size.to_s
+
+    response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == 'https') do |http|
+      http.request(request)
+    end
+
+    unless response.code.to_i.between?(200, 299)
+      begin
+        error_handler = KustoErrorHandler.new(response.body)
+        if error_handler.permanent_error?
+          @logger.error("Permanent error while uploading blob: #{error_handler.message}.Blob name: #{blob_name}")
+        end
+      rescue StandardError => e
+        @logger.error("Failed to parse error response with KustoErrorHandler: #{e.message}. Blob name: #{blob_name}")
+      end
+      raise "Blob upload failed: #{response.code} #{response.message} - #{response.body}"
+    end
+
+    [uri.to_s, blob_size]
+  end
+  # rubocop:enable Metrics/AbcSize, Metrics/MethodLength
+
+  # rubocop:disable Metrics/MethodLength
+  def prepare_ingestion_message2(db, table, data_uri, blob_size_bytes, identity_token, compression_enabled = true)
+    # Prepare the ingestion message for Azure Queue
+    additional_props = {
+      'authorizationContext' => identity_token,
+      'format' => 'multijson'
+    }
+    additional_props['CompressionType'] = 'gzip' if compression_enabled
+    {
+      'Id' => SecureRandom.uuid,
+      'BlobPath' => data_uri,
+      'RawDataSize' => blob_size_bytes,
+      'DatabaseName' => db,
+      'TableName' => table,
+      'RetainBlobOnSuccess' => true,
+      'FlushImmediately' => true,
+      'ReportLevel' => 2,  # Report both failures and successes
+      'ReportMethod' => 0, # Use Azure Queue for reporting
+      'AdditionalProperties' => additional_props
+    }.to_json
+  end
+  # rubocop:enable Metrics/MethodLength
+
+  # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+  def post_message_to_queue_http(queue_uri_with_sas, message)
+    # Post the ingestion message to Azure Queue
+    base_uri, sas_token = queue_uri_with_sas.split('?', 2)
+    base_uri = base_uri.chomp('/')
+    post_uri = URI("#{base_uri}/messages?#{sas_token}")
+    encoded_message = Base64.strict_encode64(message)
+    request = Net::HTTP::Post.new(post_uri)
+    request['Content-Type'] = 'application/xml'
+    request.body = "<QueueMessage><MessageText>#{encoded_message}</MessageText></QueueMessage>"
+    response = Net::HTTP.start(post_uri.hostname, post_uri.port, use_ssl: post_uri.scheme == 'https') do |http|
+      http.request(request)
+    end
+    {
+      code: response.code,
+      message: response.message,
+      body: response.body
+    }
+  end
+  # rubocop:enable Metrics/AbcSize, Metrics/MethodLength
+
+  def upload_data_to_blob_and_queue(raw_data, blob_name, db, table_name, compression_enabled = true)
+    # Upload data to blob and send ingestion message to queue
+    blob_uri, blob_size_bytes = upload_to_blob(@resources[:blob_sas_uri], raw_data, blob_name)
+    message = prepare_ingestion_message2(db, table_name, blob_uri, blob_size_bytes, @resources[:identity_token],
+                                         compression_enabled)
+    post_message_to_queue_http(@resources[:queue_sas_uri], message)
+    { blob_uri: blob_uri, blob_size_bytes: blob_size_bytes }
+  end
+
+  def token_provider
+    # Return the token provider from the client
+    @client.token_provider
+  end
+end

data/lib/fluent/plugin/kusto_error_handler.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+# KustoErrorHandler parses and classifies errors returned from Azure Data Explorer (Kusto) API responses.
+# Provides methods for error extraction, classification, and logging.
+class KustoErrorHandler
+  attr_reader :message
+
+  def initialize(error_response)
+    # Parse error response and extract message
+    @error = parse_error(error_response)
+    @message = @error['Message'] || @error['message'] || error_response
+  end
+
+  def permanent_error?
+    # Check if error is marked as permanent
+    @error && [true, 'true'].include?(@error['@permanent'])
+  end
+
+  # rubocop:disable Metrics/MethodLength
+  # Extracts the Kusto error type (code) from an exception or error response
+  def self.extract_kusto_error_type(error)
+    error_json = parse_error_json(error)
+    return error_json['error']['code'] if error_json && error_json['error'] && error_json['error']['code']
+
+    nil
+  end
+
+  def self.parse_error_json(error)
+    # Parse error JSON from string or exception
+    if error.is_a?(String)
+      begin
+        JSON.parse(error)
+      rescue StandardError
+        nil
+      end
+    elsif error.respond_to?(:message)
+      begin
+        JSON.parse(error.message)
+      rescue StandardError
+        nil
+      end
+    end
+  end
+  # rubocop:enable Metrics/MethodLength
+
+  # Factory method to create a KustoErrorHandler from error type and message
+  def self.from_kusto_error_type(error_type, message)
+    # You can expand this logic to map error_type to more structured fields if needed
+    error_response = { 'error' => { 'code' => error_type, 'message' => message } }.to_json
+    new(error_response)
+  end
+
+  # Handles Kusto errors and logs them appropriately
+  def self.handle_kusto_error(logger, e, unique_id)
+    kusto_error_type = extract_kusto_error_type(e)
+    if kusto_error_type
+      kusto_error = from_kusto_error_type(kusto_error_type, e.message)
+      log_kusto_data_error(logger, kusto_error)
+      log_kusto_drop_chunk(logger, kusto_error, unique_id) if kusto_error.is_permanent?
+      # Always raise the custom error if present
+      raise kusto_error if kusto_error.is_a?(StandardError)
+      raise kusto_error unless kusto_error.is_permanent?
+
+      nil
+    else
+      log_failed_ingest(logger, unique_id, e)
+      raise
+    end
+  end
+
+  # Handles errors during try_write and logs them appropriately
+  def self.handle_try_write_error(logger, e, chunk_id)
+    kusto_error_type = extract_kusto_error_type(e)
+    if kusto_error_type
+      kusto_error = from_kusto_error_type(kusto_error_type, e.message)
+      log_kusto_data_error(logger, kusto_error)
+      if kusto_error.is_permanent?
+        log_kusto_drop_chunk(logger, kusto_error, chunk_id)
+        return nil
+      end
+      # Always raise the custom error if present
+      raise kusto_error if kusto_error.is_a?(StandardError)
+      raise kusto_error unless kusto_error.is_permanent?
+
+      nil
+    else
+      logger.error(
+        "Failed to ingest chunk #{chunk_id}: #{e.full_message}"
+      )
+      raise
+    end
+  end
+
+  # Log details of a Kusto data error
+  def self.log_kusto_data_error(logger, kusto_error)
+    logger.error(
+      "KustoDataError: #{kusto_error.message} " \
+      "(Code: #{kusto_error.failure_code}, Reason: #{kusto_error.failure_sub_code}, " \
+      "Permanent: #{kusto_error.is_permanent?})"
+    )
+  end
+
+  # Log when a chunk is dropped due to a permanent error
+  def self.log_kusto_drop_chunk(logger, kusto_error, chunk_id)
+    logger.error(
+      "Dropping chunk #{chunk_id} due to permanent Kusto error: #{kusto_error.message}"
+    )
+  end
+
+  # Log failed ingestion event
+  def self.log_failed_ingest(logger, unique_id, e)
+    logger.error(
+      "Failed to ingest event to Kusto : #{unique_id}\n" \
+      "#{e.full_message}"
+    )
+  end
+
+  private
+
+  def parse_error(error_response)
+    # Parse error response JSON
+    JSON.parse(error_response)
+  rescue StandardError
+    {}
+  end
+end

data/lib/fluent/plugin/kusto_query.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+# Provides helper functions for interacting with Kusto ingestion endpoints and running Kusto API queries.
+# Includes endpoint transformation and query execution logic.
+
+require 'net/http'
+require 'uri'
+require 'json'
+require 'securerandom'
+require 'base64'
+
+def to_ingest_endpoint(data_endpoint)
+  # Convert a Kusto data endpoint to its corresponding ingest endpoint
+  data_endpoint.sub(%r{^https://}, 'https://ingest-')
+end
+
+# Runs a Kusto API query against the specified endpoint.
+# Handles both management and query endpoints, builds request, and parses response.
+def run_kusto_api_query(query, data_endpoint, token_provider, use_ingest_endpoint: false, database_name: nil)
+  access_token = token_provider.get_token
+  endpoint = use_ingest_endpoint ? to_ingest_endpoint(data_endpoint) : data_endpoint
+  path = use_ingest_endpoint ? '/v1/rest/mgmt' : '/v1/rest/query'
+  uri = URI("#{endpoint}#{path}")
+
+  http = Net::HTTP.new(uri.host, uri.port)
+  http.use_ssl = true
+
+  headers = {
+    'Authorization' => "Bearer #{access_token}",
+    'Content-Type' => 'application/json',
+    'Accept' => 'application/json',
+    'x-ms-client-version' => 'Kusto.FluentD:1.0.0'
+  }
+
+  body_hash = { csl: query }
+  body_hash[:db] = database_name if database_name
+  body = body_hash.to_json
+
+  request = Net::HTTP::Post.new(uri.request_uri, headers)
+  request.body = body
+
+  response = http.request(request)
+  unless response.code.to_i.between?(200, 299)
+    # Print error details if query fails
+    puts "Kusto query failed with status #{response.code}:"
+    puts response.body
+    begin
+      error_handler = defined?(KustoErrorHandler) ? KustoErrorHandler.new(response.body) : nil
+      puts "Permanent Kusto error: #{error_handler.message}" if error_handler&.permanent_error?
+    rescue StandardError => e
+      puts "Failed to parse error response with KustoErrorHandler: #{e.message}"
+    end
+    return response
+  end
+
+  begin
+    # Parse and return rows from response JSON
+    response_json = JSON.parse(response.body)
+    tables = response_json['Tables']
+    rows = tables && tables[0] && tables[0]['Rows']
+    rows || []
+  rescue JSON::ParserError => e
+    puts "Failed to parse JSON: #{e}"
+    puts response.body
+    response
+  end
+end