Add migrate function

This new function can be used to do migrations on both schema databases
and regular databases. It is mostly dedicated to schema migration tools.

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>

Author: haaawk

packages/libsql-client-wasm/src/wasm.ts

@@ -171,6 +171,33 @@ export class Sqlite3Client implements Client {
         }
     }
 
+    async migrate(
+        stmts: Array<InStatement>,
+    ): Promise<Array<ResultSet>> {
+        this.#checkNotClosed();
+        const db = this.#getDb();
+        try {
+            executeStmt(db, "PRAGMA foreign_keys=off", this.#intMode);
+            executeStmt(db, transactionModeToBegin("deferred"), this.#intMode);
+            const resultSets = stmts.map((stmt) => {
+                if (!inTransaction(db)) {
+                    throw new LibsqlError(
+                        "The transaction has been rolled back",
+                        "TRANSACTION_CLOSED",
+                    );
+                }
+                return executeStmt(db, stmt, this.#intMode);
+            });
+            executeStmt(db, "COMMIT", this.#intMode);
+            return resultSets;
+        } finally {
+            if (inTransaction(db)) {
+                executeStmt(db, "ROLLBACK", this.#intMode);
+            }
+            executeStmt(db, "PRAGMA foreign_keys=on", this.#intMode);
+        }
+    }
+
     async transaction(mode: TransactionMode = "write"): Promise<Transaction> {
         const db = this.#getDb();
         executeStmt(db, transactionModeToBegin(mode), this.#intMode);

packages/libsql-client/src/hrana.ts

@@ -252,7 +252,11 @@ export async function executeHranaBatch(
     version: hrana.ProtocolVersion,
     batch: hrana.Batch,
     hranaStmts: Array<hrana.Stmt>,
+    disableForeignKeys: boolean = false,
 ): Promise<Array<ResultSet>> {
+    if (disableForeignKeys) {
+        batch.step().run("PRAGMA foreign_keys=off")
+    }
     const beginStep = batch.step();
     const beginPromise = beginStep.run(transactionModeToBegin(mode));
 
@@ -282,6 +286,9 @@ export async function executeHranaBatch(
         .step()
         .condition(hrana.BatchCond.not(hrana.BatchCond.ok(commitStep)));
     rollbackStep.run("ROLLBACK").catch((_) => undefined);
+    if (disableForeignKeys) {
+        batch.step().run("PRAGMA foreign_keys=on")
+    }
 
     await batch.execute();
 

packages/libsql-client/src/http.ts

@@ -180,6 +180,40 @@ export class HttpClient implements Client {
         });
     }
 
+    async migrate(
+        stmts: Array<InStatement>,
+    ): Promise<Array<ResultSet>> {
+        return this.limit<Array<ResultSet>>(async () => {
+            try {
+                const hranaStmts = stmts.map(stmtToHrana);
+                const version = await this.#client.getVersion();
+
+                // Pipeline all operations, so `hrana.HttpClient` can open the stream, execute the batch and
+                // close the stream in a single HTTP request.
+                let resultsPromise: Promise<Array<ResultSet>>;
+                const stream = this.#client.openStream();
+                try {
+                    const batch = stream.batch(false);
+                    resultsPromise = executeHranaBatch(
+                        "deferred",
+                        version,
+                        batch,
+                        hranaStmts,
+                        true,
+                    );
+                } finally {
+                    stream.closeGracefully();
+                }
+
+                const results = await resultsPromise;
+
+                return results;
+            } catch (e) {
+                throw mapHranaError(e);
+            }
+        });
+    }
+
     async transaction(
         mode: TransactionMode = "write",
     ): Promise<HttpTransaction> {

packages/libsql-client/src/sqlite3.ts

@@ -167,6 +167,33 @@ export class Sqlite3Client implements Client {
         }
     }
 
+    async migrate(
+        stmts: Array<InStatement>,
+    ): Promise<Array<ResultSet>> {
+        this.#checkNotClosed();
+        const db = this.#getDb();
+        try {
+            executeStmt(db, "PRAGMA foreign_keys=off", this.#intMode);
+            executeStmt(db, transactionModeToBegin("deferred"), this.#intMode);
+            const resultSets = stmts.map((stmt) => {
+                if (!db.inTransaction) {
+                    throw new LibsqlError(
+                        "The transaction has been rolled back",
+                        "TRANSACTION_CLOSED",
+                    );
+                }
+                return executeStmt(db, stmt, this.#intMode);
+            });
+            executeStmt(db, "COMMIT", this.#intMode);
+            return resultSets;
+        } finally {
+            if (db.inTransaction) {
+                executeStmt(db, "ROLLBACK", this.#intMode);
+            }
+            executeStmt(db, "PRAGMA foreign_keys=on", this.#intMode);
+        }
+    }
+
     async transaction(mode: TransactionMode = "write"): Promise<Transaction> {
         const db = this.#getDb();
         executeStmt(db, transactionModeToBegin(mode), this.#intMode);

packages/libsql-client/src/ws.ts

@@ -226,6 +226,37 @@ export class WsClient implements Client {
         });
     }
 
+    async migrate(
+        stmts: Array<InStatement>,
+    ): Promise<Array<ResultSet>> {
+        return this.limit<Array<ResultSet>>(async () => {
+            const streamState = await this.#openStream();
+            try {
+                const hranaStmts = stmts.map(stmtToHrana);
+                const version = await streamState.conn.client.getVersion();
+
+                // Schedule all operations synchronously, so they will be pipelined and executed in a single
+                // network roundtrip.
+                const batch = streamState.stream.batch(version >= 3);
+                const resultsPromise = executeHranaBatch(
+                    "deferred",
+                    version,
+                    batch,
+                    hranaStmts,
+                    true,
+                );
+
+                const results = await resultsPromise;
+
+                return results;
+            } catch (e) {
+                throw mapHranaError(e);
+            } finally {
+                this._closeStream(streamState);
+            }
+        });
+    }
+
     async transaction(mode: TransactionMode = "write"): Promise<WsTransaction> {
         return this.limit<WsTransaction>(async () => {
             const streamState = await this.#openStream();

packages/libsql-core/src/api.ts

@@ -126,6 +126,39 @@ export interface Client {
         mode?: TransactionMode,
     ): Promise<Array<ResultSet>>;
 
+    /** Execute a batch of SQL statements in a transaction with PRAGMA foreign_keys=off; before and PRAGMA foreign_keys=on; after.
+     *
+     * The batch is executed in its own logical database connection and the statements are wrapped in a
+     * transaction. This ensures that the batch is applied atomically: either all or no changes are applied.
+     *
+     * The transaction mode is `"deferred"`.
+     *
+     * If any of the statements in the batch fails with an error, the batch is aborted, the transaction is
+     * rolled back and the returned promise is rejected.
+     *
+     * ```javascript
+     * const rss = await client.migrate([
+     *     // statement without arguments
+     *     "CREATE TABLE test (a INT)",
+     *
+     *     // statement with positional arguments
+     *     {
+     *         sql: "INSERT INTO books (name, author, published_at) VALUES (?, ?, ?)",
+     *         args: ["First Impressions", "Jane Austen", 1813],
+     *     },
+     *
+     *     // statement with named arguments
+     *     {
+     *         sql: "UPDATE books SET name = $new WHERE name = $old",
+     *         args: {old: "First Impressions", new: "Pride and Prejudice"},
+     *     },
+     * ]);
+     * ```
+     */
+    migrate(
+        stmts: Array<InStatement>,
+    ): Promise<Array<ResultSet>>;
+
     /** Start an interactive transaction.
      *
      * Interactive transactions allow you to interleave execution of SQL statements with your application