diff --git a/README.md b/README.md index db1fa3f0..93527477 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,13 @@ cd substrait-fiddle/ ``` Fiddle requires the [FastAPI back-end](https://github.com/voltrondata/substrait-fiddle/api) for APIs. Prior installation and operation of the service is required. ### API Service +To run it in PROD, use the complete URL, and set the environment variable. This step can be ignored in local development, +where the default user and port is used. + ``` +// for PROD +export PROD_MONGO_URL=url + cd api/ pip install -r requirements.txt uvicorn app:app --reload --port 9090 diff --git a/api/app.py b/api/app.py index 1ed3a267..99395deb 100644 --- a/api/app.py +++ b/api/app.py @@ -2,6 +2,8 @@ from fastapi.openapi.utils import get_openapi from fastapi_health import health +from motor.motor_asyncio import AsyncIOMotorCollection + import substrait_validator as sv from backend.duckdb import ( @@ -11,7 +13,7 @@ ParseFromDuckDB, ) -from shareable import MongoDBConnection, PlanData, get_mongo_conn +from shareable import MongoDBConnection, PlanData from loguru import logger @@ -19,6 +21,10 @@ duckConn = None +async def get_mongo_conn(): + return app.state.mongo_pool.initialize() + + @app.get("/") def ping(): return {"api_service": "up and running"} @@ -28,6 +34,7 @@ def ping(): async def Initialize(): global duckConn duckConn = ConnectDuckDB() + app.state.mongo_pool = MongoDBConnection() @app.get("/health/duckdb/") @@ -36,8 +43,8 @@ def CheckBackendConn(conn): @app.get("/health/mongodb/") -def CheckMongoConn(db: MongoDBConnection = Depends(get_mongo_conn)): - db.check() +def CheckMongoConn(): + app.state.mongo_pool.check() app.add_api_route("/health", health([CheckBackendConn])) @@ -87,14 +94,16 @@ async def ParseToSubstrait(data: dict): @app.post("/save/") -async def SavePlan(data: PlanData, db: MongoDBConnection = Depends(get_mongo_conn)): - response = await db.add_record(data) +async def SavePlan( + data: PlanData, db_conn: AsyncIOMotorCollection = Depends(get_mongo_conn) +): + response = await app.state.mongo_pool.add_record(db_conn, data) return response -@app.post("/fetch/", status_code=status.HTTP_200_OK) -async def FetchPlan(id: str, db: MongoDBConnection = Depends(get_mongo_conn)): - response = await db.get_record(id) +@app.post("/fetch/") +async def FetchPlan(id: str, db_conn: AsyncIOMotorCollection = Depends(get_mongo_conn)): + response = await app.state.mongo_pool.get_record(db_conn, id) if response is None: raise HTTPException(status_code=404, detail="Plan not found") return { diff --git a/api/shareable.py b/api/shareable.py index d78011b7..97a3924d 100644 --- a/api/shareable.py +++ b/api/shareable.py @@ -1,3 +1,4 @@ +import os from bson.objectid import ObjectId from pydantic import BaseModel from motor.motor_asyncio import AsyncIOMotorClient @@ -10,25 +11,30 @@ class PlanData(BaseModel): class MongoDBConnection: def __init__(self): - self.client = AsyncIOMotorClient("mongodb://localhost:27017") - self.database = self.client["plans"] - self.collection = self.database["links"] + url = os.environ.get("PROD_MONGO_URL", "mongodb://localhost:27017") + self.client = AsyncIOMotorClient(url) - async def get_record(self, id): - record = await self.collection.find_one({"_id": ObjectId(id)}) + def initialize(self): + database = self.client["plans"] + collection = database["links"] + return collection + + async def get_record(self, collection, id): + record = await collection.find_one({"_id": ObjectId(id)}) return record - async def add_record(self, data): + async def add_record(self, collection, data): data = { "json_data": data.json_string, "validator_overrides": data.validator_overrides, } - result = await self.collection.insert_one(data) + result = await collection.insert_one(data) return str(result.inserted_id) async def check(self): try: - client = AsyncIOMotorClient("mongodb://localhost:27017") + url = os.environ.get("PROD_MONGO_URL", "mongodb://localhost:27017") + client = AsyncIOMotorClient(url) except Exception as e: return False, str(e) @@ -41,16 +47,3 @@ async def check(self): False, "Links collection doesn't exist in Plan database.", ) - - async def close(self): - self.collection = None - self.database = None - self.client.close() - - -async def get_mongo_conn(): - mongo_conn = MongoDBConnection() - try: - yield mongo_conn - finally: - await mongo_conn.close()