diff --git a/CHANGELOG.md b/CHANGELOG.md index e69de29bb2d..fff63700721 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/schema/firebase-config.json b/schema/firebase-config.json index 1b1359f9c85..8eb45b14d7c 100644 --- a/schema/firebase-config.json +++ b/schema/firebase-config.json @@ -402,6 +402,9 @@ "dataconnect": { "additionalProperties": false, "properties": { + "dataDir": { + "type": "string" + }, "host": { "type": "string" }, diff --git a/src/emulator/controller.ts b/src/emulator/controller.ts index 1508a715315..c0f8172753a 100755 --- a/src/emulator/controller.ts +++ b/src/emulator/controller.ts @@ -51,7 +51,7 @@ import { Runtime, isRuntime } from "../deploy/functions/runtimes/supported"; 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"; @@ -72,8 +72,9 @@ const START_LOGGING_EMULATOR = utils.envOverride( * 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 { + // 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( @@ -81,8 +82,8 @@ export async function exportOnExit(options: any) { "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...`); } } @@ -871,19 +872,41 @@ export async function startAll( `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); } diff --git a/src/emulator/dataconnect/pgliteServer.ts b/src/emulator/dataconnect/pgliteServer.ts index f4c7b786a92..642b6b578a3 100644 --- a/src/emulator/dataconnect/pgliteServer.ts +++ b/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 { - 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((resolve) => { server.listen(port, host, () => { resolve(); }); }); - await db.waitReady; await listeningPromise; return server; } async getDb(): Promise { - 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 { + 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 { + 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; } } diff --git a/src/emulator/dataconnectEmulator.ts b/src/emulator/dataconnectEmulator.ts index 7f79bf2e7f7..4a8aa1d900d 100644 --- a/src/emulator/dataconnectEmulator.ts +++ b/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 { + if (this.postgresServer) { + await this.postgresServer.clearDb(); + } + } + + async exportData(exportPath: string): Promise { + 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 { const commandInfo = await downloadIfNecessary(Emulators.DATACONNECT); const cmd = [ diff --git a/src/emulator/hub.ts b/src/emulator/hub.ts index c926e389c03..d670656765f 100644 --- a/src/emulator/hub.ts +++ b/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; } diff --git a/src/emulator/hubExport.ts b/src/emulator/hubExport.ts index d9c0fd1b8fa..c61cd9b81bc 100644 --- a/src/emulator/hubExport.ts +++ b/src/emulator/hubExport.ts @@ -10,6 +10,7 @@ import { FirebaseError } from "../error"; import { EmulatorHub } from "./hub"; import { getDownloadDetails } from "./downloadableEmulators"; import { DatabaseEmulator } from "./databaseEmulator"; +import { DataConnectEmulator } from "./dataconnectEmulator"; import { rmSync } from "node:fs"; import { trackEmulator } from "../track"; @@ -34,12 +35,18 @@ export interface StorageExportMetadata { path: string; } +export interface DataConnectExportMetadata { + version: string; + path: string; +} + export interface ExportMetadata { version: string; firestore?: FirestoreExportMetadata; database?: DatabaseExportMetadata; auth?: AuthExportMetadata; storage?: StorageExportMetadata; + dataconnect?: DataConnectExportMetadata; } export interface ExportOptions { @@ -122,6 +129,14 @@ export class HubExport { await this.exportStorage(metadata); } + if (shouldExport(Emulators.DATACONNECT)) { + metadata.dataconnect = { + version: EmulatorHub.CLI_VERSION, + path: "dataconnect_export", + }; + await this.exportDataConnect(metadata); + } + // Make sure the export directory exists if (!fs.existsSync(this.exportPath)) { fs.mkdirSync(this.exportPath); @@ -289,6 +304,28 @@ export class HubExport { throw new FirebaseError(`Failed to export storage: ${await res.response.text()}`); } } + + private async exportDataConnect(metadata: ExportMetadata): Promise { + void trackEmulator("emulator_export", { + initiated_by: this.options.initiatedBy, + emulator_name: Emulators.DATACONNECT, + }); + + const instance = EmulatorRegistry.get(Emulators.DATACONNECT) as DataConnectEmulator; + if (!instance) { + throw new FirebaseError( + "Unable to export Data Connect emulator data: the Data Connect emulator is not running.", + ); + } + + const dataconnectExportPath = path.join(this.tmpDir, metadata.dataconnect!.path); + if (fs.existsSync(dataconnectExportPath)) { + fse.removeSync(dataconnectExportPath); + } + fs.mkdirSync(dataconnectExportPath); + + await instance.exportData(dataconnectExportPath); + } } function fetchToFile(options: http.RequestOptions, path: fs.PathLike): Promise { diff --git a/src/emulator/types.ts b/src/emulator/types.ts index 2e519fcdfdc..8492026cf85 100644 --- a/src/emulator/types.ts +++ b/src/emulator/types.ts @@ -43,6 +43,7 @@ export const IMPORT_EXPORT_EMULATORS = [ Emulators.DATABASE, Emulators.AUTH, Emulators.STORAGE, + Emulators.DATACONNECT, ]; export const ALL_SERVICE_EMULATORS = [ diff --git a/src/firebaseConfig.ts b/src/firebaseConfig.ts index 02315411016..c8a9279c934 100644 --- a/src/firebaseConfig.ts +++ b/src/firebaseConfig.ts @@ -241,6 +241,7 @@ export type EmulatorsConfig = { port?: number; postgresHost?: string; postgresPort?: number; + dataDir?: string; }; tasks?: { host?: string;