From 388758ee4d5ea72ab4c9dbd81619844a7829f02f Mon Sep 17 00:00:00 2001
From: Arik Alon <alon.arik@gmail.com>
Date: Fri, 31 May 2024 11:46:56 +0300
Subject: [PATCH] bug fix - if refresh token was revoked, runner did not
 recover

---
 .../core/sinks/robusta/dal/supabase_dal.py    | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/robusta/core/sinks/robusta/dal/supabase_dal.py b/src/robusta/core/sinks/robusta/dal/supabase_dal.py
index 71e49e32e..f8ea6af0b 100644
--- a/src/robusta/core/sinks/robusta/dal/supabase_dal.py
+++ b/src/robusta/core/sinks/robusta/dal/supabase_dal.py
@@ -5,7 +5,9 @@
 from typing import Any, Dict, List, Optional
 
 import requests
+from postgrest._sync.request_builder import SyncQueryRequestBuilder
 from postgrest.base_request_builder import BaseFilterRequestBuilder
+from postgrest.exceptions import APIError as PostgrestAPIError
 from postgrest.types import ReturnMethod
 from postgrest.utils import sanitize_param
 from supabase import create_client
@@ -70,6 +72,7 @@ def __init__(
         self.cluster = cluster_name
         options = ClientOptions(postgrest_client_timeout=SUPABASE_TIMEOUT_SECONDS, auto_refresh_token=True)
         self.client = create_client(url, key, options)
+        self.patch_postgrest_execute()
         self.email = email
         self.password = password
         self.sign_in()
@@ -78,6 +81,24 @@ def __init__(
         self.persist_events = persist_events
         self.signing_key = signing_key
 
+    def patch_postgrest_execute(self):
+        # This is somewhat hacky.
+        def execute_with_retry(_self):
+            try:
+                return self._original_execute(_self)
+            except PostgrestAPIError as exc:
+                message = exc.message or ""
+                if exc.code == "PGRST301" or "expired" in message.lower():
+                    # JWT expired. Sign in again and retry the query
+                    logging.error("JWT token expired/invalid, signing in to Supabase again")
+                    self.sign_in()
+                    return self._original_execute(_self)
+                else:
+                    raise
+
+        self._original_execute = SyncQueryRequestBuilder.execute
+        SyncQueryRequestBuilder.execute = execute_with_retry
+
     def __to_db_scanResult(self, scanResult: ScanReportRow) -> Dict[Any, Any]:
         db_sr = scanResult.dict()
         db_sr["account_id"] = self.account_id