first step in migratng

.vscode/launch.json

@@ -61,7 +61,7 @@
                 "SCRIBE_TOKEN": "${env:SCRIBE_TOKEN}",
             }
         },
-                        {
+        {
             "name": "get attestation",
             "type": "debugpy",
             "request": "launch",
@@ -76,7 +76,23 @@
             "env": {
                 "SCRIBE_TOKEN": "${env:SCRIBE_TOKEN}",
             }
+        },
+        {
+            "name": "get metadata",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "scribe_python_client.cli",
+            "console": "internalConsole",
+            "args": [
+                "--api-call",
+                "get-dataset-tables-dict",
+                "--env", "prod",
+            ],
+            "env": {
+                "SCRIBE_TOKEN": "${env:SCRIBE_TOKEN}",
+            }
         }
 
+
     ]
 }

scribe_python_client/cli.py

@@ -2,6 +2,7 @@ import argparse
 import os
 import logging
 from .client import ScribeClient
+import json
 
 def main():
     logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -21,6 +22,10 @@ def main():
         "query-lineage",
         "query-risk",
         "query-findings",
+        "get-dataset-links",
+        "get-dataset-data",
+        "get-dataset-tables-dict",
+        "get-dataset-tables-md",
     ], help="Which API call to execute")
     parser.add_argument("--api-token", required=False, default=os.getenv('SCRIBE_TOKEN'), help="Your API token from ScribeHub integrations page (defaults to SCRIBE_TOKEN environment variable)")
     parser.add_argument("--env", default="prod", choices=["prod", "dev", "test", "ci"], help="Which environment to use")
@@ -28,6 +33,9 @@ def main():
     parser.add_argument("--attestation-id", required=False, help="The ID of the attestation (required for get-attestation API call)")
     parser.add_argument("--query", required=False, help="The query to run on the dataset (required for query commands)")
     parser.add_argument("--lineage-graph-file", required=False, help="The file to save the lineage graph to (optional)")
+    parser.add_argument("--username", required=False, default = os.getenv('SUPERSET_DEV_USERNAME'), help="Superset username (required for Superset API calls)")
+    parser.add_argument("--password", required=False, default = os.getenv('SUPERSET_DEV_PASSWORD'), help="Superset password (required for Superset API calls)")
+
 
     args = parser.parse_args()
 
@@ -50,8 +58,13 @@ def main():
     base_url = "https://api.scribesecurity.com"
     if args.env != "prod":
         base_url = f"https://api.{args.env}.scribesecurity.com"
+    
+    if args.api_call in ["get-dataset-links", "get-dataset-data", "get-dataset-tables-dict", "get-dataset-tables-md"]:
+        if not args.username or not args.password:
+            logging.error("--username and --password are required for Superset table export commands.")
+            exit(1)
 
-    client = ScribeClient(api_token=args.api_token, base_url=base_url, env = args.env)
+    client = ScribeClient(api_token=args.api_token, base_url=base_url, env = args.env, username= args.username, password=args.password)
 
     logging.info(f"Executing API call: {args.api_call}")
 
@@ -86,5 +99,64 @@ def main():
     elif args.api_call == "query-findings":
         print(client.query_findings(querystr=args.query))
 
+    # Get dataset metadata section:
+    elif args.api_call == "get-dataset-links":
+        links = client.get_dataset_links()
+        if links:
+            print("Dataset Links:")
+            for link in links:
+                print(link)
+        else:
+            logging.error("No dataset links found.")
+        result = None
+    elif args.api_call == "get-dataset-data":
+        dataset_links = client.get_dataset_links()
+        if dataset_links:
+            client.get_dataset_data(dataset_links, output_file="dataset_metadata.csv")
+            logging.info("Metadata export completed successfully.")
+        else:
+            logging.error("No dataset links found.")
+        result = None
+    elif args.api_call == "get-dataset-tables-dict": 
+        dataset_links = client.get_dataset_links()
+        if dataset_links:
+            client.get_dataset_data(dataset_links, output_file="dataset_metadata.csv")
+            logging.info("Metadata export completed successfully.")
+        else:
+            logging.error("No dataset links found.")
+        result = None
+
+        tables_dict = client.get_dataset_tables_dict()
+        #Output path is set to docs/dataset_tables.json
+        output_path = os.path.join("docs", "dataset_tables.json")
+
+        with open(output_path, "w", encoding="utf-8") as f:
+            json.dump(tables_dict, f, indent=2)
+
+        logging.info(f"Saved table metadata dictionary to {output_path}")
+
+        if result is not None:
+            print(result)
+    elif args.api_call == "get-dataset-tables-md":
+        if not client.jwt_token:
+            username = input("Enter your Superset username: ")
+            password = input("Enter your Superset password: ")
+            client.username = username
+            client.password = password
+            if not client.authenticate():
+                logging.error("Superset login failed.")
+                exit(1)
+        dataset_links = client.get_dataset_links()
+        if dataset_links:
+            client.get_dataset_data(dataset_links, output_file="dataset_metadata.csv")
+            md_output = client.get_dataset_tables_md()
+            #Output path is set to docs/dataset_tables.md
+            output_path = os.path.join("docs", "dataset_tables.md")
+            with open(output_path, "w", encoding="utf-8") as f:
+                f.write(md_output)
+            logging.info(f"Saved table metadata Markdown to {output_path}")
+        else:
+            logging.error("No dataset links found.")
+
 if __name__ == "__main__":
     main()

scribe_python_client/client.py

@@ -47,7 +47,8 @@ class ScribeClient:
                 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={}, env = "prod"):
+    def __init__(self, api_token=None, ignore_env = False, base_url="https://api.scribesecurity.com", refresh_interval=600, cached_data={}, env = "prod", 
+                 username=None, password=None):
         def update_from_cached_data(cached_data):
             if "datasets" in cached_data:
                 self.datasets = cached_data["datasets"]
@@ -75,8 +76,16 @@ class ScribeClient:
         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
 
@@ -1091,4 +1100,271 @@ class ScribeClient:
             r = self.query_superset(self.scribe_findings_dataset, query, time_tokens=["timestamp"])
             if r:
                 return f"{title}\n{self.format_query_result_table(r)}"
-        return "Error running query"
+        return "Error running query"
+    
+    def superset_authenticate(self):
+        if not self.username or not self.password:
+            logging.error("Username and password must be provided for Superset login.")
+            return False
+
+        login_url = f"{self.superset_base_url}/api/v1/security/login"
+        body = {
+            "username": self.username,
+            "password": self.password,
+            "provider": "db"
+        }
+
+        response = self.send_post_request(login_url, body)
+        if not response:
+            logging.error("Login failed. No response received.")
+            return False
+
+        try:
+            result = json.loads(response)
+        except json.JSONDecodeError:
+            logging.error("Failed to parse login response.")
+            return False
+
+        access_token = result.get("access_token")
+        if access_token:
+            self.jwt_token = access_token
+            self.superset_token = access_token
+            logging.info("Login successful. Tokens obtained.")
+            return True
+        else:
+            logging.error("Login failed. No access token found.")
+            return False
+    
+    def get_dataset_links(self):
+        
+        self.superset_authenticate()
+        
+        if not self.jwt_token:
+            logging.error("No JWT token found. Please authenticate first.")
+            return []
+
+        dataset_links = []
+
+        #Use slug to directly fetch the API dashboard's datasets
+        dashboard_slug = "api-dashboard"
+        datasets_url = f"{self.superset_base_url}/api/v1/dashboard/{dashboard_slug}/datasets"
+        logging.info(f"Fetching datasets from {datasets_url}")
+
+        datasets_response = self.send_get_request(datasets_url, token=self.jwt_token)
+        if not datasets_response:
+            logging.error("Failed to fetch datasets for API dashboard.")
+            return []
+
+        try:
+            datasets_data = json.loads(datasets_response)
+            datasets = datasets_data.get("result", [])
+
+            for dataset in datasets:
+                dataset_id = dataset.get("id")
+                if dataset_id:
+                    link = f"{self.base_url}/explore/?datasource_type=table&datasource_id={dataset_id}"
+                    dataset_links.append(link)
+
+        except json.JSONDecodeError as e:
+            logging.error(f"Error parsing dataset response: {e}")
+
+        logging.info(f"Total dataset links collected: {len(dataset_links)}")
+        return dataset_links
+
+
+    def get_dataset_data(self, dataset_links, output_file="dataset_metadata.csv"):
+        records = []
+
+        for link in dataset_links:
+            #Properly parse dataset ID from the URL
+            parsed_url = urlparse(link)
+            query_params = parse_qs(parsed_url.query)
+            dataset_id_list = query_params.get("datasource_id")
+            if not dataset_id_list:
+                logging.error(f"Could not find dataset_id in link: {link}")
+                continue
+            dataset_id = dataset_id_list[0]
+
+            dataset_url = f"{self.base_url}/api/v1/dataset/{dataset_id}"
+            logging.info(f"Fetching metadata from {dataset_url}")
+
+            response = self.send_get_request(dataset_url)
+            if not response:
+                logging.error(f"Failed to fetch metadata for dataset ID {dataset_id}")
+                continue
+            try:
+                metadata_json = json.loads(response)
+                dataset_metadata = metadata_json.get("result", {})
+            except json.JSONDecodeError as e:
+                logging.error(f"Failed to parse metadata JSON for dataset ID {dataset_id}: {e}")
+                continue
+
+            verbose_name = dataset_metadata.get("verbose_name")
+            if not verbose_name or verbose_name == dataset_metadata.get("table_name"):
+                verbose_name = dataset_metadata.get("table_name", "").replace("_", " ").title()
+
+            #Get dataset metadata
+            dataset_info = {
+                "dataset_id": dataset_metadata.get("id"),
+                "catalog": dataset_metadata.get("catalog"),
+                "changed_on": dataset_metadata.get("changed_on"),
+                "table_name": dataset_metadata.get("table_name"),
+                "verbose_name": verbose_name,
+                "schema": dataset_metadata.get("schema"),
+                "datasource_type": dataset_metadata.get("datasource_type"),
+                "description": dataset_metadata.get("description"),
+            }
+
+            #Get column metadata
+            columns = dataset_metadata.get("columns", [])
+            if columns:
+                for col in columns:
+                    record = dataset_info.copy()
+                    record.update({
+                        "column_id": col.get("id"),
+                        "column_name": col.get("column_name"),
+                        "column_verbose_name": col.get("verbose_name") or col.get("column_name", "").replace("_", " ").title(),
+                        "column_type": col.get("type"),
+                        "groupby": col.get("groupby"),
+                        "filterable": col.get("filterable"),
+                        "is_dttm": col.get("is_dttm"),
+                        "column_description": col.get("description"),
+                    })
+                    records.append(record)
+            else:
+                records.append(dataset_info)
+
+        #Write to CSV
+        if records:
+            fieldnames = [
+                "dataset_id", "table_name", "verbose_name", "schema", "catalog",
+                "datasource_type", "changed_on", "description",
+                "column_id", "column_name", "column_verbose_name", "column_type",
+                "groupby", "filterable", "is_dttm", "column_description"
+            ]
+            fieldnames = [field for field in fieldnames if any(field in rec for rec in records)]
+
+            with open(output_file, mode='w', newline='', encoding='utf-8') as csvfile:
+                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+                writer.writeheader()
+                writer.writerows(records)
+
+            logging.info(f"Dataset metadata exported to {output_file}")
+        else:
+            logging.warning("No dataset metadata exported.")
+
+    def get_dataset_tables_dict(self, tables=[]) -> dict:
+        if not self.jwt_token:
+            self.superset_authenticate()
+
+        dataset_links = self.get_dataset_links()
+        if not dataset_links:
+            logging.error("No dataset links found.")
+            return {}
+
+        self.get_dataset_data(dataset_links)
+
+        #Load data from the CSV file (generated by get_dataset_data)
+        records = []
+        try:
+            with open("dataset_metadata.csv", mode="r", encoding="utf-8") as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    records.append(row)
+        except Exception as e:
+            logging.error(f"Failed to load dataset_metadata.csv: {e}")
+            return {}
+
+        #Group column data by table name
+        table_dict = {}
+        for row in records:
+            table_name = row.get("table_name")
+            if tables and table_name not in tables:
+                continue
+
+            if table_name not in table_dict:
+                table_dict[table_name] = []
+
+            table_dict[table_name].append({
+                "column_name": row.get("column_name", ""),
+                "type": row.get("column_type", ""),
+                "label": row.get("column_verbose_name") or row.get("column_name", "").replace("_", " ").title(),
+                "description": row.get("column_description", "")
+            })
+
+        return table_dict
+
+
+    def get_dataset_tables_md(self, tables=[]) -> str:
+        if not self.jwt_token:
+            self.superset_authenticate()
+
+        dataset_links = self.get_dataset_links()
+        if not dataset_links:
+            logging.error("No dataset links found.")
+            return ""
+
+        self.get_dataset_data(dataset_links)
+
+        # Load data from the CSV file
+        records = []
+        try:
+            with open("dataset_metadata.csv", mode="r", encoding="utf-8") as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    records.append(row)
+        except Exception as e:
+            logging.error(f"Failed to load dataset_metadata.csv: {e}")
+            return ""
+
+        # Group columns by table_name
+        table_columns = {}
+        for row in records:
+            table_name = row.get("table_name")
+            if tables and table_name not in tables:
+                continue
+
+            if table_name not in table_columns:
+                table_columns[table_name] = []
+
+            column_name = row.get("column_name", "")
+            col_type = row.get("column_type", "")
+            col_type = row.get("column_type", "").upper()
+            expr = row.get("expression", "") or ""
+
+            if col_type == "STRING" and (" in (" in expr.lower() or " in(" in expr.lower()):
+                like_compliant = False
+            else:
+                like_compliant = True
+
+
+            table_columns[table_name].append({
+                "column_name": column_name,
+                "label": row.get("column_verbose_name") or column_name.replace("_", " ").title(),
+                "type": col_type,
+                "description": row.get("column_description", ""),
+                "like_compliant": str(like_compliant)
+            })
+
+        # Build markdown
+        md_parts = []
+        for table_name, cols in table_columns.items():
+            md_parts.append(f"## `{table_name}` Columns\n\n")
+            md_parts.append("| Column Name | Label | Type | Description | Like-compliant |\n")
+            md_parts.append("|-------------|-------|------|-------------|----------------|\n")
+
+            for col in cols:
+                md_parts.append(
+                    f"| `{col['column_name']}` | {col['label']} | {col['type']} | "
+                    f"{col['description']} | {col['like_compliant']} |\n"
+                )
+            md_parts.append("\n")
+
+        markdown_output = "".join(md_parts)
+
+        # Write to file
+        with open("dataset_tables.md", "w", encoding="utf-8") as f:
+            f.write(markdown_output)
+
+        logging.info("Saved table metadata Markdown to dataset_tables.md")
+        return markdown_output