Update controllers/ws_controller.py

controllers/ws_controller.py

@@ -10,12 +10,12 @@ from models.models import Message, Attachment
 from fastapi import WebSocket
 from services import utils as ut
 
-def get_company_type(company_name:str)->str:
+def get_company_type(company_name:str,websocket:WebSocket)->str:
     company_types_dict ={'ao yun': 'wines and spirit', 'ardbeg': 'wines and spirit', 'belvedere': 'wines and spirit', 'bodega numanthia': 'wines and spirit', 'chandon': 'wines and spirit', 'château cheval blanc': 'wines and spirit', "château d'yquem": 'wines and spirit', 'château galoupet': 'wines and spirit', 'cheval des andes': 'wines and spirit', 'clos19': 'wines and spirit', 'cloudy bay': 'wines and spirit', 'colgin cellars': 'wines and spirit', 'dom pérignon': 'wines and spirit', 'domaine des lambrays': 'wines and spirit', 'eminente': 'wines and spirit', 'glenmorangie': 'wines and spirit', 'hennessy': 'wines and spirit', 'joseph phelps': 'wines and spirit', 'krug': 'wines and spirit', 'mercier': 'wines and spirit', 'moët & chandon': 'wines and spirit', 'newton vineyard': 'wines and spirit', 'ruinart': 'wines and spirit', 'terrazas de los andes': 'wines and spirit', 'veuve clicquot': 'wines and spirit', 'volcan de mi tierra': 'wines and spirit', 'woodinville': 'wines and spirit' , 'berluti': 'Fashion & Leather Goods', 'celine': 'Fashion & Leather Goods', 'christian dior': 'Fashion & Leather Goods', 'emilio pucci': 'Fashion & Leather Goods', 'fendi': 'Fashion & Leather Goods', 'givenchy': 'Fashion & Leather Goods', 'kenzo': 'Fashion & Leather Goods', 'loewe': 'Fashion & Leather Goods', 'loro piana': 'Fashion & Leather Goods', 'louis vuitton': 'Fashion & Leather Goods', 'marc jacobs': 'Fashion & Leather Goods', 'moynat': 'Fashion & Leather Goods', 'patou': 'Fashion & Leather Goods', 'rimowa': 'Fashion & Leather Goods','acqua di parma': 'Perfumes & Cosmetics', 'benefit cosmetics': 'Perfumes & Cosmetics', 'cha ling': 'Perfumes & Cosmetics', 'fenty beauty by rihanna': 'Perfumes & Cosmetics', 'fresh': 'Perfumes & Cosmetics', 'givenchy parfums': 'Perfumes & Cosmetics', 'guerlain': 'Perfumes & Cosmetics', 'kenzo parfums': 'Perfumes & Cosmetics', 'kvd beauty': 'Perfumes & Cosmetics', 'loewe perfumes': 'Perfumes & Cosmetics', 'maison francis kurkdjian': 'Perfumes & Cosmetics', 'make up for ever': 'Perfumes & Cosmetics', 'officine universelle buly': 'Perfumes & Cosmetics', 'olehenriksen': 'Perfumes & Cosmetics', 'parfums christian dior': 'Perfumes & Cosmetics', 'stella by stella mccartney': 'Perfumes & Cosmetics','bulgari': 'Watches & Jewelry', 'chaumet': 'Watches & Jewelry', 'fred': 'Watches & Jewelry', 'hublot': 'Watches & Jewelry', 'repossi': 'Watches & Jewelry', 'tag heuer': 'Watches & Jewelry', 'tiffany & co.': 'Watches & Jewelry', 'zenith': 'Watches & Jewelry','24s': 'Selective retailing', 'dfs': 'Selective retailing', 'la grande epicerie de paris': 'Selective retailing', 'le bon marché rive gauche': 'Selective retailing', 'sephora': 'Selective retailing','belmond': 'Other activities', 'cheval blanc': 'Other activities', 'connaissance des arts': 'Other activities', 'cova': 'Other activities', 'investir': 'Other activities', "jardin d'acclimatation": 'Other activities', 'le parisien': 'Other activities', 'les echos': 'Other activities', 'radio classique': 'Other activities', 'royal van lent': 'Other activities'}
     print(company_types_dict["louis vuitton"])
     return company_types_dict.get(company_name.lower(), 'Others')
 
-def get_messages(code: str,brand_name: Optional[str] = None) -> List[Message]:
+def get_messages(code: str,brand_name: Optional[str] = None):
     access_token = code
     g_query = f'(subject:"your order" OR subject:receipts OR subject:receipt OR subject:aankoopbon OR subject:reçu OR subject:invoice OR subject:invoices OR category:purchases) has:attachment'
     if brand_name is not None:
@@ -27,10 +27,20 @@ def get_messages(code: str,brand_name: Optional[str] = None) -> List[Message]:
     # gmail_response = requests.get(gmail_url, headers={"Authorization": f"Bearer {access_token}"})
     # gmail_data = gmail_response.json()
     # messages.append(gmail_data['messages'])
-    def fetch_message_wrapper(message_data):
+    def fetch_message_wrapper(message_data,websocket:WebSocket):
         message_id = message_data.get("id")
         if message_id:
-            return fetch_message_data(access_token, message_id)
+            message=fetch_message_data(access_token, message_id)
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+            
+                # Run the asynchronous method within the event loop
+                result = loop.run_until_complete(process_message(message,websocket,10000))
+            
+                # Close the event loop
+                loop.close()
+            
+            
         return None
 
     while True:
@@ -45,12 +55,16 @@ def get_messages(code: str,brand_name: Optional[str] = None) -> List[Message]:
         
         if "messages" in gmail_data:
             with ThreadPoolExecutor(max_workers=15) as executor:
-                futures = [executor.submit(fetch_message_wrapper, message_data) for message_data in
-                           gmail_data["messages"]]
-                for future in futures:
-                    message = future.result()
-                    if message:
-                        messages.append(message)
+                executor.submit(fetch_message_wrapper, message_data , websocket) for message_data in
+                           gmail_data["messages"]
+                
+                # futures=executor.submit(fetch_message_wrapper, message_data) for message_data in
+                #            gmail_data["messages"]
+                # for future in futures:
+                #     message = future.result()
+                #     if message:
+                #         messages.append(message)
+                        
 
         if "nextPageToken" in gmail_data:
             page_token = gmail_data["nextPageToken"]
@@ -60,6 +74,12 @@ def get_messages(code: str,brand_name: Optional[str] = None) -> List[Message]:
     print(messages)
     return messages
 
+async def process_message(message:Message, websocket:WebSocket, chunk_size:int):
+    if message:
+        message_json = message.to_json()
+        await send_message_in_chunks(websocket, message_json, chunk_size)
+        await websocket.send_text("NEXT_MESSAGE")
+
 
 def fetch_message_data(access_token: str, message_id: str) -> Message:
     message_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}"
@@ -204,22 +224,22 @@ def extract_text(html_content: str) -> str:
 
 async def websocket_main(code: str,  websocket: WebSocket,brand_name: Optional[str] = None):
     access_token = code
-    messages = get_messages(access_token,brand_name)
-    print("websocket_main")
-    print(messages)
-    # logging.info(f"brand_name:{brand_name}")
-    await websocket.send_json({"total_messages": len(messages)})
-    print("Total Length of messages")
-    print(len(messages))
-    chunk_size = 100000
-    i=0
-    for message in messages:
-        message_json = message.to_json()
+    messages = get_messages(access_token,brand_name , websocket)
+    # print("websocket_main")
+    # print(messages)
+    # # logging.info(f"brand_name:{brand_name}")
+    # await websocket.send_json({"total_messages": len(messages)})
+    # print("Total Length of messages")
+    # print(len(messages))
+    # chunk_size = 100000
+    # i=0
+    # for message in messages:
+    #     message_json = message.to_json()
         
-        logging.info(f"{i} th message")
-        i=i+1
-        await send_message_in_chunks(websocket, message_json, chunk_size)
-        await websocket.send_text("NEXT_MESSAGE")
+    #     logging.info(f"{i} th message")
+    #     i=i+1
+    #     await send_message_in_chunks(websocket, message_json, chunk_size)
+    #     await websocket.send_text("NEXT_MESSAGE")
 
     
     await websocket.close()