diff --git a/examples/tutorial_authorization.py b/examples/tutorial_authorization.py index 9d9a800f..07aa84a5 100644 --- a/examples/tutorial_authorization.py +++ b/examples/tutorial_authorization.py @@ -17,6 +17,19 @@ def main(token='my-token'): pass +def auth_per_query(token="my-token"): + """Set headers per query, most prevelantly for authorization headers.""" + client = InfluxDBClient(username=None, password=None, headers=None) + + print(f"Use authorization token {token}") + + print(f"""Response for query with token: + {client.query( + f'SHOW DATABASES', + headers={"Authorization": token})}""") + pass + + def parse_args(): """Parse the args from main.""" parser = argparse.ArgumentParser( @@ -30,3 +43,4 @@ def parse_args(): if __name__ == '__main__': args = parse_args() main(token=args.token) + auth_per_query(token=args.token) diff --git a/influxdb/client.py b/influxdb/client.py index 51a64ac3..02bd7d41 100644 --- a/influxdb/client.py +++ b/influxdb/client.py @@ -444,7 +444,8 @@ def query(self, raise_errors=True, chunked=False, chunk_size=0, - method="GET"): + method="GET", + headers=None): """Send a query to InfluxDB. .. danger:: @@ -524,7 +525,8 @@ def query(self, params=params, data=None, stream=chunked, - expected_response_code=expected_response_code + expected_response_code=expected_response_code, + headers=headers ) data = response._msgpack diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py index e511ca9b..6fa888b1 100644 --- a/influxdb/tests/client_test.py +++ b/influxdb/tests/client_test.py @@ -246,14 +246,15 @@ def test_write_gzip(self): b"cpu_load_short,host=server01,region=us-west " b"value=0.64 1257894000000000000\n" ) + compressed.seek(0) self.assertEqual( - m.last_request.body, - compressed.getvalue(), + gzip.GzipFile(fileobj=io.BytesIO(m.last_request.body)).read(), + gzip.GzipFile(fileobj=compressed).read() ) def test_write_points_gzip(self): - """Test write points for TestInfluxDBClient object.""" + """Test write points for TestInfluxDBClient object with gzip.""" with requests_mock.Mocker() as m: m.register_uri( requests_mock.POST, @@ -276,9 +277,11 @@ def test_write_points_gzip(self): b'cpu_load_short,host=server01,region=us-west ' b'value=0.64 1257894000123456000\n' ) + compressed.seek(0) + self.assertEqual( - m.last_request.body, - compressed.getvalue(), + gzip.GzipFile(fileobj=io.BytesIO(m.last_request.body)).read(), + gzip.GzipFile(fileobj=compressed).read() ) def test_write_points_toplevel_attributes(self): @@ -1498,6 +1501,52 @@ def test_auth_token(self): self.assertEqual(m.last_request.headers["Authorization"], "my-token") + def test_query_with_auth_token(self): + """Test query with custom authorization header.""" + with requests_mock.Mocker() as m: + m.register_uri( + requests_mock.GET, + "http://localhost:8086/query", + status_code=200, + text="{}", + headers={"X-Influxdb-Version": "1.2.3"} + ) + cli = InfluxDBClient(username=None, password=None, headers=None) + cli.query("SELECT * FROM foo") + self.assertEqual(m.last_request.headers.get("Authorization"), + None) + + cli.query("SELET * FROM foo", + headers={"Authorization": "my-token"}) + self.assertEqual(m.last_request.headers.get("Authorization"), + "my-token") + + def test_query_header_overwrites_client_header(self): + """Custom query authorization header must overwrite init headers.""" + with requests_mock.Mocker() as m: + m.register_uri( + requests_mock.GET, + "http://localhost:8086/query", + status_code=200, + text="{}", + headers={"X-Influxdb-Version": "1.2.3"} + ) + cli = InfluxDBClient(username=None, password=None, headers={ + "Authorization": "client-token", + "header-to-drop": "not-only-is-Authorization-overwritten", + "header-to-drop2": "headers-of-client-are-all-overwritten"}) + cli.query("SELECT * FROM foo") + self.assertEqual(m.last_request.headers.get("Authorization"), + "client-token") + + cli.query("SELET * FROM foo", + headers={"Authorization": "query-token"}) + + self.assertEqual(m.last_request.headers.get("Authorization"), + "query-token") + self.assertFalse("header-to-drop" in m.last_request.headers) + self.assertFalse("header-to-drop2" in m.last_request.headers) + class FakeClient(InfluxDBClient): """Set up a fake client instance of InfluxDBClient."""