feat: add support for headers to be passed with each emit - emitterV2

README.md

@@ -1,14 +1,20 @@
-python-logging-loki
+python-logging-loki-v2 [Based on: https://github.com/GreyZmeem/python-logging-loki.]
 ===================
 
-[![PyPI version](https://img.shields.io/pypi/v/python-logging-loki.svg)](https://pypi.org/project/python-logging-loki/)
+[![PyPI version](https://img.shields.io/pypi/v/python-logging-loki-v2.svg)](https://pypi.org/project/python-logging-loki-v2/)
 [![Python version](https://img.shields.io/badge/python-3.6%20%7C%203.7%20%7C%203.8-blue.svg)](https://www.python.org/)
 [![License](https://img.shields.io/pypi/l/python-logging-loki.svg)](https://opensource.org/licenses/MIT)
-[![Build Status](https://travis-ci.org/GreyZmeem/python-logging-loki.svg?branch=master)](https://travis-ci.org/GreyZmeem/python-logging-loki)
 
-Python logging handler for Loki.  
+[//]: # ([![Build Status](https://travis-ci.org/GreyZmeem/python-logging-loki.svg?branch=master)](https://travis-ci.org/GreyZmeem/python-logging-loki))
+
+Python logging handler for Loki.
 https://grafana.com/loki
 
+New
+===========
+0.4.0: support to headers (ability to pass tenants for multi tenant loki configuration)
+
+
 Installation
 ============
 ```bash

logging_loki/emitter.py

@@ -5,6 +5,7 @@
 import functools
 import logging
 import time
+
 from logging.config import ConvertingDict
 from typing import Any
 from typing import Dict
@@ -30,29 +31,31 @@ class LokiEmitter(abc.ABC):
     label_replace_with = const.label_replace_with
     session_class = requests.Session
 
-    def __init__(self, url: str, tags: Optional[dict] = None, auth: BasicAuth = None):
+    def __init__(self, url: str, tags: Optional[dict] = None, auth: BasicAuth = None, headers: Optional[dict] = None):
         """
         Create new Loki emitter.
 
         Arguments:
             url: Endpoint used to send log entries to Loki (e.g. `https://my-loki-instance/loki/api/v1/push`).
             tags: Default tags added to every log record.
             auth: Optional tuple with username and password for basic HTTP authentication.
-
+            headers: Optional dict with HTTP headers to send.
         """
         #: Tags that will be added to all records handled by this handler.
         self.tags = tags or {}
         #: Loki JSON push endpoint (e.g `http://127.0.0.1/loki/api/v1/push`)
         self.url = url
         #: Optional tuple with username and password for basic authentication.
         self.auth = auth
+        #: Optional headers for post request
+        self.headers = headers or {}
 
         self._session: Optional[requests.Session] = None
 
     def __call__(self, record: logging.LogRecord, line: str):
         """Send log record to Loki."""
         payload = self.build_payload(record, line)
-        resp = self.session.post(self.url, json=payload)
+        resp = self.session.post(self.url, json=payload, headers=self.headers)
         if resp.status_code != self.success_response_code:
             raise ValueError("Unexpected Loki API response status code: {0}".format(resp.status_code))
 
@@ -113,7 +116,7 @@ def build_payload(self, record: logging.LogRecord, line) -> dict:
         labels = self.build_labels(record)
         ts = rfc3339.format_microsecond(record.created)
         stream = {
-            "labels": labels,
+            "labels" : labels,
             "entries": [{"ts": ts, "line": line}],
         }
         return {"streams": [stream]}
@@ -141,3 +144,20 @@ def build_payload(self, record: logging.LogRecord, line) -> dict:
             "values": [[ts, line]],
         }
         return {"streams": [stream]}
+
+
+class LokiEmitterV2(LokiEmitterV1):
+    """
+    Emitter for Loki >= 0.4.0.
+    Enables passing additional headers to requests
+    """
+
+    def __init__(self, url: str, tags: Optional[dict] = None, auth: BasicAuth = None, headers: dict = None):
+        super().__init__(url, tags, auth, headers)
+
+    def __call__(self, record: logging.LogRecord, line: str):
+        """Send log record to Loki."""
+        payload = self.build_payload(record, line)
+        resp = self.session.post(self.url, json=payload, headers=self.headers)
+        if resp.status_code != self.success_response_code:
+            raise ValueError("Unexpected Loki API response status code: {0}".format(resp.status_code))

logging_loki/handlers.py

@@ -34,6 +34,7 @@ class LokiHandler(logging.Handler):
     emitters: Dict[str, Type[emitter.LokiEmitter]] = {
         "0": emitter.LokiEmitterV0,
         "1": emitter.LokiEmitterV1,
+        "2": emitter.LokiEmitterV2
     }
 
     def __init__(
@@ -42,6 +43,7 @@ def __init__(
         tags: Optional[dict] = None,
         auth: Optional[emitter.BasicAuth] = None,
         version: Optional[str] = None,
+        headers: Optional[dict] = None
     ):
         """
         Create new Loki logging handler.
@@ -67,7 +69,7 @@ def __init__(
         version = version or const.emitter_ver
         if version not in self.emitters:
             raise ValueError("Unknown emitter version: {0}".format(version))
-        self.emitter = self.emitters[version](url, tags, auth)
+        self.emitter = self.emitters[version](url, tags, auth, headers)
 
     def handleError(self, record):  # noqa: N802
         """Close emitter and let default handler take actions on error."""

setup.py

@@ -6,15 +6,15 @@
     long_description = fh.read()
 
 setuptools.setup(
-    name="python-logging-loki",
-    version="0.3.1",
-    description="Python logging handler for Grafana Loki.",
+    name="python-logging-loki-v2",
+    version="0.4.3",
+    description="Python logging handler for Grafana Loki, with support to headers.",
     long_description=long_description,
     long_description_content_type="text/markdown",
     license="MIT",
-    author="Andrey Maslov",
-    author_email="greyzmeem@gmail.com",
-    url="https://github.com/greyzmeem/python-logging-loki",
+    author="Roman Rapoport",
+    author_email="cryos10@gmail.com",
+    url="https://github.com/RomanR-dev/python-logging-loki",
     packages=setuptools.find_packages(exclude=("tests",)),
     python_requires=">=3.6",
     install_requires=["rfc3339>=6.1", "requests"],

tests/test_emitter_v2.py

@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+
+import logging
+
+from typing import Tuple
+from unittest.mock import MagicMock
+
+import pytest
+
+from logging_loki.emitter import LokiEmitterV2
+
+emitter_url: str = "https://example.net/loki/api/v1/push/"
+headers = {"X-Scope-OrgID": "some_tenant"}
+record_kwargs = {
+    "name"    : "test",
+    "level"   : logging.WARNING,
+    "fn"      : "",
+    "lno"     : "",
+    "msg"     : "Test",
+    "args"    : None,
+    "exc_info": None,
+}
+
+
+@pytest.fixture()
+def emitter_v2() -> Tuple[LokiEmitterV2, MagicMock]:
+    """Create v2 emitter with mocked http session."""
+    response = MagicMock()
+    response.status_code = LokiEmitterV2.success_response_code
+    session = MagicMock()
+    session().post = MagicMock(return_value=response)
+
+    instance = LokiEmitterV2(url=emitter_url, headers=headers)
+    instance.session_class = session
+
+    return instance, session
+
+
+@pytest.fixture()
+def emitter_v2_no_headers() -> Tuple[LokiEmitterV2, MagicMock]:
+    """Create v2 emitter with mocked http session."""
+    response = MagicMock()
+    response.status_code = LokiEmitterV2.success_response_code
+    session = MagicMock()
+    session().post = MagicMock(return_value=response)
+
+    instance = LokiEmitterV2(url=emitter_url)
+    instance.session_class = session
+
+    return instance, session
+
+
+def create_record(**kwargs) -> logging.LogRecord:
+    """Create test logging record."""
+    log = logging.Logger(__name__)
+    return log.makeRecord(**{**record_kwargs, **kwargs})
+
+
+def get_stream(session: MagicMock) -> dict:
+    """Return first stream item from json payload."""
+    kwargs = session().post.call_args[1]
+    streams = kwargs["json"]["streams"]
+    return streams[0]
+
+
+def get_request(session: MagicMock) -> dict:
+    kwargs = session().post.call_args[1]
+    return kwargs
+
+
+def test_record_sent_to_emitter_url(emitter_v2):
+    emitter, session = emitter_v2
+    emitter(create_record(), "")
+
+    got = session().post.call_args
+    assert got[0][0] == emitter_url
+
+
+def test_default_tags_added_to_payload(emitter_v2):
+    emitter, session = emitter_v2
+    emitter.tags = {"app": "emitter"}
+    emitter(create_record(), "")
+
+    stream = get_stream(session)
+    level = logging.getLevelName(record_kwargs["level"]).lower()
+    expected = {
+        emitter.level_tag : level,
+        emitter.logger_tag: record_kwargs["name"],
+        "app"             : "emitter",
+    }
+    assert stream["stream"] == expected
+
+
+def test_headers_added(emitter_v2):
+    emitter, session = emitter_v2
+    emitter.tags = {"app": "emitter"}
+    emitter(create_record(), "")
+
+    kwargs = get_request(session)
+    assert kwargs['headers']['X-Scope-OrgID'] == headers['X-Scope-OrgID']
+
+
+def test_no_headers_added(emitter_v2_no_headers):
+    emitter, session = emitter_v2_no_headers
+    emitter.tags = {"app": "emitter"}
+    emitter(create_record(), "")
+
+    kwargs = get_request(session)
+    assert kwargs['headers'] is not None and kwargs['headers'] == {}