startarn added

Author: Yi-Ting Lee

src/sagemaker/lineage/query.py

@@ -15,6 +15,7 @@
 
 from datetime import datetime
 from enum import Enum
+from tracemalloc import start
 from typing import Optional, Union, List, Dict
 
 from sagemaker.lineage._utils import get_resource_name_from_arn
@@ -295,6 +296,7 @@ def __init__(
         self,
         edges: List[Edge] = None,
         vertices: List[Vertex] = None,
+        startarn: List[str] = None,
     ):
         """Init for LineageQueryResult.
 
@@ -304,13 +306,17 @@ def __init__(
         """
         self.edges = []
         self.vertices = []
+        self.startarn = []
 
         if edges is not None:
             self.edges = edges
 
         if vertices is not None:
             self.vertices = vertices
 
+        if startarn is not None:
+            self.startarn = startarn
+
     def __str__(self):
         """Define string representation of ``LineageQueryResult``.
 
@@ -335,7 +341,7 @@ def __str__(self):
 
         """
         result_dict = vars(self)
-        return str({k: [vars(val) for val in v] for k, v in result_dict.items()})
+        return str({k: [str(val) for val in v] for k, v in result_dict.items()})
 
     def _covert_vertices_to_tuples(self):
         """Convert vertices to tuple format for visualizer."""
@@ -456,9 +462,8 @@ def _get_vertex(self, vertex):
             sagemaker_session=self._session,
         )
 
-    def _convert_api_response(self, response) -> LineageQueryResult:
+    def _convert_api_response(self, response, converted) -> LineageQueryResult:
         """Convert the lineage query API response to its Python representation."""
-        converted = LineageQueryResult()
         converted.edges = [self._get_edge(edge) for edge in response["Edges"]]
         converted.vertices = [self._get_vertex(vertex) for vertex in response["Vertices"]]
 
@@ -541,7 +546,9 @@ def query(
             Filters=query_filter._to_request_dict() if query_filter else {},
             MaxDepth=max_depth,
         )
-        query_response = self._convert_api_response(query_response)
+        # create query result for startarn info
+        query_result = LineageQueryResult(startarn=start_arns)
+        query_response = self._convert_api_response(query_response, query_result)
         query_response = self._collapse_cross_account_artifacts(query_response)
 
         return query_response