From 39710f873a0510f9931486fc703f69c137f87db3 Mon Sep 17 00:00:00 2001
From: Mostafa Rashed <17770919+mrashed-dev@users.noreply.github.com>
Date: Fri, 3 May 2024 12:38:24 -0400
Subject: [PATCH] Add new fields; provider for code exchange and custom_headers
 for drafts/messages (#360)

This PR adds two new fields:
- `provider` to `CodeExchangeResponse`
- `custom_headers` to `CreateDraftRequest`, `UpdateDraftRequest`, and `SendMessageRequest`
---
 CHANGELOG.md                 |  6 +++++-
 nylas/models/auth.py         |  2 ++
 nylas/models/drafts.py       | 16 ++++++++++++++++
 tests/conftest.py            |  1 +
 tests/resources/test_auth.py |  1 +
 5 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 30638355..91d2926a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,11 @@ nylas-python Changelog
 
 Unreleased
 ----------------
-* Add clean messages support
+* Added support for adding custom headers to outgoing requests
+* Added support for `provider` field in code exchange response
+* Added clean messages support
+* Added additional webhook triggers
+* Made event visibility optional to support iCloud events
 
 v6.1.1
 ----------------
diff --git a/nylas/models/auth.py b/nylas/models/auth.py
index bb6274c7..3b57ac36 100644
--- a/nylas/models/auth.py
+++ b/nylas/models/auth.py
@@ -114,6 +114,7 @@ class CodeExchangeResponse:
         refresh_token: Returned only if the code is requested using "access_type=offline".
         id_token: A JWT that contains identity information about the user. Digitally signed by Nylas.
         token_type: Always "Bearer".
+        provider: The provider that the code was exchanged with.
     """
 
     access_token: str
@@ -124,6 +125,7 @@ class CodeExchangeResponse:
     refresh_token: Optional[str] = None
     id_token: Optional[str] = None
     token_type: Optional[str] = None
+    provider: Optional[Provider] = None
 
 
 @dataclass_json
diff --git a/nylas/models/drafts.py b/nylas/models/drafts.py
index 18900fa5..3f77aaed 100644
--- a/nylas/models/drafts.py
+++ b/nylas/models/drafts.py
@@ -57,6 +57,19 @@ class TrackingOptions(TypedDict):
     thread_replies: NotRequired[bool]
 
 
+class CustomHeader(TypedDict):
+    """
+    A key-value pair representing a header that can be added to drafts and outgoing messages.
+
+    Attributes:
+        name: The name of the custom header.
+        value: The value of the custom header.
+    """
+
+    name: str
+    value: str
+
+
 class CreateDraftRequest(TypedDict):
     """
     A request to create a draft.
@@ -73,6 +86,7 @@ class CreateDraftRequest(TypedDict):
         send_at: Unix timestamp to send the message at.
         reply_to_message_id: The ID of the message that you are replying to.
         tracking_options: Options for tracking opens, links, and thread replies.
+        custom_headers: Custom headers to add to the message.
     """
 
     body: NotRequired[str]
@@ -86,6 +100,7 @@ class CreateDraftRequest(TypedDict):
     send_at: NotRequired[int]
     reply_to_message_id: NotRequired[str]
     tracking_options: NotRequired[TrackingOptions]
+    custom_headers: NotRequired[List[CustomHeader]]
 
 
 UpdateDraftRequest = CreateDraftRequest
@@ -148,6 +163,7 @@ class SendMessageRequest(CreateDraftRequest):
         send_at (NotRequired[int]): Unix timestamp to send the message at.
         reply_to_message_id (NotRequired[str]): The ID of the message that you are replying to.
         tracking_options (NotRequired[TrackingOptions]): Options for tracking opens, links, and thread replies.
+        custom_headers(NotRequired[List[CustomHeader]]): Custom headers to add to the message.
         use_draft: Whether or not to use draft support. This is primarily used when dealing with large attachments.
     """
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 6fba23db..390dc54f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -113,6 +113,7 @@ def http_client_token_exchange():
         "scope": "https://www.googleapis.com/auth/gmail.readonly profile",
         "token_type": "Bearer",
         "grant_id": "grant_123",
+        "provider": "google",
     }
     return mock_http_client
 
diff --git a/tests/resources/test_auth.py b/tests/resources/test_auth.py
index d9f0eae8..07796158 100644
--- a/tests/resources/test_auth.py
+++ b/tests/resources/test_auth.py
@@ -99,6 +99,7 @@ def test_get_token(self, http_client_token_exchange):
         assert res.scope == "https://www.googleapis.com/auth/gmail.readonly profile"
         assert res.token_type == "Bearer"
         assert res.grant_id == "grant_123"
+        assert res.provider == "google"
 
     def test_get_token_info(self, http_client_token_info):
         auth = Auth(http_client_token_info)