Cerebro/test_materialize_sql.py at master · MatthewRyanWeber/Cerebro · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# Integration test for cerebro.materialize_sql — materializing a raw blob table into
# the Cerebro schema inside SQL Server. Generates real one-page PDFs (PyMuPDF), so
# the PDF→text extraction is genuinely exercised; NER is faked for speed. Requires
# the local SQL Server + a 'CerebroMat' database (created by the SqlServerSource
# work). Fails loudly if SQL Server is unreachable (no silent skip).

import unittest

import sqlalchemy as sa
from sqlalchemy.engine import URL

from cerebro import materialize_sql as m

_CS = URL.create("mssql+pyodbc", host="localhost", database="CerebroMat",
                 query={"driver": "ODBC Driver 18 for SQL Server",
                        "Trusted_Connection": "yes", "TrustServerCertificate": "yes"}
                 ).render_as_string(hide_password=False)


def _pdf_blob(text):
    import fitz
    doc = fitz.open()
    page = doc.new_page()
    page.insert_text((72, 72), text, fontsize=12)
    b = doc.tobytes()
    doc.close()
    return b


class MaterializeSqlTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.eng = sa.create_engine(_CS)
        m.ensure_schema(cls.eng)
        with cls.eng.begin() as c:
            c.execute(sa.text("IF OBJECT_ID('mat_src') IS NOT NULL DROP TABLE mat_src"))
            c.execute(sa.text("CREATE TABLE mat_src(begin_bates NVARCHAR(50) PRIMARY KEY, file_bytes VARBINARY(MAX))"))
            for bid, txt in [("B001", "Alice Cooper met Bob Vance at Acme Corp in New York."),
                             ("B002", "Carol Diaz emailed Acme Corp about the contract.")]:
                c.execute(sa.text("INSERT INTO mat_src(begin_bates,file_bytes) VALUES(:i,:b)"),
                          {"i": bid, "b": _pdf_blob(txt)})
            for t in ("mentions", "dataset_entities", "entities", "documents", "materialize_state"):
                c.execute(sa.text(f"DELETE FROM {t}"))

    @classmethod
    def tearDownClass(cls):
        with cls.eng.begin() as c:
            c.execute(sa.text("IF OBJECT_ID('mat_src') IS NOT NULL DROP TABLE mat_src"))

    @staticmethod
    def _fake_ner(t):
        return {"organizations": ["Acme Corp"]} if "Acme" in t else {}

    def test_materialize_extracts_and_writes(self):
        res = m.materialize(_CS, _CS, table="mat_src", blob_col="file_bytes",
                            id_col="begin_bates", dataset="MatTest", ocr=False,
                            extract=self._fake_ner)
        self.assertEqual(res["ingested"], 2)
        with self.eng.connect() as c:
            self.assertEqual(c.execute(sa.text(
                "SELECT COUNT(*) FROM documents WHERE dataset='MatTest'")).scalar(), 2)
            body = c.execute(sa.text("SELECT body FROM documents WHERE id='B001'")).scalar()
            self.assertIn("Alice Cooper", body)          # real text from the generated PDF
            self.assertGreaterEqual(c.execute(sa.text("SELECT COUNT(*) FROM mentions")).scalar(), 2)
            # aggregates rebuilt → the entity is countable and dataset_entities populated
            self.assertEqual(c.execute(sa.text(
                "SELECT doc_count FROM entities WHERE name='Acme Corp'")).scalar(), 2)
            self.assertGreaterEqual(c.execute(sa.text(
                "SELECT COUNT(*) FROM dataset_entities WHERE dataset='MatTest'")).scalar(), 1)

    def test_resume_is_idempotent(self):
        m.materialize(_CS, _CS, table="mat_src", blob_col="file_bytes", id_col="begin_bates",
                      dataset="MatTest", ocr=False, extract=self._fake_ner)
        # second run: watermark is at the end → nothing new ingested
        res = m.materialize(_CS, _CS, table="mat_src", blob_col="file_bytes", id_col="begin_bates",
                            dataset="MatTest", ocr=False, extract=self._fake_ner)
        self.assertEqual(res["ingested"], 0)

    def test_doc_date_inline_and_backfill(self):
        with self.eng.begin() as c:
            c.execute(sa.text("IF OBJECT_ID('mat_dated') IS NOT NULL DROP TABLE mat_dated"))
            c.execute(sa.text("CREATE TABLE mat_dated(begin_bates NVARCHAR(50) PRIMARY KEY, file_bytes VARBINARY(MAX))"))
            for bid, txt in [("D001", "Memo dated March 15, 2024 regarding Acme Corp."),
                             ("D002", "No date here, just Acme Corp text.")]:
                c.execute(sa.text("INSERT INTO mat_dated(begin_bates,file_bytes) VALUES(:i,:b)"),
                          {"i": bid, "b": _pdf_blob(txt)})
        try:
            # doc_date is extracted from the body at materialize time
            m.materialize(_CS, _CS, table="mat_dated", blob_col="file_bytes", id_col="begin_bates",
                          dataset="DateTest", ocr=False, extract=self._fake_ner)
            with self.eng.connect() as c:
                self.assertEqual(c.execute(sa.text(
                    "SELECT doc_date FROM documents WHERE id='D001'")).scalar(), "2024-03-15")
            # backfill repairs a brain built before inline dates: NULL them, then re-derive
            with self.eng.begin() as c:
                c.execute(sa.text("UPDATE documents SET doc_date=NULL WHERE dataset='DateTest'"))
            res = m.backfill_doc_dates(_CS, "DateTest", batch=10)
            self.assertEqual(res["updated"], 2)
            self.assertEqual(res["dated"], 1)               # only D001 carries a date
            with self.eng.connect() as c:
                self.assertEqual(c.execute(sa.text("SELECT doc_date FROM documents WHERE id='D001'")).scalar(), "2024-03-15")
                self.assertEqual(c.execute(sa.text("SELECT doc_date FROM documents WHERE id='D002'")).scalar(), "")
            # idempotent: nothing left NULL → a second pass scans 0 rows
            self.assertEqual(m.backfill_doc_dates(_CS, "DateTest", batch=10)["updated"], 0)
        finally:
            # Don't pollute the shared schema — remove DateTest's docs/mentions so its
            # 'Acme Corp' mentions don't inflate the other test's global doc_count.
            with self.eng.begin() as c:
                c.execute(sa.text("IF OBJECT_ID('mat_dated') IS NOT NULL DROP TABLE mat_dated"))
                c.execute(sa.text("DELETE FROM mentions WHERE doc_id IN "
                                  "(SELECT id FROM documents WHERE dataset='DateTest')"))
                c.execute(sa.text("DELETE FROM dataset_entities WHERE dataset='DateTest'"))
                c.execute(sa.text("DELETE FROM documents WHERE dataset='DateTest'"))
                c.execute(sa.text("DELETE FROM materialize_state WHERE dataset='DateTest'"))


if __name__ == "__main__":
    unittest.main(verbosity=2)