Add support for custom storage of oauth2 tokens

hubspot3/base.py

@@ -10,7 +10,8 @@
 import urllib.parse
 import urllib.error
 import zlib
-from typing import List, Union
+from typing import Callable, List, Optional, Union
+from typing_extensions import Literal
 from hubspot3 import utils
 from hubspot3.utils import force_utf8
 from hubspot3.error import (
@@ -30,6 +31,9 @@
 class BaseClient:
     """Base abstract object for interacting with the HubSpot APIs"""
 
+    # These rules are too restrictive for the `__init__()` method and attributes.
+    # pylint: disable=too-many-arguments,too-many-instance-attributes
+
     # Controls how long we sleep for during retries, overridden by unittests
     # so tests run faster
     sleep_multiplier = 1
@@ -41,6 +45,12 @@ def __init__(
         refresh_token: str = None,
         client_id: str = None,
         client_secret: str = None,
+        oauth2_token_getter: Optional[
+            Callable[[Literal["access_token", "refresh_token"], str], str]
+        ] = None,
+        oauth2_token_setter: Optional[
+            Callable[[Literal["access_token", "refresh_token"], str, str], None]
+        ] = None,
         timeout: int = 10,
         mixins: List = None,
         api_base: str = "api.hubapi.com",
@@ -58,10 +68,18 @@ def __init__(
                 self.__class__.__bases__ = (mixin_class,) + self.__class__.__bases__
 
         self.api_key = api_key
-        self.access_token = access_token
-        self.refresh_token = refresh_token
+        # These are used as fallbacks if there aren't setters/getters, or if no remote tokens can be
+        # found. The properties without `__` prefixes should generally be used instead of these.
+        self.__access_token = access_token
+        self.__refresh_token = refresh_token
         self.client_id = client_id
         self.client_secret = client_secret
+        if (oauth2_token_getter is None) != (oauth2_token_setter is None):
+            raise HubspotBadConfig(
+                "You must either specify both the oauth2 token setter and getter, or neither."
+            )
+        self.oauth2_token_getter = oauth2_token_getter
+        self.oauth2_token_setter = oauth2_token_setter
         self.log = utils.get_log("hubspot3")
         if not disable_auth:
             if self.api_key and self.access_token:
@@ -88,8 +106,41 @@ def credentials(self):
             "api_key": self.api_key,
             "access_token": self.access_token,
             "refresh_token": self.refresh_token,
+            "oauth2_token_getter": self.oauth2_token_getter,
+            "oauth2_token_setter": self.oauth2_token_setter,
         }
 
+    @property
+    def access_token(self):
+        if self.oauth2_token_getter:
+            return (
+                self.oauth2_token_getter("access_token", self.client_id)
+                or self.__access_token
+            )
+        return self.__access_token
+
+    @access_token.setter
+    def access_token(self, access_token):
+        if self.oauth2_token_setter:
+            self.oauth2_token_setter("access_token", self.client_id, access_token)
+        self.__access_token = access_token
+
+    @property
+    def refresh_token(self):
+        if self.oauth2_token_getter:
+            return (
+                self.oauth2_token_getter("refresh_token", self.client_id)
+                or self.__refresh_token
+            )
+        return self.__refresh_token
+
+    @refresh_token.setter
+    def refresh_token(self, refresh_token):
+        if self.oauth2_token_setter:
+            self.oauth2_token_setter("refresh_token", self.client_id, refresh_token)
+        else:
+            self.__refresh_token = refresh_token
+
     def _prepare_connection_type(self):
         connection_types = {
             "http": http.client.HTTPConnection,

hubspot3/test/test_lines.py

@@ -130,6 +130,10 @@ def test_link_line_item_to_deal(self, mock_associations_client, lines_client):
         mock_instance = mock_associations_client.return_value
         lines_client.link_line_item_to_deal(1, 1)
         mock_associations_client.assert_called_with(
-            access_token=None, api_key=None, refresh_token=None
+            access_token=None,
+            api_key=None,
+            refresh_token=None,
+            oauth2_token_getter=None,
+            oauth2_token_setter=None,
         )
         mock_instance.link_line_item_to_deal.assert_called_with(1, 1)