Handle production Procore environments + saving refresh_tokens

Author: hsyyid

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tap-procore"
-version = "0.0.7"
+version = "0.0.8"
 description = "`tap-procore` is Singer tap for procore, built with the Singer SDK."
 authors = ["hotglue"]
 license = "Apache 2.0"

tap_procore/streams.py

@@ -6,8 +6,11 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Union, List, Iterable
 
+from singer.schema import Schema
 
 from singer_sdk.streams import RESTStream
+from singer_sdk.helpers._util import utc_now
+from singer_sdk.plugin_base import PluginBase as TapBaseClass
 
 
 from singer_sdk.authenticators import (
@@ -40,30 +43,64 @@ def oauth_request_body(self) -> dict:
             'grant_type': 'refresh_token',
             'client_id': self.config["client_id"],
             'client_secret': self.config["client_secret"],
-            'refresh_token': self.config["refresh_token"],
+            'refresh_token': self.config["refresh_token"] if not self.refresh_token else self.refresh_token,
             'redirect_uri': self.config["redirect_uri"]
         }
 
         return req
 
+    def update_access_token(self):
+        """Update `access_token` along with: `last_refreshed` and `expires_in`."""
+        request_time = utc_now()
+        auth_request_payload = self.oauth_request_payload
+        token_response = requests.post(self.auth_endpoint, data=auth_request_payload)
+        try:
+            token_response.raise_for_status()
+            self.logger.info("OAuth authorization attempt was successful.")
+        except Exception as ex:
+            raise RuntimeError(
+                f"Failed OAuth login, response was '{token_response.json()}'. {ex}"
+            )
+        token_json = token_response.json()
+        self.access_token = token_json["access_token"]
+        self.expires_in = token_json["expires_in"]
+        self.last_refreshed = request_time
+
+        if token_json.get("refresh_token") is not None:
+            self.refresh_token = token_json["refresh_token"]
+
 
 class ProcoreStream(RESTStream):
     """Procore stream class."""
 
+    def __init__(
+        self,
+        tap: TapBaseClass,
+        name: Optional[str] = None,
+        schema: Optional[Union[Dict[str, Any], Schema]] = None,
+        path: Optional[str] = None,
+    ):
+        """Initialize the Procore stream."""
+        super().__init__(name=name, schema=schema, tap=tap, path=path)
+        self._config = tap._config
+
     @property
     def url_base(self) -> str:
         """Return the API URL root, configurable via tap settings."""
         return "https://sandbox.procore.com/rest/v1.0" if self.config["is_sandbox"] else "https://api.procore.com/rest/v1.0"
 
     @property
     def authenticator(self) -> APIAuthenticatorBase:
-        auth_endpoint = "https://login-sandbox.procore.com/oauth/token" if self.config[
-            "is_sandbox"] else "https://login.procore.com/oauth/token"
+        if not self._config.get("authenticator"):
+            auth_endpoint = "https://login-sandbox.procore.com/oauth/token" if self.config[
+                "is_sandbox"] else "https://login.procore.com/oauth/token"
+
+            self._config["authenticator"] = ProcoreAuthenticator(
+                stream=self,
+                auth_endpoint=auth_endpoint
+            )
 
-        return ProcoreAuthenticator(
-            stream=self,
-            auth_endpoint=auth_endpoint
-        )
+        return self._config["authenticator"]
 
 
 class CompaniesStream(ProcoreStream):

tap_procore/tap.py

@@ -1,7 +1,8 @@
 """procore tap class."""
 
-from pathlib import Path
-from typing import List
+from pathlib import Path, PurePath
+from typing import List, Optional, Union
+import json
 
 from singer_sdk import Tap, Stream
 from singer_sdk.typing import (
@@ -54,10 +55,51 @@ class TapProcore(Tap):
     ).to_dict()
 
 
+    def __init__(
+        self,
+        config: Optional[Union[dict, PurePath, str, List[Union[PurePath, str]]]] = None,
+        catalog: Union[PurePath, str, dict, None] = None,
+        state: Union[PurePath, str, dict, None] = None,
+        parse_env_config: bool = False,
+    ) -> None:
+        """Initialize the tap."""
+        self.config_path = config
+        super().__init__(config=config, catalog=catalog, state=state, parse_env_config=parse_env_config)
+
+
     def discover_streams(self) -> List[Stream]:
         """Return a list of discovered streams."""
         return [stream_class(tap=self) for stream_class in STREAM_TYPES]
 
+    def sync_all(self):
+        """Sync all streams."""
+        # Do the sync
+        try:
+            super().sync_all()
+        except Exception as e:
+            self.logger.error(e)
+        finally:
+            # Update config if needed
+            self.update_config()
+
+    def update_config(self):
+        """Update config.json with new access + refresh token."""
+        self.logger.info("Updating config.")
+        path = self.config_path
+        auth = self._config.pop("authenticator", None)
+
+        if auth is not None:
+            if auth.refresh_token is not None:
+                self._config["refresh_token"] = auth.refresh_token
+            
+            if auth.access_token is not None:
+                self._config["access_token"] = auth.access_token
+
+        if isinstance(path, list):
+            path = path[0]
+
+        with open(path, 'w') as f:
+            json.dump(self._config, f, indent=4)
 
 # CLI Execution: