tor-browser

test_storage_ext_vec.js (6434B)

---

/* 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/. */

// This file tests support for the sqlite-vec extension.

function tensorToBlob(tensor) {
 return new Uint8ClampedArray(new Float32Array(tensor).buffer);
}

add_setup(async function () {
 cleanup();
});

add_task(async function test_synchronous() {
 info("Testing synchronous connection");
 let conn = getOpenedUnsharedDatabase();
 Assert.throws(
   () =>
     conn.executeSimpleSQL(
       `CREATE VIRTUAL TABLE test USING vec0(
         embedding FLOAT[4]
       );`
     ),
   /NS_ERROR_FAILURE/,
   "Should not be able to use vec without loading the extension"
 );

 await loadExtension(conn);

 conn.executeSimpleSQL(
   `
   CREATE VIRTUAL TABLE test USING vec0(
     embedding FLOAT[4]
   )
   `
 );

 let stmt = conn.createStatement(
   `
   INSERT INTO test(rowid, embedding)
   VALUES (1, :vector)
   `
 );
 stmt.bindBlobByName("vector", tensorToBlob([0.3, 0.3, 0.3, 0.3]));
 stmt.executeStep();
 stmt.reset();
 stmt.finalize();

 stmt = conn.createStatement(
   `
   SELECT
     rowid,
     distance
   FROM test
   WHERE embedding MATCH :vector
   ORDER BY distance
   LIMIT 1
   `
 );
 stmt.bindBlobByName("vector", tensorToBlob([0.3, 0.3, 0.3, 0.3]));
 Assert.ok(stmt.executeStep());
 Assert.equal(stmt.getInt32(0), 1);
 Assert.equal(stmt.getDouble(1), 0.0);
 stmt.reset();
 stmt.finalize();

 cleanup();
});

add_task(async function test_asynchronous() {
 info("Testing asynchronous connection");
 let conn = await openAsyncDatabase(getTestDB());

 await Assert.rejects(
   executeSimpleSQLAsync(
     conn,
     `
     CREATE VIRTUAL TABLE test USING vec0(
       embedding float[4]
     )
     `
   ),
   err => err.message.startsWith("no such module"),
   "Should not be able to use vec without loading the extension"
 );

 await loadExtension(conn);

 await executeSimpleSQLAsync(
   conn,
   `
   CREATE VIRTUAL TABLE test USING vec0(
     embedding float[4]
   )
   `
 );

 await asyncClose(conn);
 await IOUtils.remove(getTestDB().path, { ignoreAbsent: true });
});

add_task(async function test_clone() {
 info("Testing cloning synchronous connection loads extensions in clone");
 let conn1 = getOpenedUnsharedDatabase();
 await loadExtension(conn1);

 let conn2 = conn1.clone(false);
 conn2.executeSimpleSQL(
   `
   CREATE VIRTUAL TABLE test USING vec0(
     embedding float[4]
   )
   `
 );

 conn2.close();
 cleanup();
});

add_task(async function test_asyncClone() {
 info("Testing asynchronously cloning connection loads extensions in clone");
 let conn1 = getOpenedUnsharedDatabase();
 await loadExtension(conn1);

 let conn2 = await asyncClone(conn1, false);
 await executeSimpleSQLAsync(
   conn2,
   `
   CREATE VIRTUAL TABLE test USING vec0(
     embedding float[4]
   )
   `
 );

 await asyncClose(conn2);
 await asyncCleanup();
});

async function loadExtension(conn, ext = "vec") {
 await new Promise((resolve, reject) => {
   conn.loadExtension(ext, status => {
     if (Components.isSuccessCode(status)) {
       resolve();
     } else {
       reject(status);
     }
   });
 });
}

add_task(async function test_invariants() {
 // Test some invariants of the vec extension that we rely upon, so that if
 // the behavior changes we can catch it.
 let conn = getOpenedUnsharedDatabase();
 await loadExtension(conn);

 conn.executeSimpleSQL(
   `
   CREATE VIRTUAL TABLE vectors USING vec0(
     embedding FLOAT[4]
   )
   `
 );
 conn.executeSimpleSQL(
   `
   CREATE TABLE relations (
     rowid INTEGER PRIMARY KEY,
     content TEXT
   )
   `
 );

 let rowids = [];
 let insertRelStmt = conn.createStatement(
   `
   INSERT INTO relations (rowid, content)
   VALUES (NULL, "test")
   RETURNING rowid
   `
 );
 Assert.ok(insertRelStmt.executeStep());
 rowids.push(insertRelStmt.getInt32(0));
 insertRelStmt.reset();
 Assert.ok(insertRelStmt.executeStep());
 rowids.push(insertRelStmt.getInt32(0));
 insertRelStmt.reset();

 // Try to insert the same rowid twice in the vec table.
 let insertVecStmt = conn.createStatement(
   `
   INSERT INTO vectors (rowid, embedding)
   VALUES (:rowid, :vector)
   `
 );
 insertVecStmt.bindByName("rowid", rowids[0]);
 insertVecStmt.bindBlobByName("vector", tensorToBlob([0.1, 0.1, 0.1, 0.1]));
 insertVecStmt.executeStep();
 insertVecStmt.reset();

 let deleteStmt = conn.createStatement(
   `
   DELETE FROM vectors WHERE rowid = :rowid
   `
 );
 deleteStmt.bindByName("rowid", rowids[0]);
 deleteStmt.executeStep();
 deleteStmt.finalize();

 insertVecStmt.bindByName("rowid", rowids[0]);
 insertVecStmt.bindBlobByName("vector", tensorToBlob([0.2, 0.2, 0.2, 0.2]));
 insertVecStmt.executeStep();
 insertVecStmt.reset();

 let selectStmt = conn.createStatement(
   `
   SELECT
     rowid,
     vec_to_json(embedding)
   FROM vectors
   `
 );
 let count = 0;
 while (selectStmt.executeStep()) {
   count++;
   Assert.equal(selectStmt.getInt32(0), rowids[0]);
   Assert.equal(
     selectStmt.getUTF8String(1).replace(/(?<=[0-9])0+/g, ""),
     "[0.2,0.2,0.2,0.2]"
   );
 }
 Assert.equal(count, 1, "Should have one row in the vec table");
 selectStmt.reset();

 Assert.ok(insertRelStmt.executeStep());
 rowids.push(insertRelStmt.getInt32(0));
 insertRelStmt.finalize();
 insertVecStmt.bindByName("rowid", rowids[2]);
 insertVecStmt.bindBlobByName("vector", tensorToBlob([0.3, 0.3, 0.3, 0.3]));
 insertVecStmt.executeStep();
 insertVecStmt.finalize();

 let expected = [
   { rowid: rowids[0], vector: "[0.2,0.2,0.2,0.2]" },
   { rowid: rowids[2], vector: "[0.3,0.3,0.3,0.3]" },
 ];
 count = 0;
 for (let i = 0; selectStmt.executeStep(); i++) {
   count++;
   Assert.equal(selectStmt.getInt32(0), expected[i].rowid);
   Assert.equal(
     selectStmt.getUTF8String(1).replace(/(?<=[0-9])0+/g, ""),
     expected[i].vector
   );
 }
 Assert.equal(count, 2, "Should have two rows in the vec table");
 selectStmt.finalize();

 // TODO: In the future add testing for RETURNING and UPSERT as those are
 // currently broken. See:
 // https://github.com/asg017/sqlite-vec/issues/127
 // https://github.com/asg017/sqlite-vec/issues/229

 cleanup();
});