data_sources/connect_doc_db.py

from pymongo import MongoClient
import os
from utils import TEMP_DIR
from pymongo_schema.extract import extract_pymongo_client_schema

def connect_doc_db(connection_string, nosql_db_name, session_hash):
    try:
        # Create a MongoClient object
        client = MongoClient(connection_string)
        print("Connected to NoSQL Mongo DB")
    
        # Access a database
        db = client[nosql_db_name]

        collection_names = db.list_collection_names()

        print(collection_names)

        schema = extract_pymongo_client_schema(client)
    
        # Close the connection
        if 'client' in locals() and client:
            client.close()
            print("MongoDB Connection closed.")

        session_path = 'doc_db'

        dir_path = TEMP_DIR / str(session_hash) / str(session_path)
        os.makedirs(dir_path, exist_ok=True)

        return ["success","<p style='color:green;text-align:center;font-size:18px;'>Document database connected successful</p>", collection_names, schema]
    except Exception as e:
        print("DocDB CONNECTION ERROR")
        print(e)
        return ["error",f"<p style='color:red;text-align:center;font-size:18px;font-weight:bold;'>ERROR: {e}</p>"]