Example Script: Get table names used by data source

redash_toolbelt/examples/find-table-names.py

@@ -1,16 +1,9 @@
 import itertools, json, re
 import click
+import pytest
 from redash_toolbelt import Redash
 
 
-# This regex captures three groups:
-#
-#   0. A FROM or JOIN statement
-#   1. The whitespace character between FROM/JOIN and table name
-#   2. The table name
-PATTERN = re.compile(r"(?:FROM|JOIN)(?: )([^\s\(\)]+)", flags=re.IGNORECASE)
-
-
 def find_table_names(url, key, data_source_id):
 
     client = Redash(url, key)
@@ -29,18 +22,31 @@ def find_table_names(url, key, data_source_id):
     ]
 
     tables_by_qry = {
-        query["id"]: [
-            match
-            for match in re.findall(PATTERN, query["query"])
-            if match in schema_tables or len(schema_tables) == 0
-        ]
+        query["id"]: extract_table_names(query["query"], schema_tables)
         for query in queries
-        if re.search(PATTERN, query["query"])
     }
 
     return tables_by_qry
 
 
+def extract_table_names(str_sql, schema_tables=[]):
+
+    # This regex captures three groups:
+    #
+    #   0. A FROM or JOIN statement
+    #   1. The whitespace character(s) between FROM/JOIN and table name
+    #   2. The table name
+    PATTERN = re.compile(
+        r"(?:FROM|JOIN)(?:\s+)([^\s\(\)]+)", flags=re.IGNORECASE | re.UNICODE
+    )
+
+    return [
+        match
+        for match in re.findall(PATTERN, str_sql)
+        if match in schema_tables or len(schema_tables) == 0
+    ]
+
+
 def print_summary(tables_by_qry):
     """Builds a summary showing table names and count of queries that reference them."""
 
@@ -92,3 +98,68 @@ def main(url, key, data_source_id, detail):
 
 if __name__ == "__main__":
     main()
+
+def test_1():
+
+    sql = """
+    SELECT field FROM table0 LEFT JOIN table1 ON table0.field = table1.field
+    """
+
+    tables = extract_table_names(sql)
+
+    assert tables == ['table0', 'table1']
+
+def test_2():
+
+    sql = """
+    SELECT field FROM table0 as a LEFT JOIN table1 as b ON a.field = b.field
+    """
+
+    tables = extract_table_names(sql)
+
+    assert tables == ['table0', 'table1']
+
+def test_3():
+
+    sql = """
+    SELECT field FROM table0 a LEFT JOIN table1 b ON a.field = b.field
+    """
+
+    tables = extract_table_names(sql)
+
+    assert tables == ['table0', 'table1']
+
+def test_4():
+
+    sql = """
+    SELECT field FROM schema.table0 a LEFT JOIN schema.table1 b ON a.field = b.field
+    """
+
+    tables = extract_table_names(sql)
+
+    assert tables == ['schema.table0', 'schema.table1']
+
+def test_5():
+
+    sql = """
+    SELECT field
+    FROM
+        table0
+    LEFT JOIN
+        table1
+    """
+
+    tables = extract_table_names(sql)
+
+    assert tables == ['table0', 'table1']
+
+def test_6():
+
+    sql = """
+    SELECT field FROM table1,table0
+    WHERE table0.field = table1.field
+    """
+
+    tables = extract_table_names(sql)
+
+    assert tables == ['schema.table0', 'schema.table1']