fix(salesforce): lock-aware token refresh for RTR safety

Replaces the fire-and-forget login() with a distributed-lock-aware
implementation that coordinates with Argo (same lock protocol used by
mk-node-libs lockAwareRefreshFn). Fixes the 15-min crash where the
refresh timer reused a revoked refresh_token after Salesforce RTR
rotation.

- Acquires Argo refresh-lock before calling SF
- Re-reads credentials after lock acquire (adopts if another service
  already refreshed)
- Persists new AT + RT to Argo atomically via lock-release endpoint
- Falls back to simple (in-memory-only) login when ARGO_URL/TENANT
  env vars are absent

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: malvfr

tap_salesforce/salesforce/credentials.py

@@ -1,5 +1,8 @@
 import logging
+import os
 import threading
+import time
+import uuid as _uuid_mod
 from collections import namedtuple
 
 import requests
@@ -61,7 +64,6 @@ def from_credentials(cls, credentials, **kwargs):
 
 
 class SalesforceAuthOAuth(SalesforceAuth):
-    # The minimum expiration setting for SF Refresh Tokens is 15 minutes
     REFRESH_TOKEN_EXPIRATION_PERIOD = 900
 
     @property
@@ -70,39 +72,195 @@ def _login_body(self):
 
     @property
     def _login_url(self):
-        login_url = "https://login.salesforce.com/services/oauth2/token"
-
         if self.is_sandbox:
-            login_url = "https://test.salesforce.com/services/oauth2/token"
-
-        return login_url
+            return "https://test.salesforce.com/services/oauth2/token"
+        return "https://login.salesforce.com/services/oauth2/token"
 
     def login(self):
-        try:
-            LOGGER.info("Attempting login via OAuth2")
+        argo_url = os.environ.get("ARGO_URL")
+        tenant = os.environ.get("TENANT")
+        if argo_url and tenant:
+            self._lock_aware_login(argo_url, tenant)
+        else:
+            LOGGER.warning(
+                "SF RTR: ARGO_URL or TENANT not set — using simple login; "
+                "rotated tokens will not be persisted (RTR unsafe)"
+            )
+            self._simple_login()
+
+    # ──────────────────────────────────────────────────────────
+    # Simple login (fallback — no lock, no writeback)
+    # ──────────────────────────────────────────────────────────
 
+    def _simple_login(self):
+        resp = None
+        try:
+            LOGGER.info("Attempting login via OAuth2 (simple)")
             resp = requests.post(
                 self._login_url,
                 data=self._login_body,
                 headers={"Content-Type": "application/x-www-form-urlencoded"},
             )
-
             resp.raise_for_status()
             auth = resp.json()
-
             LOGGER.info("OAuth2 login successful")
             self._access_token = auth["access_token"]
             self._instance_url = auth["instance_url"]
+            new_rt = auth.get("refresh_token")
+            if new_rt:
+                # ponytail: in-memory only — no Argo writeback; next call uses new RT
+                self._credentials = self._credentials._replace(refresh_token=new_rt)
         except Exception as e:
             error_message = str(e)
-            if resp:
-                error_message = error_message + f", Response from Salesforce: {resp.text}"
+            if resp is not None:
+                error_message += f", Response from Salesforce: {resp.text}"
             raise Exception(error_message) from e
         finally:
             LOGGER.info("Starting new login timer")
             self.login_timer = threading.Timer(self.REFRESH_TOKEN_EXPIRATION_PERIOD, self.login)
             self.login_timer.start()
 
+    # ──────────────────────────────────────────────────────────
+    # Lock-aware login — mirrors lockAwareRefreshFn in mk-node-libs
+    # ──────────────────────────────────────────────────────────
+
+    def _lock_aware_login(self, argo_url: str, tenant: str) -> None:
+        argo_api_key = os.environ.get("ARGO_CONNECTOR_API_KEY", "")
+        lock_id = str(_uuid_mod.uuid4())
+        old_rt = self._credentials.refresh_token
+        lock_acquired = False
+
+        try:
+            LOGGER.info("SF RTR: acquiring refresh lock tenant=%s lock_id=%s", tenant, lock_id)
+            self._acquire_lock(argo_url, argo_api_key, tenant, lock_id)
+            lock_acquired = True
+            LOGGER.info("SF RTR: lock acquired tenant=%s lock_id=%s", tenant, lock_id)
+
+            # Re-read after lock: another service may have refreshed while we waited.
+            # Skip on first-ever login (self._access_token is None) — we must always
+            # call SF on startup to get a fresh AT regardless of what Argo has cached.
+            if self._access_token is not None:
+                fresh = self._read_credentials_from_argo(argo_url, argo_api_key, tenant)
+                if fresh and fresh.get("access_token") and fresh["access_token"] != self._access_token:
+                    LOGGER.info(
+                        "SF RTR: adopting tokens refreshed by another service — releasing lock "
+                        "tenant=%s lock_id=%s",
+                        tenant, lock_id,
+                    )
+                    self._access_token = fresh["access_token"]
+                    new_rt = fresh.get("refresh_token")
+                    if new_rt:
+                        self._credentials = self._credentials._replace(refresh_token=new_rt)
+                    self._instance_url = fresh.get("instance_url", self._instance_url)
+                    self._release_lock(argo_url, argo_api_key, tenant, lock_id)
+                    return
+
+            # Token still stale — call Salesforce.
+            LOGGER.info("SF RTR: calling Salesforce tenant=%s lock_id=%s", tenant, lock_id)
+            resp = requests.post(
+                self._login_url,
+                data=self._login_body,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                timeout=30,
+            )
+            resp.raise_for_status()
+            auth = resp.json()
+            LOGGER.info("SF RTR: Salesforce returned new tokens tenant=%s lock_id=%s", tenant, lock_id)
+
+            self._access_token = auth["access_token"]
+            self._instance_url = auth.get("instance_url", self._instance_url)
+            new_rt = auth.get("refresh_token")
+            if new_rt:
+                self._credentials = self._credentials._replace(refresh_token=new_rt)
+
+            # Persist tokens to Argo and release lock atomically.
+            self._release_lock_with_tokens(argo_url, argo_api_key, tenant, lock_id, old_rt, auth)
+            LOGGER.info("SF RTR: tokens persisted and lock released tenant=%s lock_id=%s", tenant, lock_id)
+
+        except Exception as e:
+            LOGGER.error("SF RTR: error during lock-aware login tenant=%s lock_id=%s: %s", tenant, lock_id, e)
+            if lock_acquired:
+                try:
+                    self._release_lock(argo_url, argo_api_key, tenant, lock_id)
+                except Exception as re:
+                    LOGGER.error("SF RTR: failed to release lock after error tenant=%s lock_id=%s: %s", tenant, lock_id, re)
+            raise
+        finally:
+            LOGGER.info("Starting new login timer")
+            self.login_timer = threading.Timer(self.REFRESH_TOKEN_EXPIRATION_PERIOD, self.login)
+            self.login_timer.start()
+
+    # ── Argo helpers ───────────────────────────────────────────
+
+    def _acquire_lock(self, argo_url: str, argo_api_key: str, tenant: str, lock_id: str) -> None:
+        max_wait_s = 25.0
+        delay_s = 0.2
+        start = time.monotonic()
+        while True:
+            resp = requests.post(
+                f"{argo_url}/v1/tenant/{tenant}/connectors/salesforce/refresh-lock",
+                json={"lockId": lock_id, "calledBy": "mk-tap-salesforce"},
+                headers={"X-Api-Key": argo_api_key},
+                timeout=10,
+            )
+            resp.raise_for_status()
+            if resp.json().get("acquired"):
+                return
+            elapsed = time.monotonic() - start
+            if elapsed + delay_s >= max_wait_s:
+                raise Exception(
+                    f"SF RTR: lock acquire timeout after {max_wait_s}s "
+                    f"(tenant={tenant}, lock_id={lock_id})"
+                )
+            time.sleep(delay_s)
+            delay_s = min(delay_s * 2, 2.0)
+
+    def _read_credentials_from_argo(self, argo_url: str, argo_api_key: str, tenant: str) -> dict:
+        resp = requests.get(
+            f"{argo_url}/v1/tenant/{tenant}/connectors/salesforce",
+            headers={"X-Api-Key": argo_api_key},
+            timeout=10,
+        )
+        resp.raise_for_status()
+        return resp.json().get("credentials") or {}
+
+    def _release_lock(self, argo_url: str, argo_api_key: str, tenant: str, lock_id: str) -> None:
+        resp = requests.put(
+            f"{argo_url}/v1/tenant/{tenant}/connectors/salesforce/refresh-lock/release",
+            json={"lockId": lock_id, "calledBy": "mk-tap-salesforce"},
+            headers={"X-Api-Key": argo_api_key},
+            timeout=10,
+        )
+        resp.raise_for_status()
+
+    def _release_lock_with_tokens(
+        self,
+        argo_url: str,
+        argo_api_key: str,
+        tenant: str,
+        lock_id: str,
+        old_rt: str,
+        tokens: dict,
+    ) -> None:
+        body = {
+            "lockId": lock_id,
+            "oldRefreshToken": old_rt,
+            "tokens": {
+                "access_token": tokens["access_token"],
+                **({"refresh_token": tokens["refresh_token"]} if tokens.get("refresh_token") else {}),
+                **({"signature": tokens["signature"]} if tokens.get("signature") else {}),
+                **({"issued_at": tokens["issued_at"]} if tokens.get("issued_at") else {}),
+            },
+            "calledBy": "mk-tap-salesforce",
+        }
+        resp = requests.put(
+            f"{argo_url}/v1/tenant/{tenant}/connectors/salesforce/refresh-lock/release",
+            json=body,
+            headers={"X-Api-Key": argo_api_key},
+            timeout=10,
+        )
+        resp.raise_for_status()
+
 
 class SalesforceAuthPassword(SalesforceAuth):
     def login(self):