File size: 5,917 Bytes
8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 c882e97 8fbb0f5 |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
import json
import sqlite3
from contextlib import asynccontextmanager
from fastapi import FastAPI, Query, HTTPException
from typing import List, Optional
from pydantic import BaseModel
from data_loader import refresh_data
import numpy as np
from pandas import Timestamp
import logging
logger = logging.getLogger(__name__)
def get_db_connection():
conn = sqlite3.connect("datasets.db")
conn.row_factory = sqlite3.Row
return conn
def setup_database():
conn = get_db_connection()
c = conn.cursor()
c.execute(
"""CREATE TABLE IF NOT EXISTS datasets
(hub_id TEXT PRIMARY KEY,
likes INTEGER,
downloads INTEGER,
tags TEXT,
created_at INTEGER,
last_modified INTEGER,
license TEXT,
language TEXT,
config_name TEXT,
column_names TEXT,
features TEXT)"""
)
c.execute("CREATE INDEX IF NOT EXISTS idx_column_names ON datasets (column_names)")
conn.commit()
conn.close()
def serialize_numpy(obj):
if isinstance(obj, np.ndarray):
return obj.tolist()
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.floating):
return float(obj)
if isinstance(obj, Timestamp):
return int(obj.timestamp())
logger.error(f"Object of type {type(obj)} is not JSON serializable")
raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
def insert_data(conn, data):
c = conn.cursor()
created_at = data.get("created_at", 0)
if isinstance(created_at, Timestamp):
created_at = int(created_at.timestamp())
last_modified = data.get("last_modified", 0)
if isinstance(last_modified, Timestamp):
last_modified = int(last_modified.timestamp())
c.execute(
"""
INSERT OR REPLACE INTO datasets
(hub_id, likes, downloads, tags, created_at, last_modified, license, language, config_name, column_names, features)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
data["hub_id"],
data.get("likes", 0),
data.get("downloads", 0),
json.dumps(data.get("tags", []), default=serialize_numpy),
created_at,
last_modified,
json.dumps(data.get("license", []), default=serialize_numpy),
json.dumps(data.get("language", []), default=serialize_numpy),
data.get("config_name", ""),
json.dumps(data.get("column_names", []), default=serialize_numpy),
json.dumps(data.get("features", []), default=serialize_numpy),
),
)
conn.commit()
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup: Load data into the database
setup_database()
logger.info("Creating database connection")
conn = get_db_connection()
logger.info("Refreshing data")
datasets = refresh_data()
for data in datasets:
insert_data(conn, data)
conn.close()
logger.info("Data refreshed")
yield
# Shutdown: You can add any cleanup operations here if needed
# For example, closing database connections, clearing caches, etc.
app = FastAPI(lifespan=lifespan)
class SearchResponse(BaseModel):
total: int
page: int
page_size: int
results: List[dict]
@app.get("/search", response_model=SearchResponse)
async def search_datasets(
columns: List[str] = Query(...),
match_all: bool = Query(False),
page: int = Query(1, ge=1),
page_size: int = Query(10, ge=1, le=1000),
):
offset = (page - 1) * page_size
conn = get_db_connection()
c = conn.cursor()
try:
if match_all:
query = """
SELECT COUNT(*) as total FROM datasets
WHERE (SELECT COUNT(*) FROM json_each(column_names)
WHERE value IN ({})) = ?
""".format(",".join("?" * len(columns)))
c.execute(query, (*columns, len(columns)))
else:
query = """
SELECT COUNT(*) as total FROM datasets
WHERE EXISTS (
SELECT 1 FROM json_each(column_names)
WHERE value IN ({})
)
""".format(",".join("?" * len(columns)))
c.execute(query, columns)
total = c.fetchone()["total"]
if match_all:
query = """
SELECT * FROM datasets
WHERE (SELECT COUNT(*) FROM json_each(column_names)
WHERE value IN ({})) = ?
LIMIT ? OFFSET ?
""".format(",".join("?" * len(columns)))
c.execute(query, (*columns, len(columns), page_size, offset))
else:
query = """
SELECT * FROM datasets
WHERE EXISTS (
SELECT 1 FROM json_each(column_names)
WHERE value IN ({})
)
LIMIT ? OFFSET ?
""".format(",".join("?" * len(columns)))
c.execute(query, (*columns, page_size, offset))
results = [dict(row) for row in c.fetchall()]
for result in results:
result["tags"] = json.loads(result["tags"])
result["license"] = json.loads(result["license"])
result["language"] = json.loads(result["language"])
result["column_names"] = json.loads(result["column_names"])
result["features"] = json.loads(result["features"])
return SearchResponse(
total=total, page=page, page_size=page_size, results=results
)
except sqlite3.Error as e:
logger.error(f"Database error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Database error: {str(e)}") from e
finally:
conn.close()
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|