From 8339c0079f72a8e00cd1edbe5c1466fdbb5f3b3d Mon Sep 17 00:00:00 2001
From: Sanjiban Sengupta <sanjiban.sg@gmail.com>
Date: Thu, 13 Jul 2023 03:10:43 +0530
Subject: [PATCH] feat: use env variables in prod

---
 README.md        |  6 ++++++
 api/app.py       | 25 +++++++++++++++++--------
 api/shareable.py | 35 ++++++++++++++---------------------
 3 files changed, 37 insertions(+), 29 deletions(-)

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()