Merge pull request #11 from elementary-data/enrich_graph_with_operational_context

Enrich graph with operational context

Author: oravi

lineage/lineage_graph.py

@@ -8,6 +8,8 @@
 from sqllineage.exceptions import SQLLineageException
 from pyvis.network import Network
 import webbrowser
+
+from lineage.query_context import QueryContext
 from lineage.utils import get_logger
 from sqllineage.models import Schema, Table
 from tqdm import tqdm
@@ -118,7 +120,9 @@ def _name_qualification(self, table: Table, database_name: str, schema_name: str
 
         return str(table).rsplit('.', 1)[-1]
 
-    def _update_lineage_graph(self, analyzed_statements: [LineageResult], database_name: str, schema_name: str) -> None:
+    def _update_lineage_graph(self, analyzed_statements: [LineageResult], query_context: QueryContext) -> None:
+        database_name = query_context.queried_database
+        schema_name = query_context.queried_schema
         for analyzed_statement in analyzed_statements:
             # Handle drop tables, if they exist in the statement
             dropped_tables = analyzed_statement.drop
@@ -139,9 +143,9 @@ def _update_lineage_graph(self, analyzed_statements: [LineageResult], database_n
             targets = {self._name_qualification(target, database_name, schema_name)
                        for target in analyzed_statement.write}
 
-            self._add_nodes_and_edges(sources, targets)
+            self._add_nodes_and_edges(sources, targets, query_context)
 
-    def _add_nodes_and_edges(self, sources: {str}, targets: {str}) -> None:
+    def _add_nodes_and_edges(self, sources: {str}, targets: {str}, query_context: QueryContext) -> None:
         if None in sources:
             sources.remove(None)
         if None in targets:
@@ -155,10 +159,10 @@ def _add_nodes_and_edges(self, sources: {str}, targets: {str}) -> None:
                 self._lineage_graph.add_nodes_from(sources)
         elif len(targets) > 0 and len(sources) == 0:
             if self._show_isolated_nodes:
-                self._lineage_graph.add_nodes_from(targets)
+                self._lineage_graph.add_nodes_from(targets, title=query_context.to_html())
         else:
             self._lineage_graph.add_nodes_from(sources)
-            self._lineage_graph.add_nodes_from(targets)
+            self._lineage_graph.add_nodes_from(targets, title=query_context.to_html())
             for source, target in itertools.product(sources, targets):
                 self._lineage_graph.add_edge(source, target)
 
@@ -191,15 +195,15 @@ def _remove_node(self, node: str) -> None:
 
     def init_graph_from_query_list(self, queries: [tuple]) -> None:
         logger.debug(f'Loading {len(queries)} queries into the lineage graph')
-        for query, database_name, schema_name in tqdm(queries, desc="Updating lineage graph", colour='green'):
+        for query, query_context in tqdm(queries, desc="Updating lineage graph", colour='green'):
             try:
                 analyzed_statements = self._parse_query(query)
             except SQLLineageException as exc:
                 logger.debug(f'SQLLineageException was raised while parsing this query -\n{query}\n'
                              f'Error was -\n{exc}.')
                 continue
 
-            self._update_lineage_graph(analyzed_statements, database_name, schema_name)
+            self._update_lineage_graph(analyzed_statements, query_context)
 
         logger.debug(f'Finished updating lineage graph!')
 

lineage/query_context.py

@@ -0,0 +1,73 @@
+from datetime import datetime
+from typing import Optional, Union
+import dateutil.parser
+
+
+class QueryContext(object):
+    def __init__(self, queried_database: Optional[str] = None, queried_schema: Optional[str] = None,
+                 query_time: Optional[datetime] = None, query_volume: Optional[int] = None,
+                 query_type: Optional[str] = None, user_name: Optional[str] = None,
+                 role_name: Optional[str] = None) -> None:
+        self.queried_database = queried_database
+        self.queried_schema = queried_schema
+        self.query_time = query_time
+        self.query_volume = query_volume
+        self.query_type = query_type
+        self.user_name = user_name
+        self.role_name = role_name
+
+    def to_dict(self) -> dict:
+        return {'queried_database': self.queried_database,
+                'queried_schema': self.queried_schema,
+                'query_time': self._query_time_to_str(self.query_time),
+                'query_volume': self.query_volume,
+                'query_type': self.query_type,
+                'user_name': self.user_name,
+                'role_name': self.role_name}
+
+    @staticmethod
+    def _query_time_to_str(query_time: Optional[datetime], fmt: str = None) -> Optional[str]:
+        if query_time is None:
+            return None
+
+        if fmt is None:
+            return query_time.isoformat()
+
+        return query_time.strftime(fmt)
+
+    @staticmethod
+    def _html_param_with_default(param: Union[str, int], default: Union[str, int] = 'unknown') -> Union[str, int]:
+        return default if param is None else param
+
+    def to_html(self) -> str:
+        query_type = self._html_param_with_default(self.query_type)
+        user_name = self._html_param_with_default(self.user_name)
+        role_name = self._html_param_with_default(self.role_name)
+        query_time = self._query_time_to_str(self.query_time, fmt='%Y-%m-%d %H:%M:%S')
+        query_volume = self._html_param_with_default(self.query_volume, 0)
+        volume_color = "DarkSlateGrey"
+        if query_volume == 0:
+            volume_color = "tomato"
+
+        return f"""
+                    <html>
+                        <body>
+                            <div style="font-family:arial;color:DarkSlateGrey;font-size:110%;">
+                                <strong>
+                                    Last update</br>
+                                </strong>
+                                <div style="min-width:62px;display:inline-block">Type:</div> {query_type}</br>
+                                <div style="min-width:62px;display:inline-block">User:</div> {user_name}</br>
+                                <div style="min-width:62px;display:inline-block">Role:</div> {role_name}</br>
+                                <div style="min-width:62px;display:inline-block">Time:</div> {query_time}</br>
+                                <div style="min-width:62px;display:inline-block;">Volume:</div> <a style="color:{volume_color}">{query_volume} rows</a>
+                            </div>
+                        </body>
+                    </html>
+        """
+
+    @staticmethod
+    def from_dict(query_context_dict: dict) -> 'QueryContext':
+        if 'query_time' in query_context_dict and query_context_dict['query_time'] is not None:
+            query_context_dict['query_time'] = dateutil.parser.parse(query_context_dict['query_time'])
+        return QueryContext(**query_context_dict)

lineage/query_history.py

@@ -1,5 +1,7 @@
 from datetime import datetime, timedelta
 from typing import Optional
+
+from lineage.query_context import QueryContext
 from lineage.utils import is_flight_mode_on
 import json
 import os
@@ -18,13 +20,18 @@ def __init__(self, con, should_export_query_history: bool = True) -> None:
     def _serialize_query_history(self, queries: [str]) -> None:
         if self.should_export_query_history:
             with open(self.LATEST_QUERY_HISTORY_FILE, 'w') as query_history_file:
-                json.dump(queries, query_history_file)
+                serialized_queries = []
+                for query, query_context in queries:
+                    serialized_queries.append((query, query_context.to_dict()))
+                json.dump(serialized_queries, query_history_file)
 
     def _deserialize_query_history(self) -> [str]:
         queries = []
         if os.path.exists(self.LATEST_QUERY_HISTORY_FILE):
             with open(self.LATEST_QUERY_HISTORY_FILE, 'r') as query_history_file:
-                queries = json.load(query_history_file)
+                deserialized_queries = json.load(query_history_file)
+                for query, query_context_dict in deserialized_queries:
+                    queries.append((query, QueryContext.from_dict(query_context_dict)))
         return queries
 
     @staticmethod
@@ -34,7 +41,7 @@ def _include_end_date(end_date: datetime) -> Optional[datetime]:
 
         return end_date
 
-    def extract_queries(self, start_date: datetime, end_date: datetime) -> [str]:
+    def extract_queries(self, start_date: datetime, end_date: datetime) -> [tuple]:
         if is_flight_mode_on():
             queries = self._deserialize_query_history()
         else:

lineage/snowflake_query_history.py

@@ -1,6 +1,7 @@
 from datetime import datetime, date, timedelta
 
 from lineage.exceptions import ConfigError
+from lineage.query_context import QueryContext
 from lineage.query_history import QueryHistory
 from lineage.utils import get_logger
 
@@ -14,7 +15,7 @@ class SnowflakeQueryHistory(QueryHistory):
     # the query was really executed on the configured db and filter it accordingly.
 
     INFORMATION_SCHEMA_QUERY_HISTORY = """
-    select query_text, database_name, schema_name
+    select query_text, database_name, schema_name, end_time, rows_produced, query_type, user_name, role_name
       from table(information_schema.query_history(
         end_time_range_start=>to_timestamp_ltz(:2),
         {end_time_range_end_expr},
@@ -31,7 +32,8 @@ class SnowflakeQueryHistory(QueryHistory):
     QUERY_HISTORY_SOURCE_INFORMATION_SCHEMA = 'information_schema'
 
     ACCOUNT_USAGE_QUERY_HISTORY = """
-    select query_text, database_name, schema_name 
+    select query_text, database_name, schema_name, end_time, rows_inserted + rows_produced, query_type, user_name, 
+    role_name
         from snowflake.account_usage.query_history 
         where end_time >= :2 and {end_time_range_end_expr} 
     and execution_status = 'SUCCESS' and query_type not in 
@@ -88,7 +90,7 @@ def _query_history_table(self, start_date: datetime, end_date: datetime) -> [tup
             logger.debug("Finished executing snowflake history query")
             rows = cursor.fetchall()
             for row in rows:
-                queries.append((row[0], row[1], row[2]))
+                queries.append((row[0], QueryContext(row[1], row[2], row[3], row[4], row[5], row[6], row[7])))
             logger.debug("Finished fetching snowflake history query results")
 
         return queries

setup.py

@@ -20,7 +20,7 @@
 setup(
     name='elementary-lineage',
     description='Presenting data lineage based on your data warehouse query history',
-    version='0.0.8',
+    version='0.0.9',
     packages=find_packages(),
     include_package_data=True,
     python_requires='>=3.6.2',

tests/test_lineage_graph.py

@@ -5,6 +5,7 @@
 from sqllineage.models import Schema
 
 from lineage.lineage_graph import LineageGraph
+from lineage.query_context import QueryContext
 
 
 def create_lineage_result(read, write):
@@ -74,13 +75,14 @@ def test_lineage_graph_add_nodes_and_edges(sources, targets, edges, show_isolate
     di_graph_mock = mock.create_autospec(nx.DiGraph)
     reference._lineage_graph = di_graph_mock
 
-    reference._add_nodes_and_edges(sources, targets)
+    empty_query_context = QueryContext()
+    reference._add_nodes_and_edges(sources, targets, empty_query_context)
 
     node_calls = []
     if len(sources) > 0:
         node_calls.append(mock.call(sources))
     if len(targets) > 0:
-        node_calls.append(mock.call(targets))
+        node_calls.append(mock.call(targets, title=empty_query_context.to_html()))
 
     edge_calls = []
     for edge in edges:
@@ -142,7 +144,7 @@ def test_lineage_graph_remove_node(show_isolated_nodes):
 def test_lineage_graph_init_graph_from_query_list_with_loops(queries, show_isolated_nodes):
     reference = LineageGraph(profile_database_name='elementary_db', show_isolated_nodes=show_isolated_nodes)
 
-    query_list = [(query, 'elementary_db', 'elementary_schema') for query in queries]
+    query_list = [(query, QueryContext('elementary_db', 'elementary_schema')) for query in queries]
     try:
         reference.init_graph_from_query_list(query_list)
     except Exception as exc: