def db(self, database: str) -> None: """ A property that returns the new using database being queried. :param database: the name of the AsyncIOMotorDatabase we want to use. """ self._db: AsyncIOMotorDatabase = AsyncIOMotorDatabase(self._client, database) self._db.with_options(**vars(self._use_consistency))
async def get_config(name: str, db: DataBase = Depends(get_db)): config_collection = db.get_collection("configs") config_res = await config_collection.find_one({"name": name}) if config_res is None: message = f"Configuration {name} not found" logger.error(message) raise HTTPException(status_code=404, detail=message) return ConfigRespSchema(config=config_res["config"])
async def add_config( config: ConfigDBSchema = Body(...), db: DataBase = Depends(get_db) ): """User Endpoint: Push a config for one, several or all of the drops""" # Create or replace config config_collection = db.get_collection("configs") await config_collection.delete_many({"name": config.name}) result = await config_collection.insert_one(jsonable_encoder(config)) return str(result.inserted_id)
def test_database(self): # Test that we can create a db directly, not just get on from # AsyncIOMotorClient. db = AsyncIOMotorDatabase(self.cx, 'motor_test') # Make sure we got the right DB and it can do an operation. self.assertEqual('motor_test', db.name) yield from db.test_collection.remove() yield from db.test_collection.insert({'_id': 1}) doc = yield from db.test_collection.find_one({'_id': 1}) self.assertEqual(1, doc['_id'])
async def test_database(self): # Test that we can create a db directly, not just get on from # AsyncIOMotorClient. db = AsyncIOMotorDatabase(self.cx, "motor_test") # Make sure we got the right DB and it can do an operation. self.assertEqual("motor_test", db.name) await db.test_collection.delete_many({}) await db.test_collection.insert_one({"_id": 1}) doc = await db.test_collection.find_one({"_id": 1}) self.assertEqual(1, doc["_id"])
async def get_firmware(name: str, db: DataBase = Depends(get_db)): firmwares_collection = db.get_collection("firmwares") firmware = await firmwares_collection.find_one({"name": name}) if firmware is None: logger.error(f"No Firmware named {name}") return HTTPException(404, detail=f"No Firmware named {name}") firmware_file = io.BytesIO(firmware["firmware"]) firmware_file.name = firmware["name"] return StreamingResponse(firmware_file, media_type="application/octet-stream")
async def add_firmware(file: bytes = File(...), name: str = Form(...), db: DataBase = Depends(get_db)): """User Endpoint: Push a firmware""" # Create or replace firmware config_collection = db.get_collection("firmwares") await config_collection.delete_many({"name": name}) result = await config_collection.insert_one({ "name": name, "firmware": bson.binary.Binary(file) }) return str(result.inserted_id)
async def add_events_data(events_struct: EventsSchema = Body(...), db: DataBase = Depends(get_db)): event_ids = [] collection = db.get_collection("events") for event in events_struct.data: if event.type is None or event.time is None: logger.warning( f"[drop={events_struct.drop}] Malformed event: " f"missing 'time' or 'type' field in event {event.dict()}") db_event = EventDBSchema.from_event(event, events_struct.drop) inserted_event = await collection.insert_one(jsonable_encoder(db_event) ) event_ids.append(str(inserted_event.inserted_id)) return len(event_ids)
async def get_versions(drop: str, db: DataBase = Depends(get_db)): collection = db.get_collection("drops") # Get the generic versions default_versions = await collection.find_one({"drop": DEFAULT_CONFIG_NAME}) if default_versions is None: default_versions = DEFAULT_CONFIG # Try to retrieve specific versions if they exist versions = await collection.find_one({"drop": drop}) if versions is None: versions = default_versions else: for field, value in versions.items(): versions[ field] = default_versions[field] if value is None else value return VersionsRespSchema(**versions)
async def run_arbitrary_pipeline( db: AsyncIOMotorDatabase, collection: str, pipeline_stages: str, limit: int = 5, ) -> list: collection = db.get_collection(collection) pipeline_stages = json.loads(pipeline_stages) if limit is not None: pipeline_stages.append({"$limit": limit}) pipeline_stages.append({"$project": {"_id": 0}}) documents = [] async for doc in collection.aggregate(pipeline_stages): documents.append(doc) return documents
async def set_versions(body: DropPostSchema = Body(...), db: DataBase = Depends(get_db)): # Over-quality # if body.config is not None: # config_collection = db.get_collection("configs") # result = await config_collection.find_one({"name": body.config}) # if result is None: # raise HTTPException( # status_code=428, detail=f"Referenced config {body.config} not found" # ) # if body.firmware is not None: # firmware_collection = db.get_collection("firmwares") # result = await firmware_collection.find_one({"name": body.firmware}) # if result is None: # raise HTTPException( # status_code=428, detail=f"Referenced firmware {body.firmware} not found" # ) drops = [] drop_collection = db.get_collection("drops") fields = {"config": body.config, "firmware": body.firmware} set_dict = {f: value for f, value in fields.items() if value is not None} if not set_dict: raise HTTPException(status_code=422, detail="Either config or firmware must be set") insert_dict = {f: value for f, value in fields.items() if value is None} for drop in body.target_drops: result = await drop_collection.update_many( {"drop": drop}, { "$set": set_dict, "$setOnInsert": { **insert_dict, "drop": drop } }, upsert=True, ) if result.raw_result["ok"]: drops.append(drop) return response_model( drops, f"{len(drops)} config/firmware versions set successfully.")
async def run_pipeline( db: AsyncIOMotorDatabase, collection: str, pipeline_stages: str, from_timestamp: datetime = None, limit: int = None, ) -> list: collection = db.get_collection(collection) pipeline_stages = json.loads(pipeline_stages) if from_timestamp is not None: pipeline_stages.append( {"$match": { "x": { "$gte": from_timestamp.isoformat() } }}) if limit is not None: pipeline_stages.append({"$sort": {"x": -1}}) pipeline_stages.append({"$limit": limit}) pipeline_stages.append({"$project": {"_id": 0}}) documents = [] async for doc in collection.aggregate(pipeline_stages): documents.append(doc) return documents
def __init__(self, database: AsyncIOMotorDatabase): self.event_collection = database.get_collection("events")
def __init__(self, database: AsyncIOMotorDatabase): self.feeds_collection = database.get_collection("feeds")