Merge pull request #58 from Flip-Liquid/flip-liquid/http_headers

Add support for arbitrary http headers

Author: hatarist

clickhouse_cli/cli.py

@@ -64,6 +64,7 @@ def __init__(
         vi_mode,
         cookie,
         insecure,
+        headers=None,
     ):
         self.config = None
 
@@ -73,6 +74,7 @@ def __init__(
         self.password = password
         self.database = database
         self.cookie = cookie
+        self.headers = headers or {}
         self.settings = {k: v[0] for k, v in parse_qs(settings).items()}
         self.format = format
         self.format_stdin = format_stdin
@@ -108,6 +110,7 @@ def connect(self):
             self.conn_timeout_retry,
             self.conn_timeout_retry_delay,
             not self.insecure,
+            headers=self.headers,
         )
 
         self.echo.print("Connecting to {host}:{port}".format(host=self.host, port=self.port))
@@ -213,6 +216,12 @@ def load_config(self):
         config_settings.update(arg_settings)
         self.settings = config_settings
 
+        config_headers = {}
+        if self.config.has_section("headers"):
+            config_headers = dict(self.config.items("headers"))
+        config_headers.update(self.headers)
+        self.headers = config_headers
+
         self.echo.colors = self.highlight
 
     def run(self, query, data):
@@ -564,6 +573,12 @@ def progress_print(self, message, percents):
 @click.option("--database", "-d", help="Database")
 @click.option("--settings", "-s", help="Query string to be sent with every query")
 @click.option("--cookie", "-c", help="Cookie header to be sent with every query")
+@click.option(
+    "--header",
+    "-H",
+    multiple=True,
+    help="Additional HTTP header to send with every request (format: 'Name: Value'). Can be specified multiple times.",
+)
 @click.option("--query", "-q", help="Query to execute")
 @click.option(
     "--insecure",
@@ -593,6 +608,7 @@ def run_cli(
     stacktrace,
     vi_mode,
     cookie,
+    header,
     version,
     files,
     insecure,
@@ -608,6 +624,18 @@ def run_cli(
     elif password:
         password = click.prompt("Password", hide_input=True, show_default=False, type=str)
 
+    # Parse --header values from "Name: Value" format into a dict.
+    # Later entries with the same name win.
+    headers = {}
+    for h in header:
+        if ": " not in h:
+            raise click.BadParameter(
+                "Headers must be in 'Name: Value' format, got: {!r}".format(h),
+                param_hint="'--header' / '-H'",
+            )
+        name, _, value = h.partition(": ")
+        headers[name] = value
+
     data_input = ()
 
     # Read from STDIN if non-interactive mode
@@ -634,6 +662,7 @@ def run_cli(
         vi_mode,
         cookie,
         insecure,
+        headers=headers,
     )
     cli.run(query, data_input)
     return 0

clickhouse_cli/clickhouse-cli.rc.sample

@@ -92,3 +92,13 @@ conn_timeout_retry_delay = 0.5
 # You can place the server-side settings here!
 
 # max_memory_usage = 20000000000
+
+
+[headers]
+# Arbitrary HTTP headers to send with every request.
+# Each key/value pair becomes one header. These can also be supplied at
+# runtime with the --header / -H flag ("Name: Value" format); CLI flags
+# take precedence over values set here.
+
+# X-My-Custom-Header = some-value
+# Authorization = Bearer my-token

clickhouse_cli/clickhouse/client.py

@@ -72,12 +72,14 @@ def __init__(
         timeout_retry=0,
         timeout_retry_delay=0.0,
         verify=True,
+        headers=None,
     ):
         self.url = url
         self.user = user
         self.password = password or ""
         self.database = database
         self.cookie = cookie
+        self.headers = headers or {}
         self.session_id = str(uuid.uuid4())
         self.cli_settings = {}
         self.stacktrace = stacktrace
@@ -113,6 +115,8 @@ def _query(
         if self.cookie:
             headers["Cookie"] = self.cookie
 
+        headers.update(self.headers)
+
         response = None
 
         if not query.endswith("\n"):

tests/test_cli.py

@@ -1,10 +1,99 @@
 import pytest
+from unittest.mock import MagicMock, patch
+
+from click.testing import CliRunner
 
 from clickhouse_cli.cli import run_cli
+from clickhouse_cli.clickhouse.client import Client
 
 
 def test_main_help():
     # Call with the --help option as a basic sanity check.
     with pytest.raises(SystemExit) as exinfo:
         run_cli(["--help", ])
     assert exinfo.value.code == 0
+
+
+def test_custom_headers_sent_with_request():
+    """tests -H --header arguments on CLI"""
+    runner = CliRunner()
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.elapsed.total_seconds.return_value = 0.1
+    mock_response.text = ""
+
+    captured = {}
+
+    def fake_request(method, url, **kwargs):
+        captured["headers"] = kwargs.get("headers", {})
+        return mock_response
+
+    with patch("requests.Session.request", side_effect=fake_request):
+        runner.invoke(run_cli, [
+            "-h", "localhost",
+            "-p", "8123",
+            "-H", "X-My-Header: test-value",
+            "-H", "X-Another: hello",
+            "-q", "SELECT 1",
+        ])
+
+    assert captured["headers"].get("X-My-Header") == "test-value"
+    assert captured["headers"].get("X-Another") == "hello"
+
+
+def test_custom_headers_stored_on_client():
+    """tests headers suppliet for client object"""
+    client = Client(
+        url="http://localhost:8123/",
+        user="default",
+        password="",
+        database="default",
+        cookie=None,
+        headers={"X-Custom": "value"},
+    )
+    assert client.headers == {"X-Custom": "value"}
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.elapsed.total_seconds.return_value = 0.0
+    mock_response.text = ""
+
+    captured = {}
+
+    def fake_request(method, url, **kwargs):
+        captured["headers"] = kwargs.get("headers", {})
+        return mock_response
+
+    with patch("requests.Session.request", side_effect=fake_request):
+        client._query("GET", "SELECT 1", {}, fmt="Null", stream=False)
+
+    assert captured["headers"].get("X-Custom") == "value"
+
+
+def test_custom_headers_override_defaults():
+    """user supplied headers override defaults"""
+    client = Client(
+        url="http://localhost:8123/",
+        user="default",
+        password="",
+        database="default",
+        cookie=None,
+        headers={"User-Agent": "my-custom-agent"},
+    )
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.elapsed.total_seconds.return_value = 0.0
+    mock_response.text = ""
+
+    captured = {}
+
+    def fake_request(method, url, **kwargs):
+        captured["headers"] = kwargs.get("headers", {})
+        return mock_response
+
+    with patch("requests.Session.request", side_effect=fake_request):
+        client._query("GET", "SELECT 1", {}, fmt="Null", stream=False)
+
+    assert captured["headers"]["User-Agent"] == "my-custom-agent"