ClickHouse/tests/integration/helpers/http_server.py

88 lines
3.2 KiB
Python
Raw Normal View History

2019-02-22 10:55:12 +00:00
# -*- coding: utf-8 -*-
import argparse
import csv
import socket
2019-02-22 10:55:12 +00:00
import ssl
2020-10-02 16:54:07 +00:00
from http.server import BaseHTTPRequestHandler, HTTPServer
2019-02-22 10:55:12 +00:00
2020-01-11 09:50:41 +00:00
# Decorator used to see if authentication works for external dictionary who use a HTTP source.
def check_auth(fn):
def wrapper(req):
2019-09-27 03:15:03 +00:00
auth_header = req.headers.get('authorization', None)
api_key = req.headers.get('api-key', None)
if not auth_header or auth_header != 'Basic Zm9vOmJhcg==' or not api_key or api_key != 'secret':
req.send_response(401)
else:
fn(req)
return wrapper
def start_server(server_address, data_path, schema, cert_path, address_family):
2019-02-22 10:55:12 +00:00
class TSVHTTPHandler(BaseHTTPRequestHandler):
@check_auth
def do_GET(self):
self.__send_headers()
self.__send_data()
@check_auth
def do_POST(self):
ids = self.__read_and_decode_post_ids()
2020-10-02 16:54:07 +00:00
print("ids=", ids)
self.__send_headers()
self.__send_data(ids)
def __send_headers(self):
2019-02-22 10:55:12 +00:00
self.send_response(200)
self.send_header('Content-type', 'text/tsv')
self.end_headers()
def __send_data(self, only_ids=None):
2019-02-22 10:55:12 +00:00
with open(data_path, 'r') as fl:
reader = csv.reader(fl, delimiter='\t')
for row in reader:
if not only_ids or (row[0] in only_ids):
2020-10-02 16:54:07 +00:00
self.wfile.write(('\t'.join(row) + '\n').encode())
2019-02-22 10:55:12 +00:00
def __read_and_decode_post_ids(self):
data = self.__read_and_decode_post_data()
2020-10-02 16:54:07 +00:00
return [_f for _f in data.split() if _f]
2019-02-22 10:55:12 +00:00
def __read_and_decode_post_data(self):
transfer_encoding = self.headers.get("Transfer-encoding")
decoded = "";
if transfer_encoding == "chunked":
while True:
2020-10-02 16:54:07 +00:00
s = self.rfile.readline().decode()
chunk_length = int(s, 16)
if not chunk_length:
break
2020-10-02 16:54:07 +00:00
decoded += self.rfile.read(chunk_length).decode()
self.rfile.readline().decode()
else:
content_length = int(self.headers.get("Content-Length", 0))
2020-10-02 16:54:07 +00:00
decoded = self.rfile.read(content_length).decode()
return decoded
if address_family == "ipv6":
HTTPServer.address_family = socket.AF_INET6
2019-02-22 10:55:12 +00:00
httpd = HTTPServer(server_address, TSVHTTPHandler)
if schema == "https":
2019-02-22 10:55:12 +00:00
httpd.socket = ssl.wrap_socket(httpd.socket, certfile=cert_path, server_side=True)
httpd.serve_forever()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Simple HTTP server returns data from file")
parser.add_argument("--host", default="localhost")
parser.add_argument("--port", default=5555, type=int)
parser.add_argument("--data-path", required=True)
parser.add_argument("--schema", choices=("http", "https"), required=True)
2019-02-22 10:55:12 +00:00
parser.add_argument("--cert-path", default="./fake_cert.pem")
parser.add_argument('--address-family', choices=("ipv4", "ipv6"), default="ipv4")
2019-02-22 10:55:12 +00:00
args = parser.parse_args()
start_server((args.host, args.port), args.data_path, args.schema, args.cert_path, args.address_family)