123 lines
3.5 KiB
Python
123 lines
3.5 KiB
Python
"""Script to move all images from MongoDB to MinIO."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
from bson import ObjectId
|
|
from dotenv import load_dotenv
|
|
from pymongo.collection import Collection
|
|
|
|
from shared.data_store import connect as connect_minio
|
|
from shared.data_store import put
|
|
from shared.database import connect
|
|
from shared.database import VisualCommunication
|
|
from shared.utils import check_env
|
|
from shared.utils import setup_logging
|
|
|
|
|
|
def list_mongo_document_ids(
|
|
collection: Collection,
|
|
) -> list[ObjectId]:
|
|
"""Get list of all VisualCommunication documents in MongoDB."""
|
|
# prepare query
|
|
query = collection.find(
|
|
filter={},
|
|
projection={'_id': True},
|
|
)
|
|
# execute query
|
|
res_list = list(query)
|
|
logging.debug('got %s document(s)', len(res_list))
|
|
# convert result
|
|
id_list = [elem['_id'] for elem in res_list]
|
|
return id_list
|
|
|
|
|
|
def get_visual_communication(
|
|
collection: Collection,
|
|
doc_id: ObjectId,
|
|
) -> VisualCommunication:
|
|
"""Get image from MongoDB."""
|
|
# prepare query
|
|
query = collection.find(
|
|
filter={'_id': doc_id},
|
|
projection={'_id': False},
|
|
)
|
|
# execute query
|
|
res_list = list(query)
|
|
logging.debug('got %s document(s)', len(res_list))
|
|
# convert result
|
|
vis_com = VisualCommunication.model_validate(res_list[0])
|
|
return vis_com
|
|
|
|
|
|
def update_visual_communication(
|
|
collection: Collection,
|
|
vis_com: VisualCommunication,
|
|
) -> None:
|
|
"""Update VisualCommunication document in MongoDB."""
|
|
try:
|
|
collection.update_one(
|
|
filter={
|
|
'name': vis_com.name,
|
|
},
|
|
update={
|
|
'$set': vis_com.model_dump(),
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
logging.error('failed saving %s', vis_com.name)
|
|
logging.debug(exc)
|
|
raise exc
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# load in env file
|
|
env_path = Path(__file__).parent.parent / 'server.env'
|
|
assert env_path.exists()
|
|
load_dotenv(env_path)
|
|
# ensure env vars set
|
|
check_env()
|
|
# setup logging
|
|
setup_logging()
|
|
# connect to minIO
|
|
minio_client = connect_minio()
|
|
# connect to MongoDB
|
|
collection, db, client = connect()
|
|
# list documents in mongoDB
|
|
id_list = list_mongo_document_ids(collection)
|
|
for doc_id in id_list:
|
|
try:
|
|
# get image from MongoDB
|
|
vis_com = get_visual_communication(collection, doc_id)
|
|
except Exception as exc:
|
|
logging.debug(exc)
|
|
logging.error('failed getting image from document: %s', doc_id)
|
|
continue
|
|
try:
|
|
# save image to buffer
|
|
buffer = BytesIO()
|
|
vis_com.image.save(buffer, 'png') # type: ignore
|
|
# put buffer in minio
|
|
object_name = put(
|
|
client=minio_client,
|
|
buffer=buffer,
|
|
)
|
|
except Exception as exc:
|
|
logging.debug(exc)
|
|
logging.error('failed saving image to minio')
|
|
continue
|
|
# update visual communication in mongodb
|
|
try:
|
|
vis_com.image = None # type: ignore
|
|
vis_com.object_name = object_name
|
|
update_visual_communication(collection, vis_com)
|
|
except Exception as exc:
|
|
logging.debug(exc)
|
|
logging.error(
|
|
'failed updating visual communication %s', vis_com.name,
|
|
)
|
|
continue
|
|
logging.debug('updated visual communication %s', vis_com.name)
|