tdfs4ds 0.2.5.4__py3-none-any.whl → 0.2.5.5__py3-none-any.whl

tdfs4ds/lineage/lineage.py

@@ -0,0 +1,409 @@
+import tdfs4ds
+from tdfs4ds import logger_safe
+
+import re
+from typing import Dict, List, Set, Iterable
+
+
+_COMMENT_SINGLE = re.compile(r"--[^\n]*")
+_COMMENT_MULTI = re.compile(r"/\*.*?\*/", re.DOTALL | re.MULTILINE)
+
+# --- add near the top with other regexes ---
+
+_RE_CREATE_TABLE_AS = re.compile(
+    r"""
+    \bCREATE\s+
+    (?:MULTISET\s+|SET\s+|VOLATILE\s+|TEMP(?:ORARY)?\s+)?   # qualifiers
+    TABLE\s+
+    (?P<name>[\w".]+)
+    \s+AS\s*\(                                              # <-- require AS (
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+
+_RE_SINGLE_QUOTED_STRING = re.compile(
+    r"""
+    '
+    (?:''|[^'])*   # escaped quote or any non-quote char
+    '
+    """,
+    re.VERBOSE | re.DOTALL,
+)
+
+def _mask_single_quoted_strings(sql: str) -> str:
+    """
+    Replace all single-quoted SQL string literals with empty string ''.
+    Preserves SQL structure while removing misleading keywords.
+    """
+    return _RE_SINGLE_QUOTED_STRING.sub("''", sql)
+
+_RE_REPLACE_VIEW_AS = re.compile(
+    r"""
+    \bREPLACE\s+VIEW\s+
+    (?P<name>[\w".]+)
+    \s+AS\b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_CREATE_OR_REPLACE_VIEW_AS = re.compile(
+    r"""
+    \b(?:CREATE|REPLACE)\s+VIEW\s+
+    (?P<name>[\w".]+)
+    \s+AS\b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_MERGE_INTO = re.compile(
+    r"""
+    \bMERGE\s+INTO\s+
+    (?P<name>[\w".]+)
+    \b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_UPDATE = re.compile(
+    r"""
+    \bUPDATE\s+
+    (?!SET\b)                 # <-- prevents UPDATE SET from matching
+    (?P<name>[\w".]+)
+    \b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+
+_RE_USING = re.compile(
+    r"""
+    \bUSING\s+(?P<name>[\w".]+)\b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_DELETE_FROM = re.compile(
+    r"""
+    \bDELETE\s+FROM\s+
+    (?P<name>[\w".]+)
+    \b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+# Capture ON <identifier> and ON (<subquery>) inside TD_* table functions
+_RE_TD_FUNC_CALL = re.compile(
+    r"""\bTD_[A-Z_0-9]+\s*\(""",
+    re.IGNORECASE,
+)
+
+_RE_TD_ON_CLAUSE = re.compile(
+    r"""
+    \bON\b\s+
+    (?P<arg>
+        \([^\)]*\)          # naive: one-level parens (good enough for your test cases)
+        |
+        [\w".]+             # identifier
+    )
+    """,
+    re.IGNORECASE | re.VERBOSE | re.DOTALL,
+)
+
+
+_RE_INSERT_INTO = re.compile(
+    r"""
+    \bINSERT\s+INTO\s+
+    (?P<name>[\w".]+)
+    \b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_CREATE_VIEW_AS = re.compile(
+    r"""
+    \b(?:CREATE|REPLACE)\s+VIEW\s+
+    (?P<name>[\w".]+)
+    \s+AS\b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_FROM_JOIN = re.compile(
+    r"""
+    \bFROM\b
+    \s+(?P<name>[\w".]+)
+    |
+    \b(?:INNER\s+JOIN|LEFT\s+JOIN|RIGHT\s+JOIN|FULL\s+OUTER\s+JOIN|CROSS\s+JOIN|JOIN)\b
+    \s+(?P<join_name>[\w".]+)
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_CTE_NAMES = re.compile(
+    r"""
+    \bWITH\b\s*(?P<first>\w+)\s+AS\s*\(
+    |,\s*(?P<next>\w+)\s+AS\s*\(
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+_RE_EXTRACT_FROM = re.compile(
+    r"""\bEXTRACT\s*\(\s*[^()]*?\bFROM\b[^()]*?\)""",
+    re.IGNORECASE | re.DOTALL,
+)
+
+_RE_USING_TABLE = re.compile(
+    r"""
+    \bUSING\s+(?P<name>[\w".]+)\b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+def _extract_using_tables(sql: str) -> List[str]:
+    out: List[str] = []
+    for m in _RE_USING_TABLE.finditer(sql):
+        name = m.group("name")
+        # Look ahead: if next non-space char is '(' then it's a TD_* USING option, not a table
+        j = m.end()
+        while j < len(sql) and sql[j].isspace():
+            j += 1
+        if j < len(sql) and sql[j] == "(":
+            continue
+        out.append(name)
+    return out
+
+
+_RE_TD_FUNC_CALL = re.compile(r"\bTD_[A-Z_0-9]+\s*\(", re.IGNORECASE)
+
+def _extract_td_function_inputs(sql: str) -> List[str]:
+    inputs: List[str] = []
+    for m in _RE_TD_FUNC_CALL.finditer(sql):
+        parens = _extract_balanced_parens(sql, m.end() - 1)  # points to '('
+        inner = parens[1:-1]  # inside TD_*( ... )
+
+        for onm in _RE_TD_ON_CLAUSE.finditer(inner):
+            arg = onm.group("arg").strip().rstrip(",);")
+
+            if arg.startswith("(") and arg.endswith(")"):
+                sub = arg[1:-1]
+                for sm in _RE_FROM_JOIN.finditer(sub):
+                    nm = sm.group("name") or sm.group("join_name")
+                    if nm:
+                        inputs.append(nm)
+            else:
+                inputs.append(arg)
+    return inputs
+
+
+
+def _mask_extract_from(sql: str) -> str:
+    # Replace EXTRACT( ... FROM ... ) with a placeholder that contains no FROM keyword
+    return _RE_EXTRACT_FROM.sub("EXTRACT(/*masked*/)", sql)
+
+_DEFAULT_DB = "<DEFAULT_DATABASE>"
+
+def _is_placeholder(part: str) -> bool:
+    return part.strip().upper() == _DEFAULT_DB.upper()
+
+def _quote_identifier(name: str) -> str:
+    name = _strip_alias(name)
+
+    # already quoted identifier: leave as-is
+    if '"' in name:
+        return name
+
+    parts = name.split(".")
+    quoted_parts = []
+    for p in parts:
+        if p == _DEFAULT_DB:
+            quoted_parts.append(_DEFAULT_DB)   # <-- do NOT quote placeholder
+        else:
+            quoted_parts.append(f'"{p}"')
+    return ".".join(quoted_parts)
+
+def _extract_balanced_parens(sql: str, open_paren_index: int) -> str:
+    """
+    Return the substring starting at open_paren_index (which must point to '(')
+    up to and including the matching closing ')'.
+    Falls back to end-of-string if unbalanced.
+    """
+    depth = 0
+    for i in range(open_paren_index, len(sql)):
+        ch = sql[i]
+        if ch == "(":
+            depth += 1
+        elif ch == ")":
+            depth -= 1
+            if depth == 0:
+                return sql[open_paren_index : i + 1]
+    return sql[open_paren_index:]
+
+
+def _qualify_default_db(name: str, default_db: str) -> str:
+    name = _strip_alias(name)
+    if "." in name or name.startswith('"') or name.startswith(default_db):
+        return name
+    return f"{default_db}.{name}"
+
+
+
+def _remove_sql_comments(sql: str) -> str:
+    sql = _COMMENT_SINGLE.sub("", sql)
+    sql = _COMMENT_MULTI.sub("", sql)
+    return sql
+
+
+def _extract_cte_names(sql: str) -> Set[str]:
+    names: Set[str] = set()
+    for m in _RE_CTE_NAMES.finditer(sql):
+        name = m.group("first") or m.group("next")
+        if name:
+            names.add(name.lower())
+    return names
+
+
+def _strip_alias(name: str) -> str:
+    return name.strip().split()[0]
+
+
+_DEFAULT_DB = "<DEFAULT_DATABASE>"
+
+def _quote_identifier(name: str) -> str:
+    name = _strip_alias(name)
+    if '"' in name:
+        return name
+
+    parts = name.split(".")
+    out = []
+    for p in parts:
+        if p == _DEFAULT_DB:
+            out.append(_DEFAULT_DB)          # keep placeholder unquoted
+        else:
+            out.append(f'"{p}"')
+    return ".".join(out)
+
+
+
+def _dedupe_preserve_order(items: Iterable[str]) -> List[str]:
+    seen: Set[str] = set()
+    out: List[str] = []
+    for x in items:
+        if x not in seen:
+            out.append(x)
+            seen.add(x)
+    return out
+
+_RE_TD_FUNC_CALL = re.compile(r"\bTD_[A-Z_0-9]+\s*\(", re.IGNORECASE)
+
+def _extract_td_function_inputs(sql: str) -> List[str]:
+    inputs: List[str] = []
+    for m in _RE_TD_FUNC_CALL.finditer(sql):
+        # m.end()-1 is the '('
+        parens = _extract_balanced_parens(sql, m.end() - 1)
+        inner = parens[1:-1]  # strip outer ()
+        for onm in _RE_TD_ON_CLAUSE.finditer(inner):
+            arg = onm.group("arg").strip()
+            if arg.startswith("(") and arg.endswith(")"):
+                sub = arg[1:-1]
+                for sm in _RE_FROM_JOIN.finditer(sub):
+                    nm = sm.group("name") or sm.group("join_name")
+                    if nm:
+                        inputs.append(nm)
+            else:
+                inputs.append(arg)
+    return inputs
+
+
+
+def analyze_sql_query(sql_query: str, *, debug: bool = False) -> Dict[str, List[str]]:
+    """
+    Extracts source and target tables/views from a SQL query.
+    """
+    sql = _remove_sql_comments(sql_query)
+    sql = _mask_single_quoted_strings(sql)
+    sql = _mask_extract_from(sql)  # <-- important fix for EXTRACT(... FROM ...)
+    cte_names = _extract_cte_names(sql)
+
+    targets_raw: List[str] = []
+    sources_raw: List[str] = []
+
+    # DELETE ... USING <table>
+    sources_raw += _extract_using_tables(sql)
+
+
+    # --- UPDATE targets (and treat them as sources too) ---
+    # IMPORTANT: _RE_UPDATE must be defined with (?!SET\b) to avoid matching MERGE's "UPDATE SET"
+    update_targets = [m.group("name") for m in _RE_UPDATE.finditer(sql)]
+    targets_raw += update_targets
+    sources_raw += update_targets  # <-- this is the "UPDATE management" you need
+
+    # --- targets ---
+    targets_raw += [m.group("name") for m in _RE_CREATE_TABLE_AS.finditer(sql)]
+    targets_raw += [m.group("name") for m in _RE_INSERT_INTO.finditer(sql)]
+    targets_raw += [m.group("name") for m in _RE_CREATE_OR_REPLACE_VIEW_AS.finditer(sql)]
+    targets_raw += [m.group("name") for m in _RE_REPLACE_VIEW_AS.finditer(sql)]
+    targets_raw += [m.group("name") for m in _RE_MERGE_INTO.finditer(sql)]
+    targets_raw += [m.group("name") for m in _RE_DELETE_FROM.finditer(sql)]
+
+    # --- sources: FROM/JOIN ---
+    for m in _RE_FROM_JOIN.finditer(sql):
+        name = (m.group("name") or m.group("join_name") or "").strip()
+        if not name:
+            continue
+
+        # If FROM/JOIN references a TD_* function call (TD_UNPIVOT(...)), don't treat TD_* as a table
+        end = m.end()
+        j = end
+        while j < len(sql) and sql[j].isspace():
+            j += 1
+        if name.upper().startswith("TD_") and j < len(sql) and sql[j] == "(":
+            continue
+
+        # Skip derived tables like FROM (SELECT ...) alias
+        if name.startswith("("):
+            continue
+
+        sources_raw.append(name)
+
+    # --- sources: TD_* ON inputs ---
+    sources_raw += _extract_td_function_inputs(sql)
+
+    # If you also have a generic USING regex, keep it (but you already have _RE_USING_TABLE above)
+    # sources_raw += [m.group("name") for m in _RE_USING.finditer(sql)]
+
+    if debug:
+        logger_safe("debug", "CTE names detected: %s", sorted(cte_names))
+        logger_safe("debug", "Raw target matches: %s", targets_raw)
+        logger_safe("debug", "Raw source matches: %s", sources_raw)
+
+    def normalize_and_filter(names: Iterable[str], *, qualify_default: bool) -> List[str]:
+        out: List[str] = []
+        for n in names:
+            n = _strip_alias(n)
+            n = n.rstrip(",);")
+
+            if not n:
+                continue
+
+            # ignore CTEs
+            if n.lower() in cte_names:
+                continue
+
+            if qualify_default:
+                n = _qualify_default_db(n, _DEFAULT_DB)
+
+            out.append(_quote_identifier(n))
+        return _dedupe_preserve_order(out)
+
+    result = {
+        "source": normalize_and_filter(sources_raw, qualify_default=True),
+        "target": normalize_and_filter(targets_raw, qualify_default=True),
+    }
+
+    if debug:
+        logger_safe("debug", "Final parsed result: %s", result)
+
+    return result
+
+