dbt数据转换与模型测试集成实践

import sqlite3from dataclasses import dataclass, fieldfrom typing import List, Optional @dataclassclass TestResult:    test_name: str    test_type: str    status: str  # "pass", "fail", "warn"    failed_rows: int = 0    message: str = "" class DataQualityTester:    """dbt风格的数据质量测试框架"""     def __init__(self, db_path: str):        self.conn = sqlite3.connect(db_path)        self.results: List[TestResult] = []     def test_unique(self, table: str, column: str) -> TestResult:        """唯一性测试：检查列值是否唯一"""        query = f"""            SELECT {column}, COUNT(*) as cnt            FROM {table}            GROUP BY {column}            HAVING cnt > 1        """        cursor = self.conn.execute(query)        duplicates = cursor.fetchall()        result = TestResult(            test_name=f"unique_{table}_{column}",            test_type="unique",            status="fail" if duplicates else "pass",            failed_rows=len(duplicates),            message=f"发现{len(duplicates)}个重复值" if duplicates else "唯一性检查通过"        )        self.results.append(result)        return result     def test_not_null(self, table: str, column: str) -> TestResult:        """非空测试：检查列是否包含NULL值"""        query = f"""            SELECT * FROM {table}            WHERE {column} IS NULL        """        cursor = self.conn.execute(query)        nulls = cursor.fetchall()        result = TestResult(            test_name=f"not_null_{table}_{column}",            test_type="not_null",            status="fail" if nulls else "pass",            failed_rows=len(nulls),            message=f"发现{len(nulls)}个空值" if nulls else "非空检查通过"        )        self.results.append(result)        return result     def test_relationships(self, parent_table: str, parent_column: str,                           child_table: str, child_column: str) -> TestResult:        """引用完整性测试：检查外键引用是否有效"""        query = f"""            SELECT c.{child_column}            FROM {child_table} c            LEFT JOIN {parent_table} p                ON c.{child_column} = p.{parent_column}            WHERE p.{parent_column} IS NULL              AND c.{child_column} IS NOT NULL        """        cursor = self.conn.execute(query)        orphans = cursor.fetchall()        result = TestResult(            test_name=f"fk_{child_table}_{child_column}_to_{parent_table}",            test_type="relationships",            status="fail" if orphans else "pass",            failed_rows=len(orphans),            message=f"发现{len(orphans)}条孤儿记录" if orphans else "引用完整性检查通过"        )        self.results.append(result)        return result     def test_custom_sql(self, test_name: str, sql: str,                        threshold: int = 0,                        warn_threshold: Optional[int] = None) -> TestResult:        """自定义SQL测试：返回行数超过阈值则失败"""        cursor = self.conn.execute(sql)        failed_rows = len(cursor.fetchall())        if warn_threshold is not None and failed_rows > warn_threshold:            status = "warn" if failed_rows <= threshold else "fail"        else:            status = "fail" if failed_rows > threshold else "pass"        result = TestResult(            test_name=test_name,            test_type="custom_sql",            status=status,            failed_rows=failed_rows,            message=f"自定义测试返回{failed_rows}行(阈值:{threshold})"        )        self.results.append(result)        return result     def generate_report(self) -> str:        """生成测试报告"""        total = len(self.results)        passed = sum(1 for r in self.results if r.status == "pass")        failed = sum(1 for r in self.results if r.status == "fail")        warned = sum(1 for r in self.results if r.status == "warn")        lines = [            "=" * 50,            f"数据质量测试报告 | 共{total}项",            f"通过: {passed} | 失败: {failed} | 警告: {warned}",            "=" * 50,        ]        for r in self.results:            icon = "✓" if r.status == "pass" else ("✗" if r.status == "fail" else "⚠")            lines.append(f"  {icon} [{r.test_type}] {r.test_name}: {r.message}")        return "\n".join(lines) # 使用示例if __name__ == "__main__":    tester = DataQualityTester(":memory:")    tester.conn.execute("""        CREATE TABLE users (            user_id INTEGER PRIMARY KEY,            name TEXT NOT NULL,            email TEXT,            dept_id INTEGER        )    """)    tester.conn.execute("""        CREATE TABLE departments (            dept_id INTEGER PRIMARY KEY,            dept_name TEXT NOT NULL        )    """)    tester.conn.execute("INSERT INTO departments VALUES (1, '工程部'), (2, '产品部')")    tester.conn.execute("INSERT INTO users VALUES (1, 'Alice', 'a@test.com', 1)")    tester.conn.execute("INSERT INTO users VALUES (2, 'Bob', NULL, 2)")    tester.conn.execute("INSERT INTO users VALUES (3, 'Charlie', 'c@test.com', 99)")    tester.conn.commit()    tester.test_unique("users", "user_id")    tester.test_not_null("users", "email")    tester.test_relationships("departments", "dept_id", "users", "dept_id")    tester.test_custom_sql(        "check_user_count",        "SELECT * FROM users WHERE user_id < 0",        threshold=0, warn_threshold=0    )    print(tester.generate_report())