332 lines
16 KiB
Python
332 lines
16 KiB
Python
|
import os
|
||
|
|
import re
|
||
|
|
import logging
|
||
|
|
import yaml
|
||
|
|
from typing import List, Dict, Any, Optional
|
||
|
|
from app.core.config import settings
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def infer_data_type(col_name: str) -> str:
|
||
|
|
"""Infer standard Lakehouse data type for dbt contract enforcement."""
|
||
|
|
name = col_name.lower().strip()
|
||
|
|
if name.endswith("_date") or name == "date":
|
||
|
|
return "date"
|
||
|
|
if name.endswith("_at") or "timestamp" in name or name.endswith("_time"):
|
||
|
|
return "timestamp with time zone"
|
||
|
|
if any(k in name for k in ["score", "rate", "amount", "percent", "price", "val", "total", "count"]):
|
||
|
|
return "numeric"
|
||
|
|
if name.endswith("_id") or name.endswith("_number") or name == "id" or "code" in name or "status" in name or "type" in name:
|
||
|
|
return "character varying(50)"
|
||
|
|
return "character varying(255)"
|
||
|
|
|
||
|
|
|
||
|
|
def generate_model_schema_dict(
|
||
|
|
model_name: str,
|
||
|
|
test_configs: Optional[Dict[str, Any]] = None,
|
||
|
|
sql_query: Optional[str] = None
|
||
|
|
) -> Dict[str, Any]:
|
||
|
|
"""Generate dbt schema dictionary for a model, configuring tests and contract enforcement."""
|
||
|
|
clean_name = model_name.replace(".sql", "")
|
||
|
|
cfg = test_configs or {}
|
||
|
|
|
||
|
|
unique_cols = cfg.get("unique_columns") or cfg.get("unique_cols") or []
|
||
|
|
if isinstance(unique_cols, str):
|
||
|
|
unique_cols = [c.strip() for c in unique_cols.split(",") if c.strip()]
|
||
|
|
|
||
|
|
not_null_cols = cfg.get("not_null_columns") or cfg.get("not_null_cols") or []
|
||
|
|
if isinstance(not_null_cols, str):
|
||
|
|
not_null_cols = [c.strip() for c in not_null_cols.split(",") if c.strip()]
|
||
|
|
|
||
|
|
accepted_values = cfg.get("accepted_values") or {}
|
||
|
|
relationships = cfg.get("relationships") or {}
|
||
|
|
enforce_contract = cfg.get("enforce_contract", False)
|
||
|
|
column_types = cfg.get("column_types") or {}
|
||
|
|
|
||
|
|
# Collect all configured columns
|
||
|
|
all_cols = set(unique_cols) | set(not_null_cols) | set(accepted_values.keys()) | set(relationships.keys()) | set(column_types.keys())
|
||
|
|
|
||
|
|
# If contract is enforced, extract all output columns from SQL query to satisfy dbt contract rules
|
||
|
|
effective_sql = sql_query or cfg.get("sql_query") or cfg.get("query") or ""
|
||
|
|
if enforce_contract and effective_sql:
|
||
|
|
m = re.search(r'SELECT\s+(.*?)\s+FROM', effective_sql, re.IGNORECASE | re.DOTALL)
|
||
|
|
if m:
|
||
|
|
for raw_col in m.group(1).split(","):
|
||
|
|
raw_col = raw_col.strip()
|
||
|
|
if not raw_col:
|
||
|
|
continue
|
||
|
|
if " as " in raw_col.lower():
|
||
|
|
col_id = re.split(r'\s+as\s+', raw_col, flags=re.IGNORECASE)[-1].strip().strip('"').strip("'")
|
||
|
|
else:
|
||
|
|
col_id = raw_col.split(".")[-1].strip().strip('"').strip("'")
|
||
|
|
if col_id and re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', col_id):
|
||
|
|
all_cols.add(col_id)
|
||
|
|
|
||
|
|
columns_list = []
|
||
|
|
for col in sorted(all_cols):
|
||
|
|
col_tests = []
|
||
|
|
if col in unique_cols:
|
||
|
|
col_tests.append("unique")
|
||
|
|
if col in not_null_cols:
|
||
|
|
col_tests.append("not_null")
|
||
|
|
if col in accepted_values and accepted_values[col]:
|
||
|
|
vals = accepted_values[col]
|
||
|
|
if isinstance(vals, str):
|
||
|
|
vals = [v.strip() for v in vals.split(",") if v.strip()]
|
||
|
|
col_tests.append({"accepted_values": {"values": vals}})
|
||
|
|
if col in relationships and relationships[col]:
|
||
|
|
rel_info = relationships[col]
|
||
|
|
if isinstance(rel_info, dict):
|
||
|
|
col_tests.append({"relationships": {
|
||
|
|
"to": rel_info.get("to", "ref('upstream')"),
|
||
|
|
"field": rel_info.get("field", "id")
|
||
|
|
}})
|
||
|
|
|
||
|
|
col_dict: Dict[str, Any] = {"name": col}
|
||
|
|
if col_tests:
|
||
|
|
col_dict["tests"] = col_tests
|
||
|
|
if col in column_types:
|
||
|
|
col_dict["data_type"] = column_types[col]
|
||
|
|
elif enforce_contract:
|
||
|
|
col_dict["data_type"] = infer_data_type(col)
|
||
|
|
columns_list.append(col_dict)
|
||
|
|
|
||
|
|
model_entry: Dict[str, Any] = {
|
||
|
|
"name": clean_name,
|
||
|
|
"description": f"dbt Model {clean_name} auto-registered with Shift-Left quality assertions"
|
||
|
|
}
|
||
|
|
|
||
|
|
if enforce_contract:
|
||
|
|
model_entry["config"] = {"contract": {"enforced": True}}
|
||
|
|
if columns_list:
|
||
|
|
model_entry["columns"] = columns_list
|
||
|
|
|
||
|
|
return model_entry
|
||
|
|
|
||
|
|
|
||
|
|
def compile_dbt_model(
|
||
|
|
model_name: str,
|
||
|
|
layer: str = "silver",
|
||
|
|
sql_query: str = "",
|
||
|
|
test_configs: Optional[Dict[str, Any]] = None
|
||
|
|
) -> Dict[str, Any]:
|
||
|
|
"""
|
||
|
|
Perform pre-flight compilation check for dbt models:
|
||
|
|
- Verifies Jinja delimiters pairing.
|
||
|
|
- Extracts {{ ref('...') }} and {{ source('...', '...') }}.
|
||
|
|
- Compiles Jinja refs into simulated physical lakehouse identifiers.
|
||
|
|
- Validates SQL structure (SELECT/WITH statements, parentheses matching).
|
||
|
|
- Generates corresponding dbt schema.yml definition with quality tests.
|
||
|
|
"""
|
||
|
|
clean_name = (model_name or "model").replace(".sql", "")
|
||
|
|
errors: List[str] = []
|
||
|
|
warnings: List[str] = []
|
||
|
|
|
||
|
|
if not sql_query or not sql_query.strip():
|
||
|
|
errors.append("SQL query cannot be empty.")
|
||
|
|
|
||
|
|
# 1. Jinja delimiter syntax checking
|
||
|
|
open_brackets = len(re.findall(r'\{\{', sql_query))
|
||
|
|
close_brackets = len(re.findall(r'\}\}', sql_query))
|
||
|
|
if open_brackets != close_brackets:
|
||
|
|
errors.append(f"Mismatched Jinja expression delimiters: {open_brackets} '{{{{' vs {close_brackets} '}}}}'.")
|
||
|
|
|
||
|
|
open_blocks = len(re.findall(r'\{%', sql_query))
|
||
|
|
close_blocks = len(re.findall(r'%\}', sql_query))
|
||
|
|
if open_blocks != close_blocks:
|
||
|
|
errors.append(f"Mismatched Jinja block delimiters: {open_blocks} '{{%' vs {close_blocks} '%}}'.")
|
||
|
|
|
||
|
|
# 2. Extract references & sources
|
||
|
|
extracted_refs = re.findall(r"ref\(\s*['\"]([^'\"]+)['\"]\s*\)", sql_query)
|
||
|
|
extracted_sources = re.findall(r"source\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]([^'\"]+)['\"]\s*\)", sql_query)
|
||
|
|
|
||
|
|
# 3. Compile simulated SQL (substituting Jinja with lakehouse catalogue paths)
|
||
|
|
catalog = os.getenv("ICEBERG_CATALOG_NAME", "dlh_lakehouse")
|
||
|
|
compiled_sql = sql_query
|
||
|
|
|
||
|
|
for ref in extracted_refs:
|
||
|
|
escaped_ref = re.escape(ref)
|
||
|
|
replacement = f"{catalog}.{layer}.{ref}" if not ref.startswith("bronze") else f"{catalog}.bronze.{ref}"
|
||
|
|
pattern = re.compile(r"\{\{\s*ref\(\s*['\"]" + escaped_ref + r"['\"]\s*\)\s*\}\}")
|
||
|
|
compiled_sql = pattern.sub(replacement, compiled_sql)
|
||
|
|
|
||
|
|
for src_schema, src_tbl in extracted_sources:
|
||
|
|
escaped_schema = re.escape(src_schema)
|
||
|
|
escaped_tbl = re.escape(src_tbl)
|
||
|
|
replacement = f"{catalog}.bronze.{src_tbl}"
|
||
|
|
pattern = re.compile(r"\{\{\s*source\(\s*['\"]" + escaped_schema + r"['\"]\s*,\s*['\"]" + escaped_tbl + r"['\"]\s*\)\s*\}\}")
|
||
|
|
compiled_sql = pattern.sub(replacement, compiled_sql)
|
||
|
|
|
||
|
|
# Clean up dbt config() macro in compiled preview
|
||
|
|
compiled_sql = re.sub(r"\{\{\s*config\([^)]*\)\s*\}\}", "-- [dbt config applied]", compiled_sql)
|
||
|
|
|
||
|
|
# 4. Basic SQL structural validation
|
||
|
|
clean_sql = re.sub(r'--.*', '', compiled_sql).strip()
|
||
|
|
if clean_sql:
|
||
|
|
if not re.search(r'\b(select|with)\b', clean_sql, re.IGNORECASE):
|
||
|
|
warnings.append("SQL query does not start with SELECT or WITH clause; ensure it produces a valid dataset.")
|
||
|
|
if clean_sql.count('(') != clean_sql.count(')'):
|
||
|
|
errors.append("Unbalanced parentheses '(' and ')' detected in SQL query.")
|
||
|
|
|
||
|
|
# 5. Generate Schema YAML testing representation
|
||
|
|
schema_dict = generate_model_schema_dict(clean_name, test_configs, sql_query)
|
||
|
|
wrapper = {"version": 2, "models": [schema_dict]}
|
||
|
|
schema_yml_str = yaml.dump(wrapper, sort_keys=False, default_flow_style=False)
|
||
|
|
|
||
|
|
is_valid = len(errors) == 0
|
||
|
|
return {
|
||
|
|
"status": "success" if is_valid else "error",
|
||
|
|
"valid": is_valid,
|
||
|
|
"model_name": clean_name,
|
||
|
|
"layer": layer,
|
||
|
|
"compiled_sql": compiled_sql.strip(),
|
||
|
|
"schema_yml": schema_yml_str,
|
||
|
|
"extracted_refs": extracted_refs,
|
||
|
|
"extracted_sources": [{"schema": s[0], "table": s[1]} for s in extracted_sources],
|
||
|
|
"warnings": warnings,
|
||
|
|
"errors": errors
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def sync_dbt_models_from_nodes(nodes: List[Dict[str, Any]]) -> List[str]:
|
||
|
|
"""Sync dbt SQL model files, schema.yml (with tests), and sources.yml physically in both dbt target directories."""
|
||
|
|
synced_files = []
|
||
|
|
|
||
|
|
target_dbt_dirs = [settings.DBT_DIR]
|
||
|
|
if hasattr(settings, "DEFAULT_DBT_DIR") and settings.DEFAULT_DBT_DIR and settings.DEFAULT_DBT_DIR != settings.DBT_DIR:
|
||
|
|
target_dbt_dirs.append(settings.DEFAULT_DBT_DIR)
|
||
|
|
|
||
|
|
for base_dbt_dir in target_dbt_dirs:
|
||
|
|
try:
|
||
|
|
models_dir = os.path.join(base_dbt_dir, "models")
|
||
|
|
os.makedirs(models_dir, exist_ok=True)
|
||
|
|
|
||
|
|
models_by_layer: Dict[str, List[Dict[str, Any]]] = {}
|
||
|
|
|
||
|
|
# 1. Sync dbt SQL models
|
||
|
|
for node in nodes:
|
||
|
|
node_type = node.get("data", {}).get("type", "")
|
||
|
|
cfg = node.get("data", {}).get("config", {})
|
||
|
|
if node_type == "dbt_transform":
|
||
|
|
layer = (cfg.get("dbt_layer") or cfg.get("sel") or "silver").lower()
|
||
|
|
model_name = cfg.get("dbt_model") or cfg.get("sel") or node.get("data", {}).get("label", "").lower().replace(" ", "_")
|
||
|
|
if not model_name:
|
||
|
|
continue
|
||
|
|
if not model_name.endswith(".sql"):
|
||
|
|
model_name += ".sql"
|
||
|
|
|
||
|
|
target_dir = os.path.join(models_dir, layer)
|
||
|
|
os.makedirs(target_dir, exist_ok=True)
|
||
|
|
|
||
|
|
sql_content = cfg.get("query") or cfg.get("sql_query_preview") or f"-- dbt Model: {model_name}\nSELECT * FROM {{{{ ref('bronze_raw') }}}}"
|
||
|
|
model_path = os.path.join(target_dir, model_name)
|
||
|
|
with open(model_path, "w", encoding="utf-8") as f:
|
||
|
|
f.write(sql_content)
|
||
|
|
logger.info("Auto-synced dbt model file to %s", model_path)
|
||
|
|
synced_files.append(model_path)
|
||
|
|
|
||
|
|
# Collect for schema.yml tests
|
||
|
|
test_cfgs = cfg.get("test_configs") or cfg.get("quality_tests")
|
||
|
|
if not test_cfgs and any(k in cfg for k in ["unique_columns", "not_null_columns", "accepted_values", "enforce_contract"]):
|
||
|
|
test_cfgs = {
|
||
|
|
"unique_columns": cfg.get("unique_columns") or [],
|
||
|
|
"not_null_columns": cfg.get("not_null_columns") or [],
|
||
|
|
"accepted_values": cfg.get("accepted_values") or {},
|
||
|
|
"enforce_contract": cfg.get("enforce_contract") or False,
|
||
|
|
}
|
||
|
|
if test_cfgs:
|
||
|
|
model_entry = generate_model_schema_dict(model_name, test_cfgs, sql_content)
|
||
|
|
models_by_layer.setdefault(layer, []).append(model_entry)
|
||
|
|
|
||
|
|
# 2. Auto-sync schema.yml per layer if test_configs are present
|
||
|
|
for layer, m_list in models_by_layer.items():
|
||
|
|
if not m_list:
|
||
|
|
continue
|
||
|
|
schema_file = os.path.join(models_dir, layer, "schema.yml")
|
||
|
|
existing_schema: Dict[str, Any] = {"version": 2, "models": []}
|
||
|
|
if os.path.exists(schema_file):
|
||
|
|
try:
|
||
|
|
with open(schema_file, "r", encoding="utf-8") as sf:
|
||
|
|
existing_schema = yaml.safe_load(sf) or {"version": 2, "models": []}
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
existing_models = {m.get("name"): m for m in existing_schema.get("models", [])}
|
||
|
|
for new_m in m_list:
|
||
|
|
existing_models[new_m["name"]] = new_m
|
||
|
|
existing_schema["models"] = list(existing_models.values())
|
||
|
|
|
||
|
|
with open(schema_file, "w", encoding="utf-8") as sf:
|
||
|
|
yaml.safe_dump(existing_schema, sf, sort_keys=False, default_flow_style=False)
|
||
|
|
logger.info("Auto-synced dbt schema tests to %s", schema_file)
|
||
|
|
synced_files.append(schema_file)
|
||
|
|
|
||
|
|
# 3. Auto-sync sources.yml from source_db nodes
|
||
|
|
sources_file = os.path.join(models_dir, "sources.yml")
|
||
|
|
source_tables_by_schema = {}
|
||
|
|
for node in nodes:
|
||
|
|
node_type = node.get("data", {}).get("type", "")
|
||
|
|
cfg = node.get("data", {}).get("config", {})
|
||
|
|
if node_type in ["source_db", "database"]:
|
||
|
|
tbl = cfg.get("table_name", "")
|
||
|
|
if tbl:
|
||
|
|
parts = tbl.split(".")
|
||
|
|
s_name = parts[0] if len(parts) > 1 else "public"
|
||
|
|
t_name = parts[1] if len(parts) > 1 else parts[0]
|
||
|
|
source_tables_by_schema.setdefault(s_name, set()).add(t_name)
|
||
|
|
|
||
|
|
if source_tables_by_schema and os.path.exists(sources_file):
|
||
|
|
try:
|
||
|
|
with open(sources_file, "r", encoding="utf-8") as sf:
|
||
|
|
src_data = yaml.safe_load(sf) or {}
|
||
|
|
|
||
|
|
sources_list = src_data.get("sources", [])
|
||
|
|
modified = False
|
||
|
|
for schema_name, tbl_set in source_tables_by_schema.items():
|
||
|
|
schema_entry = next((s for s in sources_list if s.get("schema") == schema_name), None)
|
||
|
|
if not schema_entry:
|
||
|
|
schema_entry = {"name": schema_name, "schema": schema_name, "tables": []}
|
||
|
|
sources_list.append(schema_entry)
|
||
|
|
modified = True
|
||
|
|
|
||
|
|
existing_tables = {t.get("name") for t in schema_entry.get("tables", [])}
|
||
|
|
for t in tbl_set:
|
||
|
|
if t not in existing_tables:
|
||
|
|
schema_entry.setdefault("tables", []).append({
|
||
|
|
"name": t,
|
||
|
|
"description": f"Auto-registered source table '{t}' from canvas pipeline"
|
||
|
|
})
|
||
|
|
modified = True
|
||
|
|
|
||
|
|
if modified:
|
||
|
|
src_data["sources"] = sources_list
|
||
|
|
with open(sources_file, "w", encoding="utf-8") as sf:
|
||
|
|
yaml.safe_dump(src_data, sf, sort_keys=False, default_flow_style=False)
|
||
|
|
logger.info("Auto-synced sources.yml at %s", sources_file)
|
||
|
|
except Exception as se:
|
||
|
|
logger.warning("Auto-sync sources.yml warning: %s", se)
|
||
|
|
|
||
|
|
# 4. Auto-sync bronze view models from source_db tables so downstream models can ref() them
|
||
|
|
bronze_dir = os.path.join(models_dir, "bronze")
|
||
|
|
os.makedirs(bronze_dir, exist_ok=True)
|
||
|
|
for schema_name, tbl_set in source_tables_by_schema.items():
|
||
|
|
for t_name in tbl_set:
|
||
|
|
m_name = f"bronze_{t_name}.sql" if not t_name.startswith("bronze_") else f"{t_name}.sql"
|
||
|
|
m_path = os.path.join(bronze_dir, m_name)
|
||
|
|
if not os.path.exists(m_path):
|
||
|
|
b_content = (
|
||
|
|
f"-- dbt Bronze View Model: {t_name}\n"
|
||
|
|
f"{{{{ config(materialized='view') }}}}\n\n"
|
||
|
|
f"SELECT * FROM {{{{ source('{schema_name}', '{t_name}') }}}}\n"
|
||
|
|
)
|
||
|
|
with open(m_path, "w", encoding="utf-8") as bf:
|
||
|
|
bf.write(b_content)
|
||
|
|
logger.info("Auto-synced bronze dbt model to %s", m_path)
|
||
|
|
synced_files.append(m_path)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning("dbt Model Sync Note: %s", e)
|
||
|
|
|
||
|
|
return synced_files
|