diff --git a/Sources/LoggingLoki/Batcher.swift b/Sources/LoggingLoki/Batcher.swift index 2cf593e..260c5ea 100644 --- a/Sources/LoggingLoki/Batcher.swift +++ b/Sources/LoggingLoki/Batcher.swift @@ -2,6 +2,7 @@ import Foundation class Batcher { private let session: LokiSession + private let headers: [String: String] private let lokiURL: URL private let sendDataAsJSON: Bool @@ -13,12 +14,14 @@ class Batcher { var batch: Batch? = nil - internal init(session: LokiSession, - lokiURL: URL, - sendDataAsJSON: Bool, - batchSize: Int, - maxBatchTimeInterval: TimeInterval?) { + init(session: LokiSession, + headers: [String: String], + lokiURL: URL, + sendDataAsJSON: Bool, + batchSize: Int, + maxBatchTimeInterval: TimeInterval?) { self.session = session + self.headers = headers self.lokiURL = lokiURL self.sendDataAsJSON = sendDataAsJSON self.batchSize = batchSize @@ -53,7 +56,7 @@ class Batcher { } private func sendBatch(_ batch: Batch) { - session.send(batch, url: lokiURL, sendAsJSON: sendDataAsJSON) { result in + session.send(batch, url: lokiURL, headers: headers, sendAsJSON: sendDataAsJSON) { result in if case .failure(let failure) = result { debugPrint(failure) } diff --git a/Sources/LoggingLoki/LokiLogHandler.swift b/Sources/LoggingLoki/LokiLogHandler.swift index 3513629..9891d06 100644 --- a/Sources/LoggingLoki/LokiLogHandler.swift +++ b/Sources/LoggingLoki/LokiLogHandler.swift @@ -25,6 +25,7 @@ public struct LokiLogHandler: LogHandler { /// This initializer is only used internally and for running Unit Tests. internal init(label: String, lokiURL: URL, + headers: [String: String] = [:], sendAsJSON: Bool = false, batchSize: Int = 10, maxBatchTimeInterval: TimeInterval? = 5 * 60, @@ -44,6 +45,7 @@ public struct LokiLogHandler: LogHandler { self.maxBatchTimeInterval = maxBatchTimeInterval self.session = session self.batcher = Batcher(session: self.session, + headers: headers, lokiURL: self.lokiURL, sendDataAsJSON: self.sendDataAsJSON, batchSize: self.batchSize, @@ -67,6 +69,7 @@ public struct LokiLogHandler: LogHandler { /// - Parameters: /// - label: client supplied string describing the logger. Should be unique but not enforced /// - lokiURL: client supplied Grafana Loki base URL + /// - headers: These headers will be added to all requests sent to Grafana Loki. /// - sendAsJSON: Indicates if the logs should be sent to Loki as JSON. /// This should not be required in most cases. By default this is false. /// Logs will instead be sent as snappy compressed protobuf, @@ -79,11 +82,13 @@ public struct LokiLogHandler: LogHandler { /// The option should prevent leaving logs in memory for too long without sending them. public init(label: String, lokiURL: URL, + headers: [String: String] = [:], sendAsJSON: Bool = false, batchSize: Int = 10, maxBatchTimeInterval: TimeInterval? = 5 * 60) { self.init(label: label, lokiURL: lokiURL, + headers: headers, sendAsJSON: sendAsJSON, batchSize: batchSize, session: URLSession(configuration: .ephemeral)) diff --git a/Sources/LoggingLoki/LokiSession.swift b/Sources/LoggingLoki/LokiSession.swift index 918de94..bf63313 100644 --- a/Sources/LoggingLoki/LokiSession.swift +++ b/Sources/LoggingLoki/LokiSession.swift @@ -8,6 +8,7 @@ import Snappy protocol LokiSession { func send(_ batch: Batch, url: URL, + headers: [String: String], sendAsJSON: Bool, completion: @escaping (Result) -> Void) } @@ -15,6 +16,7 @@ protocol LokiSession { extension URLSession: LokiSession { func send(_ batch: Batch, url: URL, + headers: [String: String], sendAsJSON: Bool = false, completion: @escaping (Result) -> Void) { do { @@ -49,6 +51,9 @@ extension URLSession: LokiSession { request.httpMethod = "POST" request.httpBody = data request.setValue(contentType, forHTTPHeaderField: "Content-Type") + for header in headers { + request.setValue(header.value, forHTTPHeaderField: header.key) + } let task = dataTask(with: request) { data, response, error in if let error = error { diff --git a/Tests/LoggingLokiTests/LoggingLokiTests.swift b/Tests/LoggingLokiTests/LoggingLokiTests.swift index 72b12ad..cf2b92c 100644 --- a/Tests/LoggingLokiTests/LoggingLokiTests.swift +++ b/Tests/LoggingLokiTests/LoggingLokiTests.swift @@ -8,6 +8,7 @@ class TestSession: LokiSession { func send(_ batch: Batch, url: URL, + headers: [String: String], sendAsJSON: Bool = false, completion: @escaping (Result) -> Void) { self.logs = batch.entries.first?.logEntries