Skip to content

Commit 30abb75

Browse files
committed
refine, resynch
1 parent 8ed8abc commit 30abb75

File tree

1 file changed

+50
-5
lines changed

1 file changed

+50
-5
lines changed

main.py

Lines changed: 50 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,18 @@
22
import duckdb
33
import json
44
import time
5+
import tempfile
6+
import hashlib
7+
import base64
8+
59
from flask import Flask, request, jsonify
610
from flask_httpauth import HTTPBasicAuth
711

812
app = Flask(__name__, static_folder="public", static_url_path="")
913
auth = HTTPBasicAuth()
1014

15+
dbpath = os.getenv('DBPATH', '/tmp/')
16+
1117
# Global connection
1218
conn = None
1319

@@ -18,9 +24,14 @@ def verify(username, password):
1824
print('stateless session')
1925
conn = duckdb.connect(':memory:')
2026
else:
21-
path = f"{globals()['path']}/{hash(username + password)}.db"
22-
print(f'stateful session {path}')
23-
conn = duckdb.connect(path)
27+
28+
global path
29+
os.makedirs(path, exist_ok=True)
30+
user_pass_hash = hashlib.sha256((username + password).encode()).hexdigest()
31+
db_file = os.path.join(dbpath, f"{user_pass_hash}.db")
32+
print(f'stateful session {db_file}')
33+
conn = duckdb.connect(db_file)
34+
2435
return True
2536

2637
def convert_to_ndjson(result):
@@ -54,8 +65,38 @@ def convert_to_clickhouse_jsoncompact(result, query_time):
5465

5566
return json.dumps(json_result)
5667

57-
def duckdb_query_with_errmsg(query, format):
68+
def handle_insert_query(query, format, data=None):
69+
table_name = query.split("INTO")[1].split()[0].strip()
70+
71+
temp_file_name = None
72+
if format.lower() == 'jsoneachrow' and data is not None:
73+
temp_file_name = save_to_tempfile(data)
74+
75+
if temp_file_name:
76+
try:
77+
ingest_query = f"COPY {table_name} FROM '{temp_file_name}' (FORMAT 'json')"
78+
conn.execute(ingest_query)
79+
except Exception as e:
80+
return b"", str(e).encode()
81+
finally:
82+
os.remove(temp_file_name)
83+
84+
return b"Ok", b""
85+
86+
def save_to_tempfile(data):
87+
temp_file = tempfile.NamedTemporaryFile(delete=False, mode='w+', encoding='utf-8')
88+
temp_file.write(data)
89+
temp_file.flush()
90+
temp_file.close()
91+
return temp_file.name
92+
93+
94+
def duckdb_query_with_errmsg(query, format, data=None, request_method="GET"):
5895
try:
96+
97+
if request_method == "POST" and query.strip().lower().startswith('insert into') and data:
98+
return handle_insert_query(query, format, data)
99+
59100
start_time = time.time()
60101
result = conn.execute(query)
61102
query_time = time.time() - start_time
@@ -83,18 +124,22 @@ def clickhouse():
83124
query = request.args.get('query', default="", type=str)
84125
format = request.args.get('default_format', default="JSONCompact", type=str)
85126
database = request.args.get('database', default="", type=str)
127+
data = None
86128

87129
# Log incoming request data for debugging
88130
print(f"Received request: method={request.method}, query={query}, format={format}, database={database}")
89131

90132
if not query:
91133
return app.send_static_file('play.html')
92134

135+
if request.method == "POST":
136+
data = request.get_data(as_text=True)
137+
93138
if database:
94139
query = f"ATTACH '{database}' AS db; USE db; {query}"
95140

96141
# Execute the query and capture the result and error message
97-
result, errmsg = duckdb_query_with_errmsg(query.strip(), format)
142+
result, errmsg = duckdb_query_with_errmsg(query.strip(), format, data, request.method)
98143

99144
# Handle response for HEAD requests
100145
if len(errmsg) == 0:

0 commit comments

Comments
 (0)