29787659139

Committed 20 Jul 2026 11:34PM UTC coverage: 51.578% (-1.1%) from 52.66%

Build # 29787659139

Build Type

Pull #99

github

Committed by

web-flow

Commit Message

Merge ee805adeb into 06bc84ec9

Pull Request Pull Request #99: Support CM time_metric (NS/MS/SEC/US) across reader and viz

Coverage Stats

34832 of 86278 branches covered (40.37%)

Branch coverage included in aggregate %.

1034 of 1319 new or added lines in 30 files covered. (78.39%)

5193 existing lines in 197 files now uncovered.

35349 of 49791 relevant lines covered (70.99%)

9765.54 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

20.45

/src/dftracer/utils/python/index_database.cpp

#include <dftracer/utils/python/index_database.h>
#include <dftracer/utils/python/py_errors.h>
#include <dftracer/utils/python/py_list_helpers.h>
#include <dftracer/utils/python/py_runtime_mixin.h>
#include <dftracer/utils/python/py_type_helpers.h>
#include <dftracer/utils/python/sst_distribution.h>
#include <dftracer/utils/utilities/indexer/index_database.h>
#include <dftracer/utils/utilities/indexer/index_database_sst_writer_context.h>

#include <new>
#include <string>
#include <unordered_set>
#include <vector>

using dftracer::utils::utilities::indexer::IndexDatabase;
using dftracer::utils::utilities::indexer::SstArtifactRegistry;

static void IndexDatabase_dealloc(IndexDatabaseObject *self) {
    self->db.~shared_ptr<IndexDatabase>();
    Py_TYPE(self)->tp_free((PyObject *)self);
}

static PyObject *IndexDatabase_new(PyTypeObject *type, PyObject * /*args*/,
                                   PyObject * /*kwds*/) {
    auto *self = (IndexDatabaseObject *)type->tp_alloc(type, 0);
    if (!self) return NULL;
    new (&self->db) std::shared_ptr<IndexDatabase>();
    return (PyObject *)self;
}

static int IndexDatabase_init(IndexDatabaseObject *self, PyObject *args,
                              PyObject *kwds) {
    static const char *kwlist[] = {"index_path", NULL};
    const char *index_path;
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", (char **)kwlist,
                                     &index_path)) {
        return -1;
    }
    try {
        self->db = std::make_shared<IndexDatabase>(index_path);
    } catch (const std::exception &e) {
        set_typed_py_error(e);
        return -1;
    }
    return 0;
}

static PyObject *IndexDatabase_init_schema(IndexDatabaseObject *self,
                                           PyObject * /*ignored*/) {
    if (!self->db) {
        PyErr_SetString(PyExc_RuntimeError, "IndexDatabase not initialised");
        return NULL;
    }
    if (!run_blocking([&] { self->db->init_schema(); })) return NULL;
    Py_RETURN_NONE;
}

static PyObject *IndexDatabase_register_files(IndexDatabaseObject *self,
                                              PyObject *args, PyObject *kwds) {
    static const char *kwlist[] = {"paths", "build_manifest", NULL};
    PyObject *paths_obj;
    int build_manifest = 0;
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|p", (char **)kwlist,
                                     &paths_obj, &build_manifest)) {
        return NULL;
    }
    std::vector<std::string> paths;
    if (!parse_str_list(paths_obj, "paths", paths)) return NULL;

    std::vector<int> ids;
    if (!run_blocking_r(
            [&] {
                return self->db->register_files(paths, build_manifest != 0);
            },
            ids)) {
        return NULL;
    }

    PyObject *out = PyList_New(static_cast<Py_ssize_t>(ids.size()));
    if (!out) return NULL;
    for (Py_ssize_t i = 0; i < static_cast<Py_ssize_t>(ids.size()); ++i) {
        PyList_SET_ITEM(out, i, PyLong_FromLong(ids[i]));
    }
    return out;
}

static PyObject *IndexDatabase_reserve_file_id_range(IndexDatabaseObject *self,
                                                     PyObject *args) {
    Py_ssize_t count;
    if (!PyArg_ParseTuple(args, "n", &count)) return NULL;
    if (count < 0) {
        PyErr_SetString(PyExc_ValueError, "count must be >= 0");
        return NULL;
    }
    int first;
    if (!run_blocking_r(
            [&] {
                return self->db->reserve_file_id_range(
                    static_cast<std::size_t>(count));
            },
            first)) {
        return NULL;
    }
    return PyLong_FromLong(first);
}

static PyObject *IndexDatabase_bulk_ingest(IndexDatabaseObject *self,
                                           PyObject *args, PyObject *kwds) {
    static const char *kwlist[] = {"registry", "skip_cfs", NULL};
    PyObject *registry_obj;
    PyObject *skip_cfs_obj = NULL;
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", (char **)kwlist,
                                     &registry_obj, &skip_cfs_obj)) {
        return NULL;
    }

    SstArtifactRegistry *registry = sst_artifact_registry_get(registry_obj);
    if (!registry) {
        PyErr_SetString(PyExc_TypeError,
                        "expected an SstArtifactRegistry instance");
        return NULL;
    }

    std::unordered_set<std::string> skip_cfs;
    if (skip_cfs_obj && skip_cfs_obj != Py_None) {
        PyObject *seq =
            PySequence_Fast(skip_cfs_obj, "skip_cfs must be an iterable");
        if (!seq) return NULL;
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
        for (Py_ssize_t i = 0; i < n; ++i) {
            PyObject *item = PySequence_Fast_GET_ITEM(seq, i);
            const char *s = PyUnicode_AsUTF8(item);
            if (!s) {
                Py_DECREF(seq);
                return NULL;
            }
            skip_cfs.emplace(s);
        }
        Py_DECREF(seq);
    }

    if (!run_blocking([&] { self->db->bulk_ingest(*registry, skip_cfs); }))
        return NULL;
    Py_RETURN_NONE;
}

static PyObject *IndexDatabase_write_agg_file_markers(IndexDatabaseObject *self,
                                                      PyObject *args) {
    PyObject *ids_obj;
    if (!PyArg_ParseTuple(args, "O", &ids_obj)) return NULL;

    PyObject *seq = PySequence_Fast(ids_obj, "file_ids must be an iterable");
    if (!seq) return NULL;
    Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
    std::vector<int> file_ids;
    file_ids.reserve(static_cast<std::size_t>(n));
    for (Py_ssize_t i = 0; i < n; ++i) {
        PyObject *item = PySequence_Fast_GET_ITEM(seq, i);
        long v = PyLong_AsLong(item);
        if (v == -1 && PyErr_Occurred()) {
            Py_DECREF(seq);
            return NULL;
        }
        file_ids.push_back(static_cast<int>(v));
    }
    Py_DECREF(seq);

    if (!run_blocking([&] { self->db->write_agg_file_markers(file_ids); }))
        return NULL;
    Py_RETURN_NONE;
}

static PyObject *IndexDatabase_write_agg_global_config(
    IndexDatabaseObject *self, PyObject *args, PyObject *kwds) {
    static const char *kwlist[] = {"time_interval_us", "config_hash", NULL};
    unsigned long long time_interval_us = 0;
    unsigned int config_hash = 0;
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "K|I", (char **)kwlist,
                                     &time_interval_us, &config_hash)) {
        return NULL;
    }
    if (!run_blocking([&] {
            self->db->write_agg_global_config(
                static_cast<std::uint64_t>(time_interval_us),
                static_cast<std::uint32_t>(config_hash));
        })) {
        return NULL;
    }
    Py_RETURN_NONE;
}

static PyObject *IndexDatabase_write_aggregation_tracker(
    IndexDatabaseObject *self, PyObject *args) {
    PyObject *blobs_obj;
    if (!PyArg_ParseTuple(args, "O", &blobs_obj)) return NULL;
    PyObject *seq = PySequence_Fast(blobs_obj, "blobs must be an iterable");
    if (!seq) return NULL;
    Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
    std::vector<std::string> blobs;
    blobs.reserve(static_cast<std::size_t>(n));
    for (Py_ssize_t i = 0; i < n; ++i) {
        PyObject *item = PySequence_Fast_GET_ITEM(seq, i);
        if (item == Py_None) continue;
        char *buf = nullptr;
        Py_ssize_t len = 0;
        if (PyBytes_Check(item)) {
            if (PyBytes_AsStringAndSize(item, &buf, &len) < 0) {
                Py_DECREF(seq);
                return NULL;
            }
        } else {
            Py_DECREF(seq);
            PyErr_SetString(PyExc_TypeError,
                            "blobs entries must be bytes or None");
            return NULL;
        }
        if (len > 0) blobs.emplace_back(buf, static_cast<std::size_t>(len));
    }
    Py_DECREF(seq);
    if (!run_blocking([&] { self->db->write_aggregation_tracker(blobs); }))
        return NULL;
    Py_RETURN_NONE;
}

static PyObject *IndexDatabase_rebuild_root_summaries(IndexDatabaseObject *self,
                                                      PyObject * /*ignored*/) {
    if (!run_blocking([&] { self->db->rebuild_root_summaries(); })) return NULL;
    Py_RETURN_NONE;
}

static PyObject *build_str_list(const std::vector<std::string> &v) {
    PyObject *lst = PyList_New(static_cast<Py_ssize_t>(v.size()));
    if (!lst) return NULL;
    for (Py_ssize_t i = 0; i < static_cast<Py_ssize_t>(v.size()); ++i) {
        PyObject *s =
            PyUnicode_FromString(v[static_cast<std::size_t>(i)].c_str());
        if (!s) {
            Py_DECREF(lst);
            return NULL;
        }
        PyList_SET_ITEM(lst, i, s);
    }
    return lst;
}

static PyObject *IndexDatabase_find_stale_files(IndexDatabaseObject *self,
                                                PyObject *args,
                                                PyObject *kwds) {
    static const char *kwlist[] = {"paths", NULL};
    PyObject *paths_obj;
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", (char **)kwlist,
                                     &paths_obj)) {
        return NULL;
    }
    std::vector<std::string> paths;
    if (!parse_str_list(paths_obj, "paths", paths)) return NULL;

    IndexDatabase::StaleCheckResult result;
    if (!run_blocking_r([&] { return self->db->find_stale_files(paths); },
                        result)) {
        return NULL;
    }

    PyObject *changed = build_str_list(result.changed);
    PyObject *added = build_str_list(result.added);
    PyObject *removed = build_str_list(result.removed);
    PyObject *d = PyDict_New();
    if (!changed || !added || !removed || !d) {
        Py_XDECREF(changed);
        Py_XDECREF(added);
        Py_XDECREF(removed);
        Py_XDECREF(d);
        return NULL;
    }
    PyObject *so = PyBool_FromLong(result.schema_outdated);
    PyObject *st = PyBool_FromLong(result.stale());
    PyDict_SetItemString(d, "changed", changed);
    PyDict_SetItemString(d, "added", added);
    PyDict_SetItemString(d, "removed", removed);
    PyDict_SetItemString(d, "schema_outdated", so);
    PyDict_SetItemString(d, "stale", st);
    Py_DECREF(changed);
    Py_DECREF(added);
    Py_DECREF(removed);
    Py_DECREF(so);
    Py_DECREF(st);
    return d;
}

static PyMethodDef IndexDatabase_methods[] = {
    {"init_schema", (PyCFunction)IndexDatabase_init_schema, METH_NOARGS,
     "Idempotently initialise the schema version key."},
    {"register_files", (PyCFunction)IndexDatabase_register_files,
     METH_VARARGS | METH_KEYWORDS,
     "register_files(paths, build_manifest=False) -> list[int]\n"
     "Register each path in the DEFAULT-CF file registry and return the "
     "assigned file_ids. Idempotent for files with matching hash."},
    {"find_stale_files", (PyCFunction)IndexDatabase_find_stale_files,
     METH_VARARGS | METH_KEYWORDS,
     "find_stale_files(paths) -> dict\n"
     "Stat-only (mtime + size) staleness check of the given trace paths "
     "against the index. Returns {changed, added, removed, schema_outdated, "
     "stale}."},
    {"reserve_file_id_range", (PyCFunction)IndexDatabase_reserve_file_id_range,
     METH_VARARGS,
     "reserve_file_id_range(count) -> int\n"
     "Atomically reserve `count` contiguous file_ids, return the first."},
    {"bulk_ingest", (PyCFunction)IndexDatabase_bulk_ingest,
     METH_VARARGS | METH_KEYWORDS,
     "bulk_ingest(registry, skip_cfs=None) -> None\n"
     "Ingest all SSTs collected in the SstArtifactRegistry.\n"
     "skip_cfs is an optional iterable of CF names whose SSTs are left "
     "outside the unified DB (used by distributed builds to keep "
     "AGGREGATION/SYSTEM_METRICS SSTs addressable by manifest)."},
    {"rebuild_root_summaries",
     (PyCFunction)IndexDatabase_rebuild_root_summaries, METH_NOARGS,
     "Recompute ROOT_* summary column families from per-file CFs."},
    {"write_agg_global_config",
     (PyCFunction)IndexDatabase_write_agg_global_config,
     METH_VARARGS | METH_KEYWORDS,
     "write_agg_global_config(time_interval_us, config_hash=0) -> None\n"
     "Write the AGG_GLOBAL_CONFIG_KEY marker into the AGGREGATION CF. "
     "Required for `iter_arrow_dfanalyzer_all` on distributed builds "
     "(which never materialise the key via worker SSTs) or "
     "post-consolidate indices."},
    {"write_agg_file_markers",
     (PyCFunction)IndexDatabase_write_agg_file_markers, METH_VARARGS,
     "write_agg_file_markers(file_ids) -> None\n"
     "Write per-file aggregation completion markers (\\xFF\\xFF + file_id) "
     "into the AGGREGATION CF. Required after distributed_index otherwise "
     "`ensure_indexed()` concludes aggregation is incomplete and re-runs "
     "the entire build."},
    {"write_aggregation_tracker",
     (PyCFunction)IndexDatabase_write_aggregation_tracker, METH_VARARGS,
     "write_aggregation_tracker(blobs) -> None\n"
     "Merge a list of serialized AssociationTracker bytes and write the "
     "result to the AGGREGATION CF under the `__tracker__` key."},
    {NULL}};

PyTypeObject IndexDatabaseType = {
    PyVarObject_HEAD_INIT(NULL, 0) "dftracer_utils_ext.IndexDatabase",
    sizeof(IndexDatabaseObject),
    0,
    (destructor)IndexDatabase_dealloc,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    Py_TPFLAGS_DEFAULT,
    "Handle to a .dftindex RocksDB store.",
    0,
    0,
    0,
    0,
    0,
    0,
    IndexDatabase_methods,
    0,
    0,
    0,
    0,
    0,
    0,
    0,
    (initproc)IndexDatabase_init,
    0,
    IndexDatabase_new,
};

int init_index_database(PyObject *m) {
    if (register_type(m, &IndexDatabaseType, "IndexDatabase") < 0) return -1;
    return 0;
}

1	#include <dftracer/utils/python/index_database.h>
2	#include <dftracer/utils/python/py_errors.h>
3	#include <dftracer/utils/python/py_list_helpers.h>
4	#include <dftracer/utils/python/py_runtime_mixin.h>
5	#include <dftracer/utils/python/py_type_helpers.h>
6	#include <dftracer/utils/python/sst_distribution.h>
7	#include <dftracer/utils/utilities/indexer/index_database.h>
8	#include <dftracer/utils/utilities/indexer/index_database_sst_writer_context.h>
9
10	#include <new>
11	#include <string>
12	#include <unordered_set>
13	#include <vector>
14
15	using dftracer::utils::utilities::indexer::IndexDatabase;
16	using dftracer::utils::utilities::indexer::SstArtifactRegistry;
17
18	static void IndexDatabase_dealloc(IndexDatabaseObject *self) {	4✔
19	self->db.~shared_ptr<IndexDatabase>();	4✔
20	Py_TYPE(self)->tp_free((PyObject *)self);	4✔
21	}	4✔
22
23	static PyObject IndexDatabase_new(PyTypeObject type, PyObject * /args/,	4✔
24	PyObject * /kwds/) {
25	auto self = (IndexDatabaseObject )type->tp_alloc(type, 0);	4✔
26	if (!self) return NULL;	4!
27	new (&self->db) std::shared_ptr<IndexDatabase>();	4✔
28	return (PyObject *)self;	4✔
29	}	4✔
30
31	static int IndexDatabase_init(IndexDatabaseObject self, PyObject args,	4✔
32	PyObject *kwds) {
33	static const char *kwlist[] = {"index_path", NULL};
34	const char *index_path;
35	if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", (char **)kwlist,	4!
36	&index_path)) {
37	return -1;	×
38	}
39	try {
40	self->db = std::make_shared<IndexDatabase>(index_path);	4!
41	} catch (const std::exception &e) {	4!
42	set_typed_py_error(e);	×
43	return -1;	×
44	}	×
45	return 0;	4✔
46	}	4✔
47
48	static PyObject IndexDatabase_init_schema(IndexDatabaseObject self,	×
49	PyObject * /ignored/) {
50	if (!self->db) {	×
51	PyErr_SetString(PyExc_RuntimeError, "IndexDatabase not initialised");	×
52	return NULL;	×
53	}
54	if (!run_blocking([&] { self->db->init_schema(); })) return NULL;	×
55	Py_RETURN_NONE;	×
UNCOV 56	}	×
57
58	static PyObject IndexDatabase_register_files(IndexDatabaseObject self,	×
59	PyObject args, PyObject kwds) {
60	static const char *kwlist[] = {"paths", "build_manifest", NULL};
61	PyObject *paths_obj;
62	int build_manifest = 0;	×
63	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O\|p", (char **)kwlist,	×
64	&paths_obj, &build_manifest)) {
65	return NULL;	×
66	}
67	std::vector<std::string> paths;	×
68	if (!parse_str_list(paths_obj, "paths", paths)) return NULL;	×
69
70	std::vector<int> ids;	×
71	if (!run_blocking_r(	×
72	[&] {	×
73	return self->db->register_files(paths, build_manifest != 0);	×
74	},
75	ids)) {
76	return NULL;	×
77	}
78
79	PyObject *out = PyList_New(static_cast<Py_ssize_t>(ids.size()));	×
80	if (!out) return NULL;	×
81	for (Py_ssize_t i = 0; i < static_cast<Py_ssize_t>(ids.size()); ++i) {	×
82	PyList_SET_ITEM(out, i, PyLong_FromLong(ids[i]));	×
UNCOV 83	}	×
84	return out;	×
85	}	×
86
87	static PyObject IndexDatabase_reserve_file_id_range(IndexDatabaseObject self,	×
88	PyObject *args) {
89	Py_ssize_t count;
90	if (!PyArg_ParseTuple(args, "n", &count)) return NULL;	×
91	if (count < 0) {	×
92	PyErr_SetString(PyExc_ValueError, "count must be >= 0");	×
93	return NULL;	×
94	}
95	int first;
96	if (!run_blocking_r(	×
97	[&] {	×
98	return self->db->reserve_file_id_range(	×
99	static_cast<std::size_t>(count));	×
100	},
101	first)) {
102	return NULL;	×
103	}
104	return PyLong_FromLong(first);	×
UNCOV 105	}	×
106
107	static PyObject IndexDatabase_bulk_ingest(IndexDatabaseObject self,	×
108	PyObject args, PyObject kwds) {
109	static const char *kwlist[] = {"registry", "skip_cfs", NULL};
110	PyObject *registry_obj;
111	PyObject *skip_cfs_obj = NULL;	×
112	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O\|O", (char **)kwlist,	×
113	&registry_obj, &skip_cfs_obj)) {
114	return NULL;	×
115	}
116
117	SstArtifactRegistry *registry = sst_artifact_registry_get(registry_obj);	×
118	if (!registry) {	×
119	PyErr_SetString(PyExc_TypeError,	×
120	"expected an SstArtifactRegistry instance");
121	return NULL;	×
122	}
123
124	std::unordered_set<std::string> skip_cfs;	×
125	if (skip_cfs_obj && skip_cfs_obj != Py_None) {	×
UNCOV 126	PyObject *seq =	×
127	PySequence_Fast(skip_cfs_obj, "skip_cfs must be an iterable");	×
128	if (!seq) return NULL;	×
129	Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);	×
130	for (Py_ssize_t i = 0; i < n; ++i) {	×
131	PyObject *item = PySequence_Fast_GET_ITEM(seq, i);	×
132	const char *s = PyUnicode_AsUTF8(item);	×
133	if (!s) {	×
UNCOV 134	Py_DECREF(seq);	×
135	return NULL;	×
136	}
137	skip_cfs.emplace(s);	×
UNCOV 138	}	×
UNCOV 139	Py_DECREF(seq);	×
UNCOV 140	}	×
141
142	if (!run_blocking([&] { self->db->bulk_ingest(*registry, skip_cfs); }))	×
143	return NULL;	×
144	Py_RETURN_NONE;	×
145	}	×
146
147	static PyObject IndexDatabase_write_agg_file_markers(IndexDatabaseObject self,	×
148	PyObject *args) {
149	PyObject *ids_obj;
150	if (!PyArg_ParseTuple(args, "O", &ids_obj)) return NULL;	×
151
152	PyObject *seq = PySequence_Fast(ids_obj, "file_ids must be an iterable");	×
153	if (!seq) return NULL;	×
154	Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);	×
155	std::vector<int> file_ids;	×
156	file_ids.reserve(static_cast<std::size_t>(n));	×
157	for (Py_ssize_t i = 0; i < n; ++i) {	×
158	PyObject *item = PySequence_Fast_GET_ITEM(seq, i);	×
159	long v = PyLong_AsLong(item);	×
160	if (v == -1 && PyErr_Occurred()) {	×
UNCOV 161	Py_DECREF(seq);	×
162	return NULL;	×
163	}
164	file_ids.push_back(static_cast<int>(v));	×
UNCOV 165	}	×
UNCOV 166	Py_DECREF(seq);	×
167
168	if (!run_blocking([&] { self->db->write_agg_file_markers(file_ids); }))	×
169	return NULL;	×
170	Py_RETURN_NONE;	×
171	}	×
172
173	static PyObject *IndexDatabase_write_agg_global_config(	×
174	IndexDatabaseObject self, PyObject args, PyObject *kwds) {
175	static const char *kwlist[] = {"time_interval_us", "config_hash", NULL};
176	unsigned long long time_interval_us = 0;	×
177	unsigned int config_hash = 0;	×
178	if (!PyArg_ParseTupleAndKeywords(args, kwds, "K\|I", (char **)kwlist,	×
179	&time_interval_us, &config_hash)) {
180	return NULL;	×
181	}
182	if (!run_blocking([&] {	×
183	self->db->write_agg_global_config(	×
184	static_cast<std::uint64_t>(time_interval_us),	×
185	static_cast<std::uint32_t>(config_hash));	×
186	})) {	×
187	return NULL;	×
188	}
189	Py_RETURN_NONE;	×
UNCOV 190	}	×
191
192	static PyObject *IndexDatabase_write_aggregation_tracker(	×
193	IndexDatabaseObject self, PyObject args) {
194	PyObject *blobs_obj;
195	if (!PyArg_ParseTuple(args, "O", &blobs_obj)) return NULL;	×
196	PyObject *seq = PySequence_Fast(blobs_obj, "blobs must be an iterable");	×
197	if (!seq) return NULL;	×
198	Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);	×
199	std::vector<std::string> blobs;	×
200	blobs.reserve(static_cast<std::size_t>(n));	×
201	for (Py_ssize_t i = 0; i < n; ++i) {	×
202	PyObject *item = PySequence_Fast_GET_ITEM(seq, i);	×
203	if (item == Py_None) continue;	×
204	char *buf = nullptr;	×
205	Py_ssize_t len = 0;	×
206	if (PyBytes_Check(item)) {	×
207	if (PyBytes_AsStringAndSize(item, &buf, &len) < 0) {	×
UNCOV 208	Py_DECREF(seq);	×
209	return NULL;	×
210	}
UNCOV 211	} else {	×
UNCOV 212	Py_DECREF(seq);	×
213	PyErr_SetString(PyExc_TypeError,	×
214	"blobs entries must be bytes or None");
215	return NULL;	×
216	}
217	if (len > 0) blobs.emplace_back(buf, static_cast<std::size_t>(len));	×
UNCOV 218	}	×
UNCOV 219	Py_DECREF(seq);	×
220	if (!run_blocking([&] { self->db->write_aggregation_tracker(blobs); }))	×
221	return NULL;	×
222	Py_RETURN_NONE;	×
223	}	×
224
225	static PyObject IndexDatabase_rebuild_root_summaries(IndexDatabaseObject self,	×
226	PyObject * /ignored/) {
227	if (!run_blocking([&] { self->db->rebuild_root_summaries(); })) return NULL;	×
228	Py_RETURN_NONE;	×
UNCOV 229	}	×
230
231	static PyObject *build_str_list(const std::vector<std::string> &v) {	12✔
232	PyObject *lst = PyList_New(static_cast<Py_ssize_t>(v.size()));	12✔
233	if (!lst) return NULL;	12!
234	for (Py_ssize_t i = 0; i < static_cast<Py_ssize_t>(v.size()); ++i) {	15✔
235	PyObject *s =	3✔
236	PyUnicode_FromString(v[static_cast<std::size_t>(i)].c_str());	3✔
237	if (!s) {	3!
UNCOV 238	Py_DECREF(lst);	×
239	return NULL;	×
240	}
241	PyList_SET_ITEM(lst, i, s);	3✔
242	}	3✔
243	return lst;	12✔
244	}	12✔
245
246	static PyObject IndexDatabase_find_stale_files(IndexDatabaseObject self,	4✔
247	PyObject *args,
248	PyObject *kwds) {
249	static const char *kwlist[] = {"paths", NULL};
250	PyObject *paths_obj;
251	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", (char **)kwlist,	4!
252	&paths_obj)) {
253	return NULL;	×
254	}
255	std::vector<std::string> paths;	4✔
256	if (!parse_str_list(paths_obj, "paths", paths)) return NULL;	4!
257
258	IndexDatabase::StaleCheckResult result;	4✔
259	if (!run_blocking_r([&] { return self->db->find_stale_files(paths); },	8!
260	result)) {
261	return NULL;	×
262	}
263
264	PyObject *changed = build_str_list(result.changed);	4!
265	PyObject *added = build_str_list(result.added);	4!
266	PyObject *removed = build_str_list(result.removed);	4!
267	PyObject *d = PyDict_New();	4!
268	if (!changed \|\| !added \|\| !removed \|\| !d) {	4!
269	Py_XDECREF(changed);	×
270	Py_XDECREF(added);	×
271	Py_XDECREF(removed);	×
272	Py_XDECREF(d);	×
273	return NULL;	×
274	}
275	PyObject *so = PyBool_FromLong(result.schema_outdated);	4!
276	PyObject *st = PyBool_FromLong(result.stale());	4!
277	PyDict_SetItemString(d, "changed", changed);	4!
278	PyDict_SetItemString(d, "added", added);	4!
279	PyDict_SetItemString(d, "removed", removed);	4!
280	PyDict_SetItemString(d, "schema_outdated", so);	4!
281	PyDict_SetItemString(d, "stale", st);	4!
282	Py_DECREF(changed);	4!
283	Py_DECREF(added);	4!
284	Py_DECREF(removed);	4!
285	Py_DECREF(so);	4!
286	Py_DECREF(st);	4!
287	return d;	4✔
288	}	4✔
289
290	static PyMethodDef IndexDatabase_methods[] = {
291	{"init_schema", (PyCFunction)IndexDatabase_init_schema, METH_NOARGS,
292	"Idempotently initialise the schema version key."},
293	{"register_files", (PyCFunction)IndexDatabase_register_files,
294	METH_VARARGS \| METH_KEYWORDS,
295	"register_files(paths, build_manifest=False) -> list[int]\n"
296	"Register each path in the DEFAULT-CF file registry and return the "
297	"assigned file_ids. Idempotent for files with matching hash."},
298	{"find_stale_files", (PyCFunction)IndexDatabase_find_stale_files,
299	METH_VARARGS \| METH_KEYWORDS,
300	"find_stale_files(paths) -> dict\n"
301	"Stat-only (mtime + size) staleness check of the given trace paths "
302	"against the index. Returns {changed, added, removed, schema_outdated, "
303	"stale}."},
304	{"reserve_file_id_range", (PyCFunction)IndexDatabase_reserve_file_id_range,
305	METH_VARARGS,
306	"reserve_file_id_range(count) -> int\n"
307	"Atomically reserve `count` contiguous file_ids, return the first."},
308	{"bulk_ingest", (PyCFunction)IndexDatabase_bulk_ingest,
309	METH_VARARGS \| METH_KEYWORDS,
310	"bulk_ingest(registry, skip_cfs=None) -> None\n"
311	"Ingest all SSTs collected in the SstArtifactRegistry.\n"
312	"skip_cfs is an optional iterable of CF names whose SSTs are left "
313	"outside the unified DB (used by distributed builds to keep "
314	"AGGREGATION/SYSTEM_METRICS SSTs addressable by manifest)."},
315	{"rebuild_root_summaries",
316	(PyCFunction)IndexDatabase_rebuild_root_summaries, METH_NOARGS,
317	"Recompute ROOT_* summary column families from per-file CFs."},
318	{"write_agg_global_config",
319	(PyCFunction)IndexDatabase_write_agg_global_config,
320	METH_VARARGS \| METH_KEYWORDS,
321	"write_agg_global_config(time_interval_us, config_hash=0) -> None\n"
322	"Write the AGG_GLOBAL_CONFIG_KEY marker into the AGGREGATION CF. "
323	"Required for `iter_arrow_dfanalyzer_all` on distributed builds "
324	"(which never materialise the key via worker SSTs) or "
325	"post-consolidate indices."},
326	{"write_agg_file_markers",
327	(PyCFunction)IndexDatabase_write_agg_file_markers, METH_VARARGS,
328	"write_agg_file_markers(file_ids) -> None\n"
329	"Write per-file aggregation completion markers (\\xFF\\xFF + file_id) "
330	"into the AGGREGATION CF. Required after distributed_index otherwise "
331	"`ensure_indexed()` concludes aggregation is incomplete and re-runs "
332	"the entire build."},
333	{"write_aggregation_tracker",
334	(PyCFunction)IndexDatabase_write_aggregation_tracker, METH_VARARGS,
335	"write_aggregation_tracker(blobs) -> None\n"
336	"Merge a list of serialized AssociationTracker bytes and write the "
337	"result to the AGGREGATION CF under the `__tracker__` key."},
338	{NULL}};
339
340	PyTypeObject IndexDatabaseType = {
341	PyVarObject_HEAD_INIT(NULL, 0) "dftracer_utils_ext.IndexDatabase",
342	sizeof(IndexDatabaseObject),
343	0,
344	(destructor)IndexDatabase_dealloc,
345	0,
346	0,
347	0,
348	0,
349	0,
350	0,
351	0,
352	0,
353	0,
354	0,
355	0,
356	0,
357	0,
358	0,
359	Py_TPFLAGS_DEFAULT,
360	"Handle to a .dftindex RocksDB store.",
361	0,
362	0,
363	0,
364	0,
365	0,
366	0,
367	IndexDatabase_methods,
368	0,
369	0,
370	0,
371	0,
372	0,
373	0,
374	0,
375	(initproc)IndexDatabase_init,
376	0,
377	IndexDatabase_new,
378	};
379
380	int init_index_database(PyObject *m) {	1✔
381	if (register_type(m, &IndexDatabaseType, "IndexDatabase") < 0) return -1;	1!
382	return 0;	1✔
383	}	1✔

llnl / dftracer-utils / 29787659139

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous