diff --git a/code/database.py b/code/database.py index 3a3f935..710cc55 100644 --- a/code/database.py +++ b/code/database.py @@ -7,6 +7,7 @@ import pandas as pd import psycopg2 from dotenv import load_dotenv from sqlalchemy import create_engine +from pymongo import MongoClient class Database(object): @@ -222,4 +223,10 @@ def get_untranslated(dt): 'id': int(df['id'][0]), 'content': str(df['content'][0]) } - return val_dict \ No newline at end of file + return val_dict + +def connect_mongo(ip_addr: str = '192.168.1.2', db_name: str = 'twitter', collection_name: str = 'tweet'): + client = MongoClient(ip_addr) + db = client[db_name] + coll = db[collection_name] + return coll