From 7cf9049763b976f87c341e2498d1c263a3fb02b4 Mon Sep 17 00:00:00 2001
From: Haochen Pan <haochenpan@uchicago.edu>
Date: Thu, 7 Dec 2023 14:14:48 -0600
Subject: [PATCH] token store

---
 diaspora_event_sdk/sdk/kafka_client.py        |  2 +-
 .../sdk/login_manager/tokenstore.py           | 34 +++----------------
 2 files changed, 6 insertions(+), 30 deletions(-)

diff --git a/diaspora_event_sdk/sdk/kafka_client.py b/diaspora_event_sdk/sdk/kafka_client.py
index 779a19b..16cd37d 100644
--- a/diaspora_event_sdk/sdk/kafka_client.py
+++ b/diaspora_event_sdk/sdk/kafka_client.py
@@ -103,7 +103,7 @@ def consumer_connection_test(result):
     start_time = time.time()
     while len(result) < 2:  # two tests
         if retry_count > 0:
-            print("Block until connected... retry count:", retry_count, ", time passed:", int(time.time() - start_time), "seconds")
+            print(f"Block until connected or timed out ({max_minutes} minutes)... retry count:", retry_count, ", time passed:", int(time.time() - start_time), "seconds")
         producer_connection_test(result)
         consumer_connection_test(result)
         retry_count += 1
diff --git a/diaspora_event_sdk/sdk/login_manager/tokenstore.py b/diaspora_event_sdk/sdk/login_manager/tokenstore.py
index b931e68..1e9857d 100644
--- a/diaspora_event_sdk/sdk/login_manager/tokenstore.py
+++ b/diaspora_event_sdk/sdk/login_manager/tokenstore.py
@@ -18,29 +18,10 @@ def _home() -> pathlib.Path:
     return pathlib.Path.home()
 
 
-def invalidate_old_config() -> None:
-    token_file = _home() / ".funcx" / "credentials" / "funcx_sdk_tokens.json"
-
-    if token_file.exists():
-        try:
-            auth_client = internal_auth_client()
-            with open(token_file) as fp:
-                data = json.load(fp)
-            for token_data in data.values():
-                if "access_token" in token_data:
-                    auth_client.oauth2_revoke_token(token_data["access_token"])
-                if "refresh_token" in token_data:
-                    auth_client.oauth2_revoke_token(
-                        token_data["refresh_token"])
-        finally:
-            os.remove(token_file)
-
-
-def ensure_compute_dir() -> pathlib.Path:
-    legacy_dirname = _home() / ".funcx"
-    dirname = _home() / ".globus_compute"
-
-    user_dir = os.getenv("GLOBUS_COMPUTE_USER_DIR")
+def ensure_diaspora_dir() -> pathlib.Path:
+    dirname = _home() / ".diaspora"
+
+    user_dir = os.getenv("DIASPORA_USER_DIR")
     if user_dir:
         dirname = pathlib.Path(user_dir)
 
@@ -51,9 +32,6 @@ def ensure_compute_dir() -> pathlib.Path:
             f"Error creating directory {dirname}, "
             "please remove or rename the conflicting file"
         )
-    elif legacy_dirname.is_dir() and not user_dir:
-        legacy_dirname.replace(dirname)
-        legacy_dirname.symlink_to(dirname, target_is_directory=True)
     else:
         dirname.mkdir(mode=0o700, parents=True, exist_ok=True)
 
@@ -61,7 +39,7 @@ def ensure_compute_dir() -> pathlib.Path:
 
 
 def _get_storage_filename():
-    datadir = ensure_compute_dir()
+    datadir = ensure_diaspora_dir()
     return os.path.join(datadir, "storage.db")
 
 
@@ -90,8 +68,6 @@ def get_token_storage_adapter(*, environment: str | None = None) -> SQLiteAdapte
     # when initializing the token storage adapter, check if the storage file exists
     # if it does not, then use this as a flag to clean the old config
     fname = _get_storage_filename()
-    if not os.path.exists(fname):
-        invalidate_old_config()
     # namespace is equal to the current environment
     return SQLiteAdapter(
         fname,