Import + export support for the data connect emulator

CHANGELOG.md

@@ -0,0 +1,2 @@
+- Added `--import` and `emulators:export` support to the Data Connect emulator.
+- Added `firebase.json#emulators.dataconnect.dataDir`. When set, Data Connect data will be persisted to the configured directory between emulator runs.

schema/firebase-config.json

@@ -402,6 +402,9 @@
                 "dataconnect": {
                     "additionalProperties": false,
                     "properties": {
+                        "dataDir": {
+                            "type": "string"
+                        },
                         "host": {
                             "type": "string"
                         },

src/emulator/controller.ts

@@ -51,7 +51,7 @@
 import { AuthEmulator, SingleProjectMode } from "./auth";
 import { DatabaseEmulator, DatabaseEmulatorArgs } from "./databaseEmulator";
 import { EventarcEmulator } from "./eventarcEmulator";
-import { DataConnectEmulator } from "./dataconnectEmulator";
+import { DataConnectEmulator, DataConnectEmulatorArgs } from "./dataconnectEmulator";
 import { FirestoreEmulator, FirestoreEmulatorArgs } from "./firestoreEmulator";
 import { HostingEmulator } from "./hostingEmulator";
 import { PubsubEmulator } from "./pubsubEmulator";
@@ -70,19 +70,20 @@
 
 /**
  * Exports emulator data on clean exit (SIGINT or process end)
  * @param options
  */
-export async function exportOnExit(options: any) {
-  const exportOnExitDir = options.exportOnExit;
+export async function exportOnExit(options: Options): Promise<void> {
+  // Note: options.exportOnExit is coerced to a string before this point in commandUtils.ts#setExportOnExitOptions
+  const exportOnExitDir = options.exportOnExit as string;
   if (exportOnExitDir) {
     try {
       utils.logBullet(
         `Automatically exporting data using ${FLAG_EXPORT_ON_EXIT_NAME} "${exportOnExitDir}" ` +
           "please wait for the export to finish...",
       );
       await exportEmulatorData(exportOnExitDir, options, /* initiatedBy= */ "exit");
-    } catch (e: any) {
-      utils.logWarning(e);
+    } catch (e: unknown) {
+      utils.logWarning(`${e}`);
       utils.logWarning(`Automatic export to "${exportOnExitDir}" failed, going to exit now...`);
     }
   }
@@ -90,10 +91,10 @@
 
 /**
  * Hook to do things when we're exiting cleanly (this does not include errors). Will be skipped on a second SIGINT
  * @param options
  */
 export async function onExit(options: any) {
   await exportOnExit(options);
 }
 
 /**
@@ -112,9 +113,9 @@
 
 /**
  * Filters a list of emulators to only those specified in the config
  * @param options
  */
 export function filterEmulatorTargets(options: { only: string; config: any }): Emulators[] {
   let targets = [...ALL_SERVICE_EMULATORS];
   targets.push(Emulators.EXTENSIONS);
   targets = targets.filter((e) => {
@@ -871,19 +872,41 @@
         `TODO: Add support for multiple services in the Data Connect emulator. Currently emulating first service ${config[0].source}`,
       );
     }
-    const configDir = config[0].source;
-    const dataConnectEmulator = new DataConnectEmulator({
+
+    const args: DataConnectEmulatorArgs = {
       listen: listenForEmulator.dataconnect,
       projectId,
       auto_download: true,
-      configDir,
+      configDir: config[0].source,
       rc: options.rc,
       config: options.config,
       autoconnectToPostgres: true,
       postgresListen: listenForEmulator["dataconnect.postgres"],
       enable_output_generated_sdk: true, // TODO: source from arguments
       enable_output_schema_extensions: true,
-    });
+    };
+
+    if (exportMetadata.dataconnect) {
+      utils.assertIsString(options.import);
+      const importDirAbsPath = path.resolve(options.import);
+      const exportMetadataFilePath = path.resolve(
+        importDirAbsPath,
+        exportMetadata.dataconnect.path,
+      );
+
+      EmulatorLogger.forEmulator(Emulators.DATACONNECT).logLabeled(
+        "BULLET",
+        "dataconnect",
+        `Importing data from ${exportMetadataFilePath}`,
+      );
+      args.importPath = exportMetadataFilePath;
+      void trackEmulator("emulator_import", {
+        initiated_by: "start",
+        emulator_name: Emulators.DATACONNECT,
+      });
+    }
+
+    const dataConnectEmulator = new DataConnectEmulator(args);
     await startEmulator(dataConnectEmulator);
   }
 

src/emulator/dataconnect/pgliteServer.ts

@@ -1,12 +1,14 @@
 // https://github.com/supabase-community/pg-gateway
 
-import { PGlite } from "@electric-sql/pglite";
+import { PGlite, PGliteOptions } from "@electric-sql/pglite";
 // Unfortunately, we need to dynamically import the Postgres extensions.
 // They are only available as ESM, and if we import them normally,
 // our tsconfig will convert them to requires, which will cause errors
 // during module resolution.
 const { dynamicImport } = require(true && "../../dynamicImport");
 import * as net from "node:net";
+import * as fs from "fs";
+
 import {
   getMessages,
   type PostgresConnection,
@@ -19,11 +21,13 @@ import { logger } from "../../logger";
 export class PostgresServer {
   private username: string;
   private database: string;
+  private dataDirectory?: string;
+  private importPath?: string;
 
-  public db: PGlite | undefined;
+  public db: PGlite | undefined = undefined;
   public async createPGServer(host: string = "127.0.0.1", port: number): Promise<net.Server> {
-    const db: PGlite = await this.getDb();
-    await db.waitReady;
+    const getDb = this.getDb.bind(this);
+
     const server = net.createServer(async (socket) => {
       const connection: PostgresConnection = await fromNodeSocket(socket, {
         serverVersion: "16.3 (PGlite 0.2.0)",
@@ -34,6 +38,7 @@ export class PostgresServer {
           if (!isAuthenticated) {
             return;
           }
+          const db = await getDb();
           const result = await db.execProtocolRaw(data);
           // Extended query patch removes the extra Ready for Query messages that
           // pglite wrongly sends.
@@ -50,41 +55,71 @@ export class PostgresServer {
         server.emit("error", err);
       });
     });
+
     const listeningPromise = new Promise<void>((resolve) => {
       server.listen(port, host, () => {
         resolve();
       });
     });
-    await db.waitReady;
     await listeningPromise;
     return server;
   }
 
   async getDb(): Promise<PGlite> {
-    if (this.db) {
-      return this.db;
+    if (!this.db) {
+      // Not all schemas will need vector installed, but we don't have an good way
+      // to swap extensions after starting PGLite, so we always include it.
+      const vector = (await dynamicImport("@electric-sql/pglite/vector")).vector;
+      const uuidOssp = (await dynamicImport("@electric-sql/pglite/contrib/uuid_ossp")).uuid_ossp;
+      const pgliteArgs: PGliteOptions = {
+        username: this.username,
+        database: this.database,
+        debug: 0,
+        extensions: {
+          vector,
+          uuidOssp,
+        },
+        dataDir: this.dataDirectory,
+      };
+      if (this.importPath) {
+        logger.debug(`Importing from ${this.importPath}`);
+        const rf = fs.readFileSync(this.importPath);
+        const file = new File([rf], this.importPath);
+        pgliteArgs.loadDataDir = file;
+      }
+      this.db = await PGlite.create(pgliteArgs);
+      await this.db.waitReady;
     }
-    // Not all schemas will need vector installed, but we don't have an good way
-    // to swap extensions after starting PGLite, so we always include it.
-    const vector = (await dynamicImport("@electric-sql/pglite/vector")).vector;
-    const uuidOssp = (await dynamicImport("@electric-sql/pglite/contrib/uuid_ossp")).uuid_ossp;
-    return PGlite.create({
-      username: this.username,
-      database: this.database,
-      debug: 0,
-      extensions: {
-        vector,
-        uuidOssp,
-      },
-      // TODO:  Use dataDir + loadDataDir to implement import/export.
-      // dataDir?: string;
-      // loadDataDir?: Blob | File;
-    });
+    return this.db;
+  }
+
+  public async clearDb(): Promise<void> {
+    const db = await this.getDb();
+    await db.query(`
+DO $do$
+BEGIN
+   EXECUTE
+   (SELECT 'TRUNCATE TABLE ' || string_agg(oid::regclass::text, ', ') || ' CASCADE'
+    FROM   pg_class
+    WHERE  relkind = 'r'
+    AND    relnamespace = 'public'::regnamespace
+   );
+END
+$do$;`);
+  }
+
+  public async exportData(exportPath: string): Promise<void> {
+    const db = await this.getDb();
+    const dump = await db.dumpDataDir();
+    const arrayBuff = await dump.arrayBuffer();
+    fs.writeFileSync(exportPath, new Uint8Array(arrayBuff));
   }
 
-  constructor(database: string, username: string) {
+  constructor(database: string, username: string, dataDirectory?: string, importPath?: string) {
     this.username = username;
     this.database = database;
+    this.dataDirectory = dataDirectory;
+    this.importPath = importPath;
   }
 }
 

src/emulator/dataconnectEmulator.ts

@@ -1,6 +1,7 @@
 import * as childProcess from "child_process";
 import { EventEmitter } from "events";
 import * as clc from "colorette";
+import * as path from "path";
 
 import { dataConnectLocalConnString } from "../api";
 import { Constants } from "./constants";
@@ -31,6 +32,7 @@ export interface DataConnectEmulatorArgs {
   postgresListen?: ListenSpec[];
   enable_output_schema_extensions: boolean;
   enable_output_generated_sdk: boolean;
+  importPath?: string;
 }
 
 export interface DataConnectGenerateArgs {
@@ -49,6 +51,7 @@ export const dataConnectEmulatorEvents = new EventEmitter();
 export class DataConnectEmulator implements EmulatorInstance {
   private emulatorClient: DataConnectEmulatorClient;
   private usingExistingEmulator: boolean = false;
+  private postgresServer: PostgresServer | undefined;
 
   constructor(private args: DataConnectEmulatorArgs) {
     this.emulatorClient = new DataConnectEmulatorClient();
@@ -102,11 +105,15 @@ export class DataConnectEmulator implements EmulatorInstance {
           `FIREBASE_DATACONNECT_POSTGRESQL_STRING is set to ${clc.bold(connStr)} - using that instead of starting a new database`,
         );
       } else if (pgHost && pgPort) {
-        const pgServer = new PostgresServer(dbId, "postgres");
-        const server = await pgServer.createPGServer(pgHost, pgPort);
+        const dataDirectory = this.args.config.get("emulators.dataconnect.dataDir");
+        const postgresDumpPath = this.args.importPath
+          ? path.join(this.args.importPath, "postgres.tar.gz")
+          : undefined;
+        this.postgresServer = new PostgresServer(dbId, "postgres", dataDirectory, postgresDumpPath);
+        const server = await this.postgresServer.createPGServer(pgHost, pgPort);
         const connectableHost = connectableHostname(pgHost);
         connStr = `postgres://${connectableHost}:${pgPort}/${dbId}?sslmode=disable`;
-        server.on("error", (err) => {
+        server.on("error", (err: any) => {
           if (err instanceof FirebaseError) {
             this.logger.logLabeled("ERROR", "Data Connect", `${err}`);
           } else {
@@ -170,6 +177,22 @@ export class DataConnectEmulator implements EmulatorInstance {
     return Emulators.DATACONNECT;
   }
 
+  async clearData(): Promise<void> {
+    if (this.postgresServer) {
+      await this.postgresServer.clearDb();
+    }
+  }
+
+  async exportData(exportPath: string): Promise<void> {
+    if (this.postgresServer) {
+      await this.postgresServer.exportData(path.join(exportPath, "postgres.tar.gz"));
+    } else {
+      throw new FirebaseError(
+        "The Data Connect emulator is currently connected to a separate Postgres instance. Export is not supported.",
+      );
+    }
+  }
+
   static async generate(args: DataConnectGenerateArgs): Promise<string> {
     const commandInfo = await downloadIfNecessary(Emulators.DATACONNECT);
     const cmd = [

src/emulator/hub.ts

@@ -11,6 +11,7 @@ import { EmulatorRegistry } from "./registry";
 import { FunctionsEmulator } from "./functionsEmulator";
 import { ExpressBasedEmulator } from "./ExpressBasedEmulator";
 import { PortName } from "./portUtils";
+import { DataConnectEmulator } from "./dataconnectEmulator";
 import { isVSCodeExtension } from "../vsCodeUtils";
 
 // We use the CLI version from package.json
@@ -36,6 +37,7 @@ export class EmulatorHub extends ExpressBasedEmulator {
   static PATH_DISABLE_FUNCTIONS = "/functions/disableBackgroundTriggers";
   static PATH_ENABLE_FUNCTIONS = "/functions/enableBackgroundTriggers";
   static PATH_EMULATORS = "/emulators";
+  static PATH_CLEAR_DATA_CONNECT = "/dataconnect/clearData";
 
   /**
    * Given a project ID, find and read the Locator file for the emulator hub.
@@ -164,6 +166,24 @@ export class EmulatorHub extends ExpressBasedEmulator {
       res.status(200).json({ enabled: true });
     });
 
+    app.post(EmulatorHub.PATH_CLEAR_DATA_CONNECT, async (req, res) => {
+      if (req.headers.origin) {
+        res.status(403).json({
+          message: `Clear Data Connect cannot be triggered by external callers.`,
+        });
+      }
+      utils.logLabeledBullet("emulators", `Clearing data from Data Connect data sources.`);
+
+      const instance = EmulatorRegistry.get(Emulators.DATACONNECT) as DataConnectEmulator;
+      if (!instance) {
+        res.status(400).json({ error: "The Data Connect emulator is not running." });
+        return;
+      }
+
+      await instance.clearData();
+      res.status(200).send("Data cleared");
+    });
+
     return app;
   }