sqlh 0.2.6__tar.gz → 0.2.8__tar.gz

{sqlh-0.2.6 → sqlh-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sqlh
-Version: 0.2.6
+Version: 0.2.8
 Summary: A lightweight SQL lineage analysis library for tracking table dependencies in data pipelines
 Keywords: sql,lineage,data-pipeline,dag,dependency,database,etl,data-engineering
 Maintainer: Perry DU

{sqlh-0.2.6 → sqlh-0.2.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlh"
-version = "0.2.6"
+version = "0.2.8"
 maintainers = [
   {name = "Perry DU", email = "duneite@gmail.com"}
 ]

{sqlh-0.2.6 → sqlh-0.2.8}/sqlh/__init__.py

@@ -11,9 +11,10 @@ from .utils import (
     search_related_tables,
     search_related_upstream_tables,
     visualize_dag,
+    table_count
 )
 
-__version__ = "0.2.6"
+__version__ = "0.2.8"
 
 __all__ = [
     "split_sql",
@@ -29,4 +30,5 @@ __all__ = [
     "search_related_tables",
     "search_command_json",
     "visualize_dag",
+    "table_count"
 ]

{sqlh-0.2.6 → sqlh-0.2.8}/sqlh/cli.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import sys
 from pathlib import Path
 
@@ -18,6 +19,7 @@ from .utils import (
     search_related_root_tables,
     search_related_tables,
     search_related_upstream_tables,
+    table_count,
     visualize_dag,
 )
 
@@ -79,6 +81,7 @@ def arg_parse():
     search_direction.add_argument("--downstream", action="store_true", help="search downstream tables (dependents)")
     search_direction.add_argument("--all", action="store_true", help="search both upstream and downstream tables")
     search_parser.add_argument("-t", "--table", help="table name to search", required=True)
+    search_parser.add_argument("-d", "--depth", help="search depth", default=None, type=int)
     search_parser.add_argument("-h", "--help", action="help", default=argparse.SUPPRESS, help="show this help message")
 
     # web 子命令
@@ -91,6 +94,15 @@ def arg_parse():
     web_parser.add_argument("--html-path", help="html file path for visualization", default=".")
     web_parser.add_argument("-h", "--help", action="help", default=argparse.SUPPRESS, help="show this help message")
 
+    # table-count 子命令
+    table_count_parser = subparsers.add_parser(
+        "table-count",
+        parents=[parent_parser],
+        help="count the number of tables",
+        add_help=False,
+    )
+    table_count_parser.add_argument("-t", "--table", help="table name to search")
+    table_count_parser.add_argument("-h", "--help", action="help", default=argparse.SUPPRESS, help="show this help message")
     return parser.parse_args()
 
 
@@ -127,10 +139,10 @@ def main():
             output = search_related_root_tables(sql_stmt_str, args.table)
             sub_command_arg = "--root"
         elif args.upstream:
-            output = search_related_upstream_tables(sql_stmt_str, args.table)
+            output = search_related_upstream_tables(sql_stmt_str, args.table, args.depth)
             sub_command_arg = "--upstream"
         elif args.downstream:
-            output = search_related_downstream_tables(sql_stmt_str, args.table)
+            output = search_related_downstream_tables(sql_stmt_str, args.table, args.depth)
             sub_command_arg = "--downstream"
         elif args.all:
             output = search_related_tables(sql_stmt_str, args.table)
@@ -151,3 +163,21 @@ def main():
         print(f"open web page: {html_file_path}")
         visualize_dag(get_all_dag(sql_stmt_str), template_type="dagre", filename=html_file_path)
         return
+
+    elif args.command == "table-count":
+        table_count_lst = table_count(sql_stmt_str, args.table)
+        if args.output_format == "text":
+            for table, count in table_count_lst:
+                print(f"{table}: {count}")
+
+        elif args.output_format == "json":
+            result = {
+                "status": "ok",
+                "command": "table-count",
+                "data": table_count_lst,
+                "meta": {"table_count": len(table_count_lst)},
+            }
+            print(json.dumps(result, indent=2, ensure_ascii=False))
+        else:
+            print(f"Error: Not Supported output format: {args.output_format}")
+            sys.exit(1)

{sqlh-0.2.6 → sqlh-0.2.8}/sqlh/core/graph.py

@@ -267,12 +267,13 @@ class DagGraph:
 
         return False
 
-    def find_upstream(self, node: str) -> FindResult:
+    def find_upstream(self, node: str, depth: int | None = None) -> FindResult:
         """
         查找所有上游依赖的边
 
         Args:
             node: 目标节点
+            depth: 最大深度，None 表示全链
 
         Returns:
             上游边的集合
@@ -280,12 +281,17 @@ class DagGraph:
         if node not in self.__nodes:
             return NodeNotFoundException(f"节点不存在:{node}")
 
-        queue = deque([node])
+        queue = deque([(node, 0)])  # (node, current_depth)
         visited = set([node])
         all_relations = []
 
         while queue:
-            current = queue.popleft()
+            # current = queue.popleft()
+            current, current_depth = queue.popleft()
+            # 检查深度限制
+            if depth is not None and current_depth >= depth:
+                continue
+
             # 使用反向邻接表直接查找上游节点，提升性能
             predecessors = self.__reverse_adjacency_list.get(current, set())
             for predecessor in predecessors:
@@ -293,15 +299,16 @@ class DagGraph:
                 all_relations.append(edge)
                 if predecessor not in visited:
                     visited.add(predecessor)
-                    queue.append(predecessor)
+                    queue.append((predecessor, current_depth + 1))
         return DagGraph(edges=all_relations)
 
-    def find_downstream(self, node: str) -> FindResult:
+    def find_downstream(self, node: str, depth: int | None = None) -> FindResult:
         """
         查找所有下游依赖的边
 
         Args:
             node: 起始节点
+            depth: 最大深度，None 表示全链
 
         Returns:
             下游边的集合
@@ -309,12 +316,17 @@ class DagGraph:
         if node not in self.__nodes:
             return NodeNotFoundException(f"节点不存在:{node}")
 
-        queue = deque([node])
+        queue = deque([(node, 0)])  # (node, current_depth)
         visited = set([node])
         all_relations = []
 
         while queue:
-            current = queue.popleft()
+            current, current_depth = queue.popleft()
+
+            # 检查深度限制
+            if depth is not None and current_depth >= depth:
+                continue
+
             # 使用邻接表直接查找下游节点，提升性能
             neighbors = self.__adjacency_list.get(current, set())
             for neighbor in neighbors:
@@ -322,7 +334,7 @@ class DagGraph:
                 all_relations.append(edge)
                 if neighbor not in visited:
                     visited.add(neighbor)
-                    queue.append(neighbor)
+                    queue.append((neighbor, current_depth + 1))
 
         return DagGraph(edges=all_relations)
 
@@ -377,9 +389,7 @@ class DagGraph:
 
         # 替换模板变量
         html_content = template.safe_substitute(
-            title="DAG Visualization",
-            mermaid_content=mermaid_content,
-            lineage_data=lineage_data
+            title="DAG Visualization", mermaid_content=mermaid_content, lineage_data=lineage_data
         )
 
         return html_content

{sqlh-0.2.6 → sqlh-0.2.8}/sqlh/tests/test_utils.py

@@ -26,7 +26,7 @@ def test_get_all_root_tables():
 
 
 def test_search_related_upstream_tables():
-    a = utils.search_related_upstream_tables(sql_stmt_str, "dws.dws_cy_cust_ltst_active_rec")
+    a = utils.search_related_upstream_tables(sql_stmt_str, "dws.xxx")
     if isinstance(a, Tuple):
         print(utils.list_command_text(a[0]))
     else:
@@ -34,7 +34,7 @@ def test_search_related_upstream_tables():
 
 
 def test_search_related_downstream_tables():
-    a = utils.search_related_downstream_tables(sql_stmt_str, "ods.ods_plr_dwm_hr_xy_shop_performance_all")
+    a = utils.search_related_downstream_tables(sql_stmt_str, "ods_hive.ods_order")
     if isinstance(a, Tuple):
         print(utils.list_command_text(a[0]))
     else:
@@ -42,7 +42,7 @@ def test_search_related_downstream_tables():
 
 
 def test_search_related_tables():
-    a = utils.search_related_tables(sql_stmt_str, "dim.dim_shopinfo")
+    a = utils.search_related_tables(sql_stmt_str, "ods_hive.ods_order")
     if isinstance(a, Tuple):
         print(utils.list_command_text(a[0]))
     else:
@@ -78,9 +78,15 @@ def test_search_command_json():
 
 
 def test_search_related_root_tables():
-    output = utils.search_related_root_tables(sql_stmt_str, "ods_hive.ods_ec_staff")
+    output = utils.search_related_root_tables(sql_stmt_str, "ods_hive.ods_order")
     # print(utils.list_command_text(output[0]))
     if isinstance(output, Tuple):
         print(utils.list_command_text(output[0]))
     else:
         print(output)
+
+
+def test_table_count():
+    print(utils.table_count(sql_stmt_str))
+    for table, count in utils.table_count(sql_stmt_str, "ods_hive.ods_order"):
+        print(f"{table}: {count}")

{sqlh-0.2.6 → sqlh-0.2.8}/sqlh/utils.py

@@ -21,6 +21,7 @@ Example:
 
 import json
 import re
+from collections import Counter
 from pathlib import Path
 from typing import List, Literal, Tuple, Union
 
@@ -197,7 +198,7 @@ def search_related_root_tables(sql_stmt_str: str, target_table: str) -> SearchRe
         return related_graph
 
 
-def search_related_upstream_tables(sql_stmt_str: str, target_table: str) -> SearchResult:
+def search_related_upstream_tables(sql_stmt_str: str, target_table: str, depth: int | None = None) -> SearchResult:
     """
     从目标表向上追溯,获取所有上游依赖表(包含所有中间表)
 
@@ -214,7 +215,7 @@ def search_related_upstream_tables(sql_stmt_str: str, target_table: str) -> Sear
         (上游表列表, 上游依赖子图)
     """
     _, _, dg = __build_tables_and_graph(sql_stmt_str)
-    related_graph = dg.find_upstream(target_table)
+    related_graph = dg.find_upstream(target_table, depth=depth)
     if isinstance(related_graph, DagGraph):
         if related_graph.empty:
             return [], DagGraph()
@@ -229,7 +230,7 @@ def search_related_upstream_tables(sql_stmt_str: str, target_table: str) -> Sear
         return related_graph
 
 
-def search_related_downstream_tables(sql_stmt_str: str, target_table: str) -> SearchResult:
+def search_related_downstream_tables(sql_stmt_str: str, target_table: str, depth: int | None = None) -> SearchResult:
     """
     从目标表向下追溯,获取所有下游依赖表(包含所有中间表)
 
@@ -246,7 +247,7 @@ def search_related_downstream_tables(sql_stmt_str: str, target_table: str) -> Se
         (下游表列表, 下游依赖子图)
     """
     _, _, dg = __build_tables_and_graph(sql_stmt_str)
-    related_graph = dg.find_downstream(target_table)
+    related_graph = dg.find_downstream(target_table, depth=depth)
     if isinstance(related_graph, DagGraph):
         if related_graph.empty:
             return [], DagGraph()
@@ -261,7 +262,7 @@ def search_related_downstream_tables(sql_stmt_str: str, target_table: str) -> Se
         return related_graph
 
 
-def search_related_tables(sql_stmt_str: str, target_table: str) -> SearchResult:
+def search_related_tables(sql_stmt_str: str, target_table: str, depth: int | None = None) -> SearchResult:
     """
     从目标表双向追溯,获取所有相关表(上游+下游,不包含自身)
 
@@ -278,9 +279,9 @@ def search_related_tables(sql_stmt_str: str, target_table: str) -> SearchResult:
         (相关表列表, 完整依赖子图)
     """
     _, _, dg = __build_tables_and_graph(sql_stmt_str)
-    related_graph_upstream = dg.find_upstream(target_table)
+    related_graph_upstream = dg.find_upstream(target_table, depth=depth)
     if isinstance(related_graph_upstream, DagGraph):
-        related_graph_downstream = dg.find_downstream(target_table)
+        related_graph_downstream = dg.find_downstream(target_table, depth=depth)
         if isinstance(related_graph_downstream, DagGraph):
             new_graph = related_graph_downstream.union(related_graph_upstream)
 
@@ -363,3 +364,15 @@ def search_command_text(search_result: SearchResult) -> str:
             return "- " + "\n- ".join(sorted(related_tables))
     else:
         return str(search_result)
+
+
+def table_count(sql_stmt_str: str, table_name: str | None = None):
+    tables = get_all_tables(sql_stmt_str)
+    counter = Counter()
+    if table_name:
+        counter[table_name] = sql_stmt_str.count(table_name.lower())
+        return [(table_name, counter[table_name])]
+    else:
+        for table in tables:
+            counter[table] += sql_stmt_str.count(table.lower())
+        return [(table, counter[table]) for table in sorted(counter, key=lambda x: counter[x], reverse=True)]