File size: 1,431 Bytes
6b22e52
 
 
 
 
 
 
 
 
 
b6a7027
6b22e52
18ea353
445008a
6b22e52
 
a803993
 
 
 
 
 
 
 
6b22e52
 
 
b7ee991
6b22e52
 
 
 
 
 
 
 
 
 
430e1ed
b6a7027
 
6b22e52
 
 
 
 
 
 
 
 
 
 
 
d1a0183
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
""" Backend server for the frontend app 
This file contains the endpoints that can be called via HTTP
"""

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
import duckdb
import pyarrow as pa
from uvicorn import run
from fastapi.staticfiles import StaticFiles

from fire import Fire

app = FastAPI()

# origins = ["*"]
# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=origins,
#     allow_credentials=True,
#     allow_methods=["*"],
#     allow_headers=["*"],
# )

# setup global connection to the database with a table
con = duckdb.connect()
con.query("""CREATE TABLE flights AS FROM 'flights-10m.parquet'""")


@app.get("/query/{sql_query:path}")
async def query(sql_query: str):
    global con
    sql_query = sql_query.replace("count(*)", "count(*)::INT")
    result = con.query(sql_query).arrow()
    return Response(arrow_to_bytes(result), media_type="application/octet-stream")


app.mount("/", StaticFiles(directory="dist", html=True), name="static")


def arrow_to_bytes(table: pa.Table):
    sink = pa.BufferOutputStream()
    with pa.RecordBatchStreamWriter(sink, table.schema) as writer:
        writer.write_table(table)
    bytes = sink.getvalue().to_pybytes()
    return bytes


def serve(port=8000, host="localhost"):
    run(app, port=port, host=host)


if __name__ == "__main__":
    Fire(serve)  # so I can run cli args with it