cache fix

.vscode/launch.json

@@ -85,7 +85,7 @@
             "console": "internalConsole",
             "args": [
                 "--api-call",
-                "get-dataset-tables-md",
+                "get-dataset-tables-dict",
                 "--env", "prod",
             ],
             "env": {

scribe_python_client/client.py

@@ -7,7 +7,7 @@ import jwt
 import urllib #Needed for avoiding using requests, which modified the S3 presigned URL
 from urllib.parse import urlparse, parse_qs
 import csv
-
+import pathlib
 
 from .utils import convert_timestamps, send_post_request, send_get_request, decode_jwt_token, querystr_to_json, att2statement
 from .lineage_graph import create_lineage_graph 
@@ -49,112 +49,113 @@ class ScribeClient:
                 logging.error(f"Error parsing timestamp: {e}")
                 return True
             return (datetime.datetime.now() - last_update).days > 1
-def __init__(self, api_token=None, ignore_env = False, base_url="https://api.scribesecurity.com", refresh_interval=600, cached_data=None, env = "prod",
+        
+    def __init__(self, api_token=None, ignore_env = False, base_url="https://api.scribesecurity.com", refresh_interval=600, cached_data=None, env = "prod",
              username=None, password=None):
-    import pathlib
-
-    def update_from_cached_data(cached_data):
-        if "datasets" in cached_data:
-            self.datasets = cached_data["datasets"]
-            self.dataset_ids = self.create_dataset_id_map(self.datasets)
-        else:
-            logging.error("No datasets in provided cached data")
-        # product_list intentionally not loaded from cache
-        if "team_name" in cached_data:
-            self.team_name = cached_data["team_name"]
-        else:
-            logging.error("No team name in provided cached data")
-        if "team_info" in cached_data:
-            self.team_info = cached_data["team_info"]
-        else:
-            logging.error("No team info in provided cached data")
-
-    # Default cache file path
-    cache_dir = os.path.join(os.getcwd(), ".scribe")
-    cache_file = os.path.join(cache_dir, "cache.json")
-    loaded_cache = None
+    
 
-    # Only load cache if not explicitly provided
-    if cached_data is None:
-        try:
-            if os.path.exists(cache_file):
-                # Check file age
-                file_mtime = datetime.datetime.fromtimestamp(os.path.getmtime(cache_file))
-                if (datetime.datetime.now() - file_mtime).days <= 1:
-                    with open(cache_file, "r") as f:
-                        loaded_cache = json.load(f)
-                else:
-                    logging.info("Cache file is older than one day, will refresh cache.")
+        def update_from_cached_data(cached_data):
+            if "datasets" in cached_data:
+                self.datasets = cached_data["datasets"]
+                self.dataset_ids = self.create_dataset_id_map(self.datasets)
             else:
-                # Ensure .scribe directory exists
-                pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
-        except Exception as e:
-            logging.error(f"Error loading cache file: {e}")
-    else:
-        loaded_cache = cached_data
-
-    if ignore_env:
-        self.api_token = api_token
-    else:
-        self.api_token = api_token or os.getenv('SCRIBE_TOKEN', None)
-    if not self.api_token:
-        logging.error("No API token provided. Please set the SCRIBE_TOKEN environment variable or pass it as an argument.")
-        return
-    self.username = username
-    self.password = password
-    self.env = env
-    self.base_url = base_url
-    
-    if env == "prod":
-        self.superset_base_url = "https://superset.scribesecurity.com"
-    else:
-        self.superset_base_url = f"https://internal-superset.{env}.scribesecurity.com/login/"
-    
-    self.refresh_interval = refresh_interval
-    self.last_refresh = datetime.datetime.now() - datetime.timedelta(days=5)  # Initialize to a time in the past
-
-    self.jwt_token = None
-    self.superset_token = None
-    self.dataset_ids = None
-    self.datasets = None
-    self.product_list = None
-    self.team_id = None
-    self.team_name = None
-    self.team_info = None
-    p = None
-    if self.api_token:
-        self.last_refresh = datetime.datetime.now() # - datetime.timedelta(hours=5)
-        self.refresh_data(force=True)
-        
-        # Use loaded_cache if available and not stale, else refresh and save
-        if not loaded_cache or self.need_cache_data_update(loaded_cache):
-            d = self.get_datasets()
-            if not d:
-                logging.error("Refresh: Failed to update dataset IDs")
+                logging.error("No datasets in provided cached data")
+            # product_list intentionally not loaded from cache
+            if "team_name" in cached_data:
+                self.team_name = cached_data["team_name"]
             else:
-                self.datasets = d
-                self.dataset_ids = self.create_dataset_id_map(d)
-                p = self.get_products(force=True)
-            if not p:
-                logging.error("Refresh: Failed to update product list")
-                logging.error(f"Using previous product list\n{self.product_list}")
+                logging.error("No team name in provided cached data")
+            if "team_info" in cached_data:
+                self.team_info = cached_data["team_info"]
             else:
-                self.product_list = p
+                logging.error("No team info in provided cached data")
 
-            t = self.get_team_info()
-            if not t:
-                logging.error("Failed to get team data.")
+        # Default cache file path
+        cache_dir = os.path.join(os.getcwd(), ".scribe")
+        cache_file = os.path.join(cache_dir, "cache.json")
+        loaded_cache = None
 
-            # Save cache (excluding product_list)
+        # Only load cache if not explicitly provided
+        if cached_data is None:
             try:
-                pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
-                cache_data = self.get_chache_data()
-                with open(cache_file, "w") as f:
-                    json.dump(cache_data, f, indent=2)
+                if os.path.exists(cache_file):
+                    # Check file age
+                    file_mtime = datetime.datetime.fromtimestamp(os.path.getmtime(cache_file))
+                    if (datetime.datetime.now() - file_mtime).days <= 1:
+                        with open(cache_file, "r") as f:
+                            loaded_cache = json.load(f)
+                    else:
+                        logging.info("Cache file is older than one day, will refresh cache.")
+                else:
+                    # Ensure .scribe directory exists
+                    pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
             except Exception as e:
-                logging.error(f"Error saving cache file: {e}")
+                logging.error(f"Error loading cache file: {e}")
+        else:
+            loaded_cache = cached_data
+
+        if ignore_env:
+            self.api_token = api_token
+        else:
+            self.api_token = api_token or os.getenv('SCRIBE_TOKEN', None)
+        if not self.api_token:
+            logging.error("No API token provided. Please set the SCRIBE_TOKEN environment variable or pass it as an argument.")
+            return
+        self.username = username
+        self.password = password
+        self.env = env
+        self.base_url = base_url
+        
+        if env == "prod":
+            self.superset_base_url = "https://superset.scribesecurity.com"
         else:
-            update_from_cached_data(loaded_cache)
+            self.superset_base_url = f"https://internal-superset.{env}.scribesecurity.com/login/"
+        
+        self.refresh_interval = refresh_interval
+        self.last_refresh = datetime.datetime.now() - datetime.timedelta(days=5)  # Initialize to a time in the past
+
+        self.jwt_token = None
+        self.superset_token = None
+        self.dataset_ids = None
+        self.datasets = None
+        self.product_list = None
+        self.team_id = None
+        self.team_name = None
+        self.team_info = None
+        p = None
+        if self.api_token:
+            self.last_refresh = datetime.datetime.now() # - datetime.timedelta(hours=5)
+            self.refresh_data(force=True)
+            
+            # Use loaded_cache if available and not stale, else refresh and save
+            if not loaded_cache or self.need_cache_data_update(loaded_cache):
+                d = self.get_datasets()
+                if not d:
+                    logging.error("Refresh: Failed to update dataset IDs")
+                else:
+                    self.datasets = d
+                    self.dataset_ids = self.create_dataset_id_map(d)
+                    p = self.get_products(force=True)
+                if not p:
+                    logging.error("Refresh: Failed to update product list")
+                    logging.error(f"Using previous product list\n{self.product_list}")
+                else:
+                    self.product_list = p
+
+                t = self.get_team_info()
+                if not t:
+                    logging.error("Failed to get team data.")
+
+                # Save cache (excluding product_list)
+                try:
+                    pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
+                    cache_data = self.get_chache_data()
+                    with open(cache_file, "w") as f:
+                        json.dump(cache_data, f, indent=2)
+                except Exception as e:
+                    logging.error(f"Error saving cache file: {e}")
+            else:
+                update_from_cached_data(loaded_cache)
 
     def validate_dataset(self, dataset_name):
         if not self.dataset_ids: