Update test.py

test.py

@@ -149,85 +149,135 @@ async def auth_google(request: Request):
     # print(len(data_new))
     return {"attachment_count":attachment_no,"attachment_content":data_new}
 
-@app.get("/auth1/google")
-async def auth_google(request: Request):
+@app.websocket("/ws")
+async def auth_google(websocket: WebSocket):
     #This code is basically the authorization code and this authorization code helps us to get the access token with the required scopes that we have set .
     #We require the gmail.readonly scopes that requires verification of our application and all.
-
-    filter_query = None
-    if state:
-        state_params = state.split('&')
-        for param in state_params:
-            if param.startswith('filter_query='):
-                filter_query = param.split('=')[1]
-
-
-    token_url = "https://accounts.google.com/o/oauth2/token"
-    print(code)
-    data = {
-        "code": code,
-        "client_id": GOOGLE_CLIENT_ID,
-        "client_secret": GOOGLE_CLIENT_SECRET,
-        "redirect_uri": GOOGLE_REDIRECT_URI,
-        "grant_type": "authorization_code",
-    }
-    response = requests.post(token_url, data=data)
-    access_token = response.json().get("access_token")
-    print(response.json())
-    user_info = requests.get("https://www.googleapis.com/oauth2/v1/userinfo", headers={"Authorization": f"Bearer {access_token}"})
-
-    gmail_url = "https://www.googleapis.com/gmail/v1/users/me/messages?q=subject:receipt&maxResults=2"
-    gmail_response = requests.get(gmail_url, headers={"Authorization": f"Bearer {access_token}"})
-    messages = gmail_response.json().get("messages", [])
-    print(messages)
-    print("Printing gmail response")
-    print(gmail_response.json())
-    attachment_no = 0
-    data_new = [[]]
-
-    for i,message in enumerate(messages) :
-
-
-        if message:
-            message_id = message.get("id")
-            print(message_id)
-            if message_id:
-                message_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}"
-                message_response = requests.get(message_url, headers={"Authorization": f"Bearer {access_token}"})
-                message_data = message_response.json()
-
-                # Check for parts in the message payload
-                if "payload" in message_data and "parts" in message_data["payload"]:
-                    for part in message_data["payload"]["parts"]:
-                        if "body" in part and "attachmentId" in part["body"]:
-                            attachment_id = part["body"]["attachmentId"]
-                            attachment_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}/attachments/{attachment_id}"
-                            attachment_response = requests.get(attachment_url, headers={"Authorization": f"Bearer {access_token}"})
-                            attachment_data = attachment_response.json()
-                            data = attachment_data.get("data")
-                            filename = part.get("filename", "untitled.txt")
-
-                            if data:
-                                data_new.append([data.encode("utf-8")])
-                                # Decode base64-encoded attachment data
-                                attachment_content = base64.urlsafe_b64decode(data.encode("UTF-8"))
-                                # print(attachment_content)
-                                print("yo")
-                                print(data_new)
-                                attachment_no+=1
+    # raw_body = await request.body()
+    # return {"data":Yo Yo"}
+    # data = await request.json()
+    # code = data.get("access_token")
+    code = await websocket.receive_json()
+    # code = raw_body.decode()
+    def get_messages():
+        print()
+        # print(code)
+        access_token = code
+        print("printing access_token")
+        print(access_token)
+        page_token = None
+        messages = []
+        jobs_query = "subject:receipt OR subject:receipts has:attachment"
+
+        while True:
+            # Construct Gmail API request with pageToken
+            print("into the gmail")
+            gmail_url = f"https://www.googleapis.com/gmail/v1/users/me/messages?q={jobs_query}"
+
+            if page_token:
+                gmail_url += f"&pageToken={page_token}"
+
+            gmail_response = requests.get(gmail_url, headers={"Authorization": f"Bearer {access_token}"})
+            print(gmail_response)
+            gmail_data = gmail_response.json()
+
+            # Check if there are messages in the response
+            if "messages" in gmail_data:
+                messages.extend(gmail_data["messages"])
+
+            # Check if there are more pages
+            if "nextPageToken" in gmail_data:
+                page_token = gmail_data["nextPageToken"]
+            else:
+                break  # No more pages, exit the loop
+        print("returning the messages")
+
+        unique_thread_ids = set()
+
+        filtered_data_list = []
+
+        for entry in messages:
+            thread_id = entry['threadId']
+            if thread_id not in unique_thread_ids:
+                unique_thread_ids.add(thread_id)
+                filtered_data_list.append(entry)
+
+
+        print(messages)
+        print(filtered_data_list)
+        return filtered_data_list
+    
+    async def event_generator():
+        print(code)
+        access_token = code
+        messages=get_messages()
+        print(len(messages))
+        await websocket.send(json.dumps({"total_messages": len(messages)}))
+
+
+        attachments = []
+        prev_data=""
+        data_new={}
+        attachment_no=0
+        batch_size = 5
+        prev_filename = None
+
+        for i,message in enumerate(messages) :
+            print(i)
+            print(message)
+
+            if message:
+                message_id = message.get("id")
+                thread_id = message.get("threadId")
+                print(message_id)
+                if message_id:
+                    message_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}"
+                    message_response = requests.get(message_url, headers={"Authorization": f"Bearer {access_token}"})
+                    message_data = message_response.json()
+
+                    # Check for parts in the message payload
+                    if "payload" in message_data and "parts" in message_data["payload"]:
+                        for part in message_data["payload"]["parts"]:
+                            if "body" in part and "attachmentId" in part["body"]:
+                                attachment_id = part["body"]["attachmentId"]
+                                print(attachment_id)
+                                attachment_url = f"https://www.googleapis.com/gmail/v1/users/me/messages/{message_id}/attachments/{attachment_id}"
+                                attachment_response = requests.get(attachment_url, headers={"Authorization": f"Bearer {access_token}"})
+                                attachment_data = attachment_response.json()
+                                data = attachment_data.get("data",{})
+                                filename = part.get("filename", "untitled.txt")
+
+                                # print("Print the data json response for that gmail")
+                                print(filename)
+                                
+                                
+                                # print(attachment_data)
+                                # json_str = json.dumps(attachment_data, indent=2)
+
+                                # with subprocess.Popen(["less"], stdin=subprocess.PIPE) as less_process:
+                                #     less_process.communicate(input=json_str.encode("utf-8"))
+                                if data:
+                                    data_new[filename]=str(data[:10])
+                                    # data_new.append({filename:data.encode("UTF-8")})
+                                    # Decode base64-encodattaed attachment data
+                                    # print(filename)s
+                                    attachment_content = base64.urlsafe_b64decode(data.encode("UTF-8"))
+                                    # print(attachment_content)
+                                    print("yo")
+                                    # print(data_new)
+                                    attachment_no+=1
+                                    await websocket.send_json({filename:data[:10]})
+
+                                    # yield f"data: {str(data_new)}\n\n"
+                                    
+                                    # data_new={}
 
 
-                                # Save the attachment to a file
-                                #Here add the function or class to save the attachments to Google Cloud
-                                # save_path = f"/Users/omkarmalpure/Desktop/Current Work/fastapi_gmail/attachments/{filename}"
-                                # with open(save_path, "wb") as file:
-                                #     file.write(attachment_content)
+                                
 
-                                    
-                                    
 
-    return {"user_info": user_info.json(), "query": filter_query , "data":attachment_no , "attachment_content":data_new , "attachment_count":attachment_no}
 
+        await websocket.send_text("Successfully sent all the Data !")
 @app.get("/token")
 async def get_token(token: str = Depends(oauth2_scheme)):
     return jwt.decode(token, GOOGLE_CLIENT_SECRET, algorithms=["HS256"])