Merge pull request #10 from elementary-data/table_filters

Table filters

Author: oravi

lineage/lineage_graph.py

@@ -3,6 +3,7 @@
 
 import networkx as nx
 import sqlparse
+from lineage.exceptions import ConfigError
 from sqllineage.core import LineageAnalyzer, LineageResult
 from sqllineage.exceptions import SQLLineageException
 from pyvis.network import Network
@@ -16,7 +17,10 @@
 GRAPH_VISUALIZATION_OPTIONS = """{
             "edges": {
                 "color": {
-                    "inherit": true
+                  "color": "rgba(23,107,215,1)",
+                  "highlight": "rgba(23,107,215,1)",
+                  "hover": "rgba(23,107,215,1)",
+                  "inherit": false
                 },
                 "dashes": true,
                 "smooth": {
@@ -38,8 +42,12 @@
                 }
             },
             "interaction": {
+                "hover": true,
                 "navigationButtons": true,
-                "multiselect": true
+                "multiselect": true,
+                "keyboard": {
+                    "enabled": true
+                }
             },
             "physics": {
                 "enabled": false,
@@ -53,6 +61,12 @@
 
 
 class LineageGraph(object):
+    UPSTREAM_DIRECTION = 'upstream'
+    DOWNSTREAM_DIRECTION = 'downstream'
+    BOTH_DIRECTIONS = 'both'
+    SELECTED_NODE_COLOR = '#0925C7'
+    SELECTED_NODE_TITLE = 'Selected table<br/>'
+
     def __init__(self, profile_database_name: str, profile_schema_name: str = None, show_isolated_nodes: bool = False,
                  full_table_names: bool = False) -> None:
         self._lineage_graph = nx.DiGraph()
@@ -189,6 +203,48 @@ def init_graph_from_query_list(self, queries: [tuple]) -> None:
 
         logger.debug(f'Finished updating lineage graph!')
 
+    def filter_on_table(self, selected_table: str, direction: str = None, depth: int = None) -> None:
+        logger.debug(f'Filtering lineage graph on table - {selected_table}')
+        resolved_selected_table_name = self._name_qualification(Table(selected_table), self._profile_database_name,
+                                                                self._profile_schema_name)
+        logger.debug(f'Qualified table name - {resolved_selected_table_name}')
+        if resolved_selected_table_name is None:
+            raise ConfigError(f'Could not resolve table name - {selected_table}, please make sure to '
+                              f'specify a table name that exists in the database configured in your profiles file.')
+
+        if direction == self.DOWNSTREAM_DIRECTION:
+            self._lineage_graph = self._downstream_graph(resolved_selected_table_name, depth)
+        elif direction == self.UPSTREAM_DIRECTION:
+            self._lineage_graph = self._upstream_graph(resolved_selected_table_name, depth)
+        elif direction == self.BOTH_DIRECTIONS:
+            downstream_graph = self._downstream_graph(resolved_selected_table_name, depth)
+            upstream_graph = self._upstream_graph(resolved_selected_table_name, depth)
+            self._lineage_graph = nx.compose(upstream_graph, downstream_graph)
+        else:
+            raise ConfigError(f'Direction must be one of the following - {self.UPSTREAM_DIRECTION}|'
+                              f'{self.DOWNSTREAM_DIRECTION}|{self.BOTH_DIRECTIONS}, '
+                              f'Got - {direction} instead.')
+
+        self._update_selected_node_attributes(resolved_selected_table_name)
+        logger.debug(f'Finished filtering lineage graph on table - {selected_table}')
+        pass
+
+    def _downstream_graph(self, source_node: str, depth: Optional[int]) -> nx.DiGraph:
+        logger.debug(f'Building a downstream graph for - {source_node}, depth - {depth}')
+        return nx.bfs_tree(G=self._lineage_graph, source=source_node, depth_limit=depth)
+
+    def _upstream_graph(self, target_node: str, depth: Optional[int]) -> nx.DiGraph:
+        logger.debug(f'Building an upstream graph for - {target_node}, depth - {depth}')
+        reversed_lineage_graph = self._lineage_graph.reverse(copy=True)
+        return nx.bfs_tree(G=reversed_lineage_graph, source=target_node, depth_limit=depth).reverse(copy=False)
+
+    def _update_selected_node_attributes(self, selected_node: str) -> None:
+        if self._lineage_graph.has_node(selected_node):
+            node = self._lineage_graph.nodes[selected_node]
+            node_title = node.get('title', '')
+            node.update({'color': self.SELECTED_NODE_COLOR,
+                         'title': self.SELECTED_NODE_TITLE + node_title})
+
     def draw_graph(self, should_open_browser: bool = True) -> None:
         # Visualize the graph
         net = Network(height="100%", width="100%", directed=True)

lineage/main.py

@@ -94,20 +94,51 @@ def handle_parse_result(self, ctx, opts, args):
     default=False,
     help="Indicates if the lineage should filter only on tables in the configured profile's schema."
 )
+@click.option(
+    '--table', '-t',
+    type=str,
+    help="Filter on a table to see upstream and downstream dependencies of this table (see also direction param)."
+         " Table name format could be a full name like <db_name>.<schema_name>.<table_name>, a partial name like "
+         "<schema_name>.<table_name> or only a table name <table_name>. If the database name wasn't part of the name "
+         "the profiles database name will be used, if the schema name wasn't part of the name the profiles schema name "
+         "will be used.",
+    default=None
+)
+@click.option('--direction',
+              type=click.Choice([LineageGraph.UPSTREAM_DIRECTION, LineageGraph.DOWNSTREAM_DIRECTION,
+                                 LineageGraph.BOTH_DIRECTIONS]),
+              help="Sets direction of dependencies when filtering on a specific table (default is both, "
+                   "meaning showing both upstream and downstream dependencies of this table).",
+              default='both',
+              cls=RequiredIf,
+              required_if='table')
+@click.option('--depth',
+              type=int,
+              help="Sets how many levels of dependencies to show when filtering on a specific table "
+                   "(default is showing all levels of dependencies).",
+              default=None,
+              cls=RequiredIf,
+              required_if='table')
 def main(start_date: datetime, end_date: datetime, profiles_dir: str, profile_name: str, open_browser: bool,
-         export_query_history: bool, full_table_names: bool, ignore_schema: bool) -> None:
+         export_query_history: bool, full_table_names: bool, ignore_schema: bool, table: str, direction: str,
+         depth: int) -> None:
     """
     For more details check out our documentation here - https://docs.elementary-data.com/
     """
     click.echo(f"Any feedback and suggestions are welcomed! join our community here - "
                f"https://bit.ly/slack-elementary\n")
+
     query_history = QueryHistoryFactory(profiles_dir, profile_name, export_query_history).create_query_history()
     queries = query_history.extract_queries(start_date, end_date)
+
     lineage_graph = LineageGraph(show_isolated_nodes=False,
                                  profile_database_name=query_history.get_database_name(),
                                  profile_schema_name=query_history.get_schema_name() if not ignore_schema else None,
                                  full_table_names=full_table_names)
     lineage_graph.init_graph_from_query_list(queries)
+    if table is not None:
+        lineage_graph.filter_on_table(table, direction, depth)
+
     lineage_graph.draw_graph(should_open_browser=open_browser)
 
 

tests/test_lineage_graph.py

@@ -200,3 +200,77 @@ def test_lineage_graph_name_qualification(resolved_table_name, show_full_table_n
     reference = LineageGraph(profile_database_name='elementary_db', profile_schema_name='elementary_sc',
                              full_table_names=show_full_table_names)
     assert reference._name_qualification(Table(resolved_table_name), '', '') == expected_result
+
+
+def compare_edges(directed_graph: nx.DiGraph, edges: [set]) -> bool:
+    for edge in directed_graph.edges:
+        if (edge[0], edge[1]) not in edges:
+            return False
+    return True
+
+
+def create_directed_graph_from_edge_list(edges: [list]) -> nx.DiGraph:
+    G = nx.DiGraph()
+    for edge in edges:
+        G.add_node(edge[0])
+        G.add_node(edge[1])
+        G.add_edge(edge[0], edge[1])
+    return G
+
+
+@pytest.mark.parametrize("edges, selected_node, depth, expected_remaining_edges", [
+    (nx.path_graph(5).edges, 3, None, {(3, 4)}),
+    (nx.path_graph(5).edges, 3, 1, {(3, 4)}),
+    (nx.path_graph(7).edges, 3, 1, {(3, 4)}),
+    ([(0, 1), (1, 2), (0, 3), (3, 4), (3, 2), (2, 5), (4, 6)], 3, None, {(3, 4), (3, 2), (2, 5), (4, 6)}),
+    ([(0, 1), (1, 2), (0, 3), (3, 4), (3, 2), (2, 5), (4, 6)], 3, 1, {(3, 4), (3, 2)})
+])
+def test_lineage_graph_downstream_graph(edges, selected_node, depth, expected_remaining_edges):
+    reference = LineageGraph(profile_database_name='elementary_db')
+    reference._lineage_graph.add_edges_from(edges)
+    reference._lineage_graph = reference._downstream_graph(selected_node, depth)
+    assert compare_edges(reference._lineage_graph, expected_remaining_edges)
+
+
+@pytest.mark.parametrize("edges, selected_node, depth, expected_remaining_edges", [
+    (nx.path_graph(5).edges, 3, None, {(0, 1), (1, 2), (2, 3)}),
+    (nx.path_graph(5).edges, 3, 1, {(2, 3)}),
+    (nx.path_graph(7).edges, 3, 1, {(2, 3)}),
+    ([(0, 1), (1, 2), (2, 3), (0, 3), (3, 4), (2, 5), (4, 6)], 3, None, {(0, 3), (2, 3), (1, 2), (0, 1)}),
+    ([(0, 1), (1, 2), (2, 3), (0, 3), (3, 4), (2, 5), (4, 6)], 3, 1, {(0, 3), (2, 3)}),
+])
+def test_lineage_graph_upstream_graph(edges, selected_node, depth, expected_remaining_edges):
+    reference = LineageGraph(profile_database_name='elementary_db')
+    reference._lineage_graph.add_edges_from(edges)
+    reference._lineage_graph = reference._upstream_graph(selected_node, depth)
+    assert compare_edges(reference._lineage_graph, expected_remaining_edges)
+
+
+@pytest.mark.parametrize("profile_database_name, profile_schema_name, full_table_names, edges, selected_node, "
+                         "direction, depth, expected_remaining_edges", [
+                            ('db', 'sc', True, [('db.sc.t1', 'db.sc.t2'), ('db.sc.t1', 'db.sc.t3')],
+                             't3', 'upstream', None,
+                             {('db.sc.t1', 'db.sc.t3')}),
+                            ('db', 'sc', True, [('db.sc.t1', 'db.sc.t2'), ('db.sc.t1', 'db.sc.t3')],
+                             'db.sc.t3', 'upstream', None,
+                             {('db.sc.t1', 'db.sc.t3')}),
+                            ('db', 'sc', True, [('db.sc.t1', 'db.sc.t2'), ('db.sc.t1', 'db.sc.t3'),
+                                                ('db.sc.t3', 'db.sc.t4')],
+                             'sc.t1', 'downstream', 1,
+                             {('db.sc.t1', 'db.sc.t3'), ('db.sc.t1', 'db.sc.t2')}),
+                            ('db', 'sc', False, [('t1', 't2'), ('t1', 't3')],
+                             't3', 'upstream', None,
+                             {('t1', 't3')}),
+                            ('db', 'sc', False, [('t1', 't2'), ('t1', 't3'), ('t3', 't4'), ('t4', 't5'), ('t2', 't4'),
+                                                 ('t4', 't6'), ('t6', 't8')],
+                             't4', 'both', 1,
+                             {('t3', 't4'), ('t4', 't5'), ('t2', 't4'), ('t4', 't6')}),
+])
+def test_lineage_graph_filter_on_table(profile_database_name, profile_schema_name, full_table_names, edges,
+                                       selected_node, direction, depth, expected_remaining_edges):
+    reference = LineageGraph(profile_database_name=profile_database_name, profile_schema_name=profile_schema_name,
+                             full_table_names=full_table_names)
+    reference._lineage_graph = create_directed_graph_from_edge_list(edges)
+    reference.filter_on_table(selected_node, direction, depth)
+    assert compare_edges(reference._lineage_graph, expected_remaining_edges)
+