Gmail changes to auth/readme/remove cache

gmail/.env-template

@@ -1,3 +1,5 @@
 GMAIL_USER_ID=test@example.com
+GMAIL_SERVICE_ACCOUNT_INFO=
 GMAIL_SEARCH_LIMIT=5
+# Connector Authorization
 GMAIL_CONNECTOR_API_KEY=

gmail/README.md

@@ -1,31 +1,59 @@
-# Gmail Connector
+# Gmail Quick Start Connector
 
-Connects Cohere to Google Gmail. It uses the Gmail messages API to search
-for emails.
+Connects Cohere to Google Mail. It uses the Gmail Python SDK to search for emails.
 
-## Configuration
+## Authentication
 
-The following configuration variables should be set as environment variables, or put into a `.env` file
-to control the behaviour of this connector:
+This connector supports two types of authentication: Service Account and OAuth.
 
+For Service Account authentication, domain-wide delegation is required, and can only search one user's mail at a time.
+
+### Service Account
+
+Service Account authentication requires two environment variables:
+`GMAIL_SERVICE_ACCOUNT_INFO`: Containing the JSON of your Service Account's credentials.
+`GMAIL_USER_ID`: Containing the email of the user who's emails you would like to search.
+
+#### `GMAIL_SERVICE_ACCOUNT_INFO`
+
+The `GMAIL_SERVICE_ACCOUNT_INFO` variable should contain the JSON content of the service account credentials file. To get the credentials file, follow these steps:
+
+1. [Create a project in Google Cloud Console](https://cloud.google.com/resource-manager/docs/creating-managing-projects).
+2. [Create a service account](https://cloud.google.com/iam/docs/creating-managing-service-accounts) and [activate the Google Drive API](https://console.cloud.google.com/apis/api/drive.googleapis.com) in the Google Cloud Console. Make sure that the user(s) you want to search are permitted to use the service account.
+3. [Create a service account key](https://cloud.google.com/iam/docs/creating-managing-service-account-keys) and download the credentials file as JSON. The credentials file should look like this:
+
+```json
+{
+  "type": "service_account",
+  "project_id": "{project id}",
+  "private_key_id": "{private_key_id}",
+  "private_key": "{private_key}",
+  "client_email": "{client_email}",
+  "client_id": "{client_id}",
+  "auth_uri": "{auth_uri}",
+  "token_uri": "{token_uri}",
+  "auth_provider_x509_cert_url": "{auth_provider_x509_cert_url}",
+  "client_x509_cert_url": "{client_x509_cert_url}",
+  "universe_domain": "{universe_domain}"
+}
 ```
-GMAIL_USER_ID=scott@lightsonsoftware.com
-GMAIL_MAX_RESULTS=10
-```
 
-## Credentials
+4. Convert the JSON credentials to a string and save the result in the `GMAIL_SERVICE_ACCOUNT_INFO` environment variable.
+5. On the Service Accounts page on GCP, copy the client ID value for your newly created service account, you will then need a super administrator user account to access [API Controls](https://admin.google.com/ac/accountchooser?continue=https://admin.google.com/ac/owl) and click on `Manage Domain Wide Delegation` > `Add new` and paste the client ID from earlier, then add the `https://www.googleapis.com/auth/gmail.readonly` to the OAuth Scopes field. Finally, click Authorize.
+
+### OAuth
+
+When using OAuth for authentication, the connector does not require any additional environment variables. Instead, the OAuth flow should occur outside of the Connector and Cohere's API will forward the user's access token to this connector through the `Authorization` header.
 
-To use this connector, you will need to turn on the Gmail API in a Google Cloud project. There are a couple of steps to follow, please refer to the following [guide.](https://developers.google.com/gmail/api/quickstart/python)
+With OAuth the connector will be able to search any e-mails the user has access to.
 
-You will need to:
+#### `GMAIL_CONNECTOR_API_KEY`
 
-1. Enable the API
-2. Configure Oauth consent: use the user with the same email as the `GMAIL_USER_ID` environment variable
-3. Create credentials for your app, and download the credentials.json file, making sure that you've whitelisted your hosted server's address as the redirect uri for the credentials
+The `GMAIL_CONNECTOR_API_KEY` should contain an API key for the connector. This value must be present in the `Authorization` header for all requests to the connector.
 
-Once this is done, save the `credentials.json` file into this directory.
+#### Optional
 
-This connector will also require a `token.json` file that is generated automatically once you authorize the app. After you deploy it, head to your hosted connector server's `/ui` page. For example, `https://myconnector.com/ui`, and trigger an initial `/search` request. This will bring up the Google OAuth page, where you can login to the same email as step 2 earlier.
+Optionally, you can modify the `GMAIL_SEARCH_LIMIT` variable to change the number of maximum results obtained by a search query.
 
 ## Development
 

gmail/provider/app.py

@@ -14,6 +14,9 @@ def search(body):
     except UpstreamProviderError as error:
         logger.error(f"Upstream search error: {error.message}")
         abort(502, error.message)
+    except AssertionError as error:
+        logger.error(f"Gmail connector config error: {error}")
+        abort(502, f"Gmail connector config error: {error}")
     return {"results": data}, 200, {"X-Connector-Id": app.config.get("APP_ID")}
 
 

gmail/provider/client.py

@@ -2,66 +2,65 @@
 import os
 
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from functools import lru_cache
-from flask import current_app as app
+from flask import request, current_app as app
 from google.auth.transport.requests import Request
+from google.oauth2 import service_account
 from google.oauth2.credentials import Credentials
-from google_auth_oauthlib.flow import InstalledAppFlow
 from googleapiclient.discovery import build
 from googleapiclient.errors import HttpError
 
 from . import UpstreamProviderError
 
 logger = logging.getLogger(__name__)
 
-CACHE_SIZE = 256
-
-client = None
+AUTHORIZATION_HEADER = "Authorization"
+BEARER_PREFIX = "Bearer "
+DEFAULT_SEARCH_LIMIT = 5
+USER_ME = "me"
 
 
 class GoogleMailClient:
+    FORMAT = "full"
     SCOPES = [
         "https://www.googleapis.com/auth/gmail.readonly",
     ]
 
-    def __init__(self, user_id, search_limit):
+    def __init__(self, service_account_info, access_token, user_id, search_limit):
         self.user_id = user_id
         self.search_limit = search_limit
+        self.service = build(
+            "gmail",
+            "v1",
+            credentials=self._request_credentials(service_account_info, access_token),
+        )
 
-        credentials = None
-
-        # Handle Authentication
-        if os.path.exists("token.json"):
-            logger.debug("Found token.json file")
-            credentials = Credentials.from_authorized_user_file(
-                "token.json", self.SCOPES
+    def _request_credentials(self, service_account_info=None, access_token=None):
+        if service_account_info is not None:
+            logger.debug("Using Service Account credentials")
+            credentials = service_account.Credentials.from_service_account_info(
+                service_account_info, scopes=self.SCOPES
             )
-
-        # If there are no (valid) credentials available, let the user log in.
-        if not credentials or credentials.expired or not credentials.valid:
-            logger.debug("No valid credentials found")
-
-            if credentials and credentials.expired and credentials.refresh_token:
+            if credentials.expired or not credentials.valid:
                 credentials.refresh(Request())
-            else:
-                flow = InstalledAppFlow.from_client_secrets_file(
-                    "credentials.json", self.SCOPES
-                )
-                credentials = flow.run_local_server(port=0)
 
-            # Save the credentials for the next run
-            with open("token.json", "w") as token:
-                token.write(credentials.to_json())
+            # For Service Account auth, need to set user email for delegated access
+            credentials_delegated = credentials.with_subject(self.user_id)
 
-        self.service = build("gmail", "v1", credentials=credentials)
+            return credentials_delegated
+        elif access_token is not None:
+            logger.debug("Using Oauth credentials")
+            return Credentials(access_token)
+        else:
+            raise UpstreamProviderError(
+                "No Service Account or Oauth credentials provided"
+            )
 
     def _request(self, request):
         try:
             return request.execute()
         except HttpError as http_error:
             raise UpstreamProviderError(message=str(http_error)) from http_error
 
-    @lru_cache(maxsize=CACHE_SIZE)
     def search_mail(self, query):
         request = (
             self.service.users()
@@ -77,12 +76,11 @@ def search_mail(self, query):
 
         return search_results
 
-    @lru_cache(maxsize=CACHE_SIZE)
     def get_message(self, message_id):
         request = (
             self.service.users()
             .messages()
-            .get(format="full", userId=self.user_id, id=message_id)
+            .get(format=self.FORMAT, userId=self.user_id, id=message_id)
         )
 
         message = self._request(request)
@@ -110,11 +108,23 @@ def batch_get_messages(self, message_ids):
 
 
 def get_client():
-    global client
+    service_account_info = app.config.get("SERVICE_ACCOUNT_INFO", None)
+    access_token = get_access_token()
+    user_id = app.config.get("USER_ID")
+    search_limit = app.config.get("SEARCH_LIMIT", DEFAULT_SEARCH_LIMIT)
+
+    if service_account_info is None and access_token is None:
+        raise AssertionError("No service account or oauth credentials provided")
+
+    # Using Oauth, use "me" user for current authenticated user
+    if service_account_info is None and access_token is not None:
+        user_id = USER_ME
+
+    return GoogleMailClient(service_account_info, access_token, user_id, search_limit)
 
-    if client is None:
-        assert (user_id := app.config.get("USER_ID")), "GMAIL_USER_ID must be set"
-        search_limit = app.config.get("SEARCH_LIMIT", 5)
-        client = GoogleMailClient(user_id, search_limit)
 
-    return client
+def get_access_token():
+    authorization_header = request.headers.get(AUTHORIZATION_HEADER, "")
+    if authorization_header.startswith(BEARER_PREFIX):
+        return authorization_header.removeprefix(BEARER_PREFIX)
+    return None