Merge pull request #11 from scribe-security/sh-8291-risk-dataset

scribe_python_client/client.py

@@ -929,19 +929,6 @@ class ScribeClient:
 
     }
 
-    def create_lineage_graph(self, lineage_data):
-            # Determine allowed asset types from the data.
-        # all_types = {record.get("asset_type") for record in lineage_data if record.get("asset_type")}
-        # # Exclude unwanted types.
-        # exclude_types = {"user", "secret", "org_secret", "commit", "actor", "build_machine"}
-        # include_types = all_types - exclude_types
-        
-        # # Define the mapping for layering.
-         
-        # # Generate and display the network visualization.
-        # fig = lineage_graph.visualize_network(lineage_data, include_types, self.types_to_position)
-        logging.error("create_lineage_graph is not supported in the python package.")
-        return None
     
     include_fields = [
         "platform_name", "platform_type",
@@ -1039,6 +1026,24 @@ class ScribeClient:
                 return f"{title}\n{self.format_query_result_table(r)}"
         return "Error running query"
 
+    def get_lineage_graph(self, lineage_data, output_file):
+        """
+        Generate a lineage graph from the lineage data and save it to the specified output file.
+        
+        Parameters:
+            lineage_data (list): List of lineage records.
+            output_file (str): Path to save the generated graph.
+        
+        Returns:
+            str: Full path of the saved graph file.
+        """
+        if not lineage_data:
+            logging.error("No lineage data provided for graph generation.")
+            return "Error: No lineage data provided."
+        
+        full_path_filename = create_lineage_graph(lineage_data, output_file=output_file)
+        return full_path_filename
+    
     def query_lineage(self, querystr, title="", graph_filename=None) -> str:
         logging.debug(f"Query lineage with: {querystr}")
         query = self.querystr_to_json(querystr)

scribe_python_client/version.py

@@ -1 +1 @@
-__version__ = "0.2.1"
+__version__ = "0.2.2"