Skip to content

Commit 382d86e

Browse files
shredddCopiloteecavanna
authored
Implement API endpoints related to fetching BERtron data (#54)
Initial support for API endpoints. Note that we aren't using the Pydantic LinkML model for BERTron yet because the DB expects things to be in GeoJSON. This is converted on ingest, but once the schema supports GeoJSON it should simplify things. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: eecavanna <134325062+eecavanna@users.noreply.github.com>
1 parent 189b4cf commit 382d86e

1 file changed

Lines changed: 270 additions & 16 deletions

File tree

‎src/server.py‎

Lines changed: 270 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1-
from fastapi import FastAPI, HTTPException
1+
from fastapi import FastAPI, HTTPException, Query
22
import uvicorn
33
from fastapi.responses import RedirectResponse
44
from pymongo import MongoClient
5+
from typing import Optional, Dict, Any
6+
from pydantic import BaseModel, Field
57

68
# Connect to MongoDB.
79
# TODO: Get these values from environment variables instead of hard-coding them.
@@ -23,23 +25,16 @@ def get_health():
2325
return {"web_server": "ok", "database": is_database_healthy}
2426

2527

26-
# TODO: Delete this endpoint once we're confident in our database connection.
27-
@app.get("/experimental/database_names")
28-
def get_database_names():
29-
r"""Get the names of all databases in the MongoDB server."""
30-
db_names = mongo_client.list_database_names()
31-
return {"database_names": db_names}
28+
@app.get("/bertron")
29+
def get_all_entities():
30+
r"""Get all documents from the entities collection."""
31+
db = mongo_client["bertron"]
3232

33+
# Check if the collection exists
34+
if "entities" not in db.list_collection_names():
35+
raise HTTPException(status_code=404, detail="Entities collection not found")
3336

34-
@app.get("/database/{db_name}/{collection_name}")
35-
def get_collection_data(db_name: str, collection_name: str):
36-
r"""Get all documents from a specified collection in a specified database."""
37-
db = mongo_client[db_name]
38-
# Check if the collection exists in the database
39-
if collection_name not in db.list_collection_names():
40-
raise HTTPException(status_code=404, detail="Collection not found")
41-
42-
collection = db[collection_name]
37+
collection = db["entities"]
4338
documents = list(collection.find({}))
4439

4540
# Remove the MongoDB '_id' field from each document for JSON serialization
@@ -49,5 +44,264 @@ def get_collection_data(db_name: str, collection_name: str):
4944
return {"documents": documents}
5045

5146

47+
class MongoDBQuery(BaseModel):
48+
filter: Dict[str, Any] = Field(default={}, description="MongoDB find query filter")
49+
projection: Optional[Dict[str, Any]] = Field(
50+
default=None, description="Fields to include or exclude"
51+
)
52+
skip: Optional[int] = Field(
53+
default=0, ge=0, description="Number of documents to skip"
54+
)
55+
limit: Optional[int] = Field(
56+
default=100, ge=1, le=1000, description="Maximum number of documents to return"
57+
)
58+
sort: Optional[Dict[str, int]] = Field(
59+
default=None, description="Sort criteria (1 for ascending, -1 for descending)"
60+
)
61+
62+
63+
@app.post("/bertron/find")
64+
def find_entities(query: MongoDBQuery):
65+
r"""Execute a MongoDB find operation on the entities collection with filter, projection, skip, limit, and sort options.
66+
67+
Example query body:
68+
{
69+
"filter": {"field": "value", "number_field": {"$gt": 100}},
70+
"projection": {"field1": 1, "field2": 1},
71+
"skip": 0,
72+
"limit": 100,
73+
"sort": {"field1": 1, "field2": -1}
74+
}
75+
"""
76+
db = mongo_client["bertron"]
77+
78+
# Check if the collection exists
79+
if "entities" not in db.list_collection_names():
80+
raise HTTPException(status_code=404, detail="Entities collection not found")
81+
82+
collection = db["entities"]
83+
84+
try:
85+
# Execute find with query parameters
86+
cursor = collection.find(filter=query.filter, projection=query.projection)
87+
88+
# Apply skip, limit, and sort if provided
89+
if query.sort:
90+
cursor = cursor.sort(list(query.sort.items()))
91+
if query.skip:
92+
cursor = cursor.skip(query.skip)
93+
if query.limit:
94+
cursor = cursor.limit(query.limit)
95+
96+
# Convert cursor to list and remove MongoDB _id
97+
documents = list(cursor)
98+
for doc in documents:
99+
doc.pop("_id", None)
100+
101+
return {"documents": documents, "count": len(documents)}
102+
103+
except Exception as e:
104+
raise HTTPException(status_code=400, detail=f"Query error: {str(e)}")
105+
106+
107+
@app.get("/bertron/geo/nearby")
108+
def find_nearby_entities(
109+
latitude: float = Query(
110+
..., ge=-90, le=90, description="Center latitude in degrees"
111+
),
112+
longitude: float = Query(
113+
..., ge=-180, le=180, description="Center longitude in degrees"
114+
),
115+
radius_meters: float = Query(..., gt=0, description="Search radius in meters"),
116+
):
117+
r"""Find entities within a specified radius of a geographic point using MongoDB's $near operator.
118+
119+
This endpoint uses MongoDB's geospatial $near query which requires a 2dsphere index
120+
on the coordinates field for optimal performance.
121+
122+
Example: /bertron/geo/nearby?latitude=47.6062&longitude=-122.3321&radius_meters=10000
123+
"""
124+
db = mongo_client["bertron"]
125+
126+
# Check if the collection exists
127+
if "entities" not in db.list_collection_names():
128+
raise HTTPException(status_code=404, detail="Entities collection not found")
129+
130+
collection = db["entities"]
131+
132+
try:
133+
# Build the $near geospatial query
134+
geo_filter = {
135+
"coordinates": {
136+
"$near": {
137+
"$geometry": {
138+
"type": "Point",
139+
"coordinates": [
140+
longitude,
141+
latitude,
142+
], # MongoDB uses [lng, lat] format
143+
},
144+
"$maxDistance": radius_meters,
145+
}
146+
}
147+
}
148+
149+
# Execute find with geospatial filter and fixed projection
150+
cursor = collection.find(
151+
filter=geo_filter,
152+
projection={
153+
"id": 1,
154+
"name": 1,
155+
"uri": 1,
156+
"ber_data_source": 1,
157+
"coordinates": 1,
158+
},
159+
)
160+
161+
# Convert cursor to list and remove MongoDB _id
162+
documents = list(cursor)
163+
for doc in documents:
164+
doc.pop("_id", None)
165+
166+
return {
167+
"documents": documents,
168+
"count": len(documents),
169+
"query_type": "nearby",
170+
"center": {"latitude": latitude, "longitude": longitude},
171+
"radius_meters": radius_meters,
172+
}
173+
174+
except Exception as e:
175+
raise HTTPException(status_code=400, detail=f"Nearby query error: {str(e)}")
176+
177+
178+
@app.get("/bertron/geo/bbox")
179+
def find_entities_in_bounding_box(
180+
southwest_lat: float = Query(
181+
..., ge=-90, le=90, description="Southwest corner latitude"
182+
),
183+
southwest_lng: float = Query(
184+
..., ge=-180, le=180, description="Southwest corner longitude"
185+
),
186+
northeast_lat: float = Query(
187+
..., ge=-90, le=90, description="Northeast corner latitude"
188+
),
189+
northeast_lng: float = Query(
190+
..., ge=-180, le=180, description="Northeast corner longitude"
191+
),
192+
):
193+
r"""Find entities within a bounding box using MongoDB's $geoWithin operator.
194+
195+
This endpoint finds all entities whose coordinates fall within the specified
196+
rectangular bounding box defined by southwest and northeast corners.
197+
198+
Example: /bertron/geo/bbox?southwest_lat=47.5&southwest_lng=-122.4&northeast_lat=47.7&northeast_lng=-122.2
199+
"""
200+
db = mongo_client["bertron"]
201+
202+
# Check if the collection exists
203+
if "entities" not in db.list_collection_names():
204+
raise HTTPException(status_code=404, detail="Entities collection not found")
205+
206+
collection = db["entities"]
207+
208+
try:
209+
# Validate bounding box coordinates
210+
if southwest_lat >= northeast_lat:
211+
raise HTTPException(
212+
status_code=400,
213+
detail="Southwest latitude must be less than northeast latitude",
214+
)
215+
if southwest_lng >= northeast_lng:
216+
raise HTTPException(
217+
status_code=400,
218+
detail="Southwest longitude must be less than northeast longitude",
219+
)
220+
221+
# Build the $geoWithin bounding box query
222+
geo_filter = {
223+
"coordinates": {
224+
"$geoWithin": {
225+
"$box": [
226+
[
227+
southwest_lng,
228+
southwest_lat,
229+
], # MongoDB uses [lng, lat] format
230+
[northeast_lng, northeast_lat],
231+
]
232+
}
233+
}
234+
}
235+
236+
# Execute find with geospatial filter and fixed projection
237+
cursor = collection.find(
238+
filter=geo_filter,
239+
projection={
240+
"id": 1,
241+
"name": 1,
242+
"uri": 1,
243+
"ber_data_source": 1,
244+
"coordinates": 1,
245+
},
246+
)
247+
248+
# Convert cursor to list and remove MongoDB _id
249+
documents = list(cursor)
250+
for doc in documents:
251+
doc.pop("_id", None)
252+
253+
return {
254+
"documents": documents,
255+
"count": len(documents),
256+
"query_type": "bounding_box",
257+
"bounding_box": {
258+
"southwest": {"latitude": southwest_lat, "longitude": southwest_lng},
259+
"northeast": {"latitude": northeast_lat, "longitude": northeast_lng},
260+
},
261+
}
262+
263+
except Exception as e:
264+
raise HTTPException(
265+
status_code=400, detail=f"Bounding box query error: {str(e)}"
266+
)
267+
268+
269+
@app.get("/bertron/{id}")
270+
def get_entity_by_id(id: str):
271+
r"""Get a single entity by its ID.
272+
273+
Example: /bertron/emsl:12345
274+
"""
275+
db = mongo_client["bertron"]
276+
277+
# Check if the collection exists
278+
if "entities" not in db.list_collection_names():
279+
raise HTTPException(status_code=404, detail="Entities collection not found")
280+
281+
collection = db["entities"]
282+
283+
try:
284+
# Find the entity by ID with fixed projection
285+
document = collection.find_one(
286+
filter={"id": id},
287+
)
288+
289+
if not document:
290+
raise HTTPException(
291+
status_code=404, detail=f"Entity with id '{id}' not found"
292+
)
293+
294+
# Remove MongoDB _id
295+
document.pop("_id", None)
296+
297+
return document
298+
299+
except HTTPException:
300+
# Re-raise HTTP exceptions
301+
raise
302+
except Exception as e:
303+
raise HTTPException(status_code=400, detail=f"Query error: {str(e)}")
304+
305+
52306
if __name__ == "__main__":
53307
uvicorn.run(app, host="0.0.0.0", port=8000)

0 commit comments

Comments
 (0)