Source code

Revision control

Copy as Markdown

Other Tools

Test Info: Warnings

/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
"use strict";
const { PlacesSemanticHistoryDatabase } = ChromeUtils.importESModule(
"resource://gre/modules/PlacesSemanticHistoryDatabase.sys.mjs"
);
// Must be divisible by 8.
const EMBEDDING_SIZE = 64;
add_setup(async function () {
// Initialize Places.
Assert.equal(
PlacesUtils.history.databaseStatus,
PlacesUtils.history.DATABASE_STATUS_CREATE,
"Places initialized."
);
});
add_task(async function test_check_schema() {
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
let conn = await db.getConnection();
Assert.equal(
await conn.getSchemaVersion(),
db.currentSchemaVersion,
"Schema version should match."
);
await db.closeConnection();
await db.removeDatabaseFiles();
});
add_task(async function test_replace_on_downgrade() {
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
let conn = await db.getConnection();
let originalSchemaVersion = db.currentSchemaVersion;
await db.setCurrentSchemaVersionForTests(originalSchemaVersion + 1);
await db.closeConnection();
await db.setCurrentSchemaVersionForTests(originalSchemaVersion);
conn = await db.getConnection();
Assert.equal(
await conn.getSchemaVersion(),
db.currentSchemaVersion,
"Schema version should have been reset."
);
await db.closeConnection();
await db.removeDatabaseFiles();
});
add_task(async function test_broken_schema() {
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
let conn = await db.getConnection();
await conn.execute("DROP TABLE vec_history_mapping");
await db.closeConnection();
conn = await db.getConnection();
let rows = await conn.execute("SELECT COUNT(*) FROM vec_history_mapping");
Assert.equal(
rows[0].getResultByIndex(0),
0,
"Schema should have been reset."
);
await db.closeConnection();
await db.removeDatabaseFiles();
});
add_task(async function test_corruptdb() {
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
// Move a corrupt database file in place.
await IOUtils.copy(
do_get_file("../maintenance/corruptDB.sqlite").path,
db.databaseFilePath
);
let conn = await db.getConnection();
Assert.equal(
await conn.getSchemaVersion(),
db.currentSchemaVersion,
"Schema version should have been set."
);
await db.closeConnection();
await db.removeDatabaseFiles();
});
add_task(async function test_differentDimensionsReplacesDatabase() {
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
let conn = await db.getConnection();
Assert.ok(
!!(
await conn.execute(
`SELECT INSTR(sql, :needle) > 0
FROM sqlite_master WHERE name = 'vec_history'`,
{ needle: `FLOAT[${EMBEDDING_SIZE}]` }
)
)[0].getResultByIndex(0),
"Check embeddings size for the table"
);
await db.closeConnection();
db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE + 16,
fileName: "places_semantic.sqlite",
});
conn = await db.getConnection();
Assert.ok(
!!(
await conn.execute(
`SELECT INSTR(sql, :needle) > 0
FROM sqlite_master WHERE name = 'vec_history'`,
{ needle: `FLOAT[${EMBEDDING_SIZE + 16}]` }
)
)[0].getResultByIndex(0),
"Check that the database was replaced"
);
await db.closeConnection();
await db.removeDatabaseFiles();
});
add_task(async function test_healthydb() {
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
await db.getConnection();
await db.closeConnection();
// Check database creation time won't change when reopening, as that would
// indicate the database file was replaced.
let creationTime = (await IOUtils.stat(db.databaseFilePath)).creationTime;
db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
await db.getConnection();
await db.closeConnection();
Assert.equal(
creationTime,
(await IOUtils.stat(db.databaseFilePath)).creationTime,
"Database creation time should not change."
);
await db.removeDatabaseFiles();
});
add_task(async function test_defragmentation() {
Services.fog.initializeFOG();
let db = new PlacesSemanticHistoryDatabase({
embeddingSize: EMBEDDING_SIZE,
fileName: "places_semantic.sqlite",
});
// First connection creates the schema, the second connection does a database
// health check, which includes fragmentation monitoring.
let conn = await db.getConnection();
await db.closeConnection();
conn = await db.getConnection();
Assert.less(
Glean.places.databaseSemanticHistoryWastedPercentage.testGetValue(),
10,
"Wasted space should be less than 10%"
);
info("insert tensors and cause fragementation");
const HOW_MANY_TENSORS = 2000;
await conn.executeTransaction(async () => {
// Generate an empty chunk.
for (let i = 1; i <= HOW_MANY_TENSORS; i++) {
await insertTensor(conn, i, Array(EMBEDDING_SIZE).fill(Number(`0.${i}`)));
}
// Remove most of the tensors, but not all of them.
for (let i = 1; i <= HOW_MANY_TENSORS; i++) {
if (i % 10) {
await removeTensor(conn, i);
}
}
});
await db.closeConnection();
info("Reopening the connection to trigger defragmentation");
conn = await db.getConnection();
Assert.greater(
Glean.places.databaseSemanticHistoryWastedPercentage.testGetValue(),
10,
"Wasted space should be more than 10%"
);
Assert.equal(
Glean.places.databaseSemanticHistoryDefragmentTime.testGetValue().count,
1,
"There should be one defragmentation event"
);
Assert.greater(
Glean.places.databaseSemanticHistoryDefragmentTime.testGetValue().sum,
0,
"Defragmentation time should be greater than 0"
);
info("Check rowids were preserved in defragmentation");
for (let i = 10; i <= HOW_MANY_TENSORS; i += 10) {
await checkTensor(conn, i, Array(EMBEDDING_SIZE).fill(Number(`0.${i}`)));
}
await db.closeConnection();
conn = await db.getConnection();
Assert.less(
Glean.places.databaseSemanticHistoryWastedPercentage.testGetValue(),
10,
"Wasted space should be less than 10%"
);
info("Check old tables were removed");
Assert.equal(
(
await conn.execute(
`SELECT count(*) FROM sqlite_master WHERE name LIKE :suffix`,
{
suffix: "%_old",
}
)
)[0].getResultByIndex(0),
0,
"There should not be 'old' tables"
);
await db.closeConnection();
info("Reopen the connection and check fragmentation has been resolved");
conn = await db.getConnection();
Assert.less(
Glean.places.databaseSemanticHistoryWastedPercentage.testGetValue(),
10,
"Wasted space should be less than 10%"
);
await db.closeConnection();
await db.removeDatabaseFiles();
});
async function insertTensor(conn, rowid, tensor) {
await conn.execute(
`
INSERT INTO vec_history (rowid, embedding, embedding_coarse)
VALUES (:rowid, :vector, vec_quantize_binary(:vector))
`,
{
rowid,
vector: PlacesUtils.tensorToSQLBindable(tensor),
}
);
}
async function removeTensor(conn, rowid) {
await conn.execute(
`
DELETE FROM vec_history WHERE rowid = :rowid
`,
{
rowid,
}
);
}
async function checkTensor(conn, rowid, tensor) {
let rows = await conn.execute(
`
SELECT 1 FROM vec_history WHERE rowid = :rowid AND embedding = :vector
`,
{
rowid,
vector: PlacesUtils.tensorToSQLBindable(tensor),
}
);
Assert.equal(rows.length, 1, "Tensor with specified rowid should exist");
}