Cleanly shut down PGlite when shutting down Data Connect

CHANGELOG.md

@@ -1,7 +1,8 @@
 - Changes default CF3 runtime to nodejs22 (#8037)
 - Fixed an issue where `--import` would error for the Data Connect emulator if `dataDir` was also set.
 - Fixed an issue where `firebase init dataconnect` errored when importing a schema with no GQL files.
-- CF3 callables can now be annotate with a genkit action they are serving (#8039)
-- HTTPS functions can now be upgraded to HTTPS Callable functions (#8039)
-- Update default tsconfig to support more modern defaults (#8039)
+- Fixed an issue where the Data Connect emulator would not cleanly shut down Postgres and corrupt data. (#8044)
+- CF3 callables can now be annotate with a genkit action they are serving. (#8039)
+- HTTPS functions can now be upgraded to HTTPS Callable functions. (#8039)
+- Update default tsconfig to support more modern defaults. (#8039)
 - Update the Firebase Data Connect local toolkit to v1.7.5, which includes a fix for Kotlin codegen that ensures that generated XxxKeys.kt files include the required `@file:UseSerializers(UUIDSerializer::class)` annotation. (#8058)

src/emulator/dataconnect/pgliteServer.ts

@@ -1,11 +1,11 @@
 // https://github.com/supabase-community/pg-gateway
 
-import { PGlite, PGliteOptions } from "@electric-sql/pglite";
+import { DebugLevel, 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";
 
@@ -18,6 +18,7 @@
 import { fromNodeSocket } from "./pg-gateway/platforms/node";
 import { logger } from "../../logger";
 import { hasMessage } from "../../error";
+
 export const TRUNCATE_TABLES_SQL = `
 DO $do$
 BEGIN
@@ -35,12 +36,15 @@
   private database: string;
   private dataDirectory?: string;
   private importPath?: string;
+  private debug: DebugLevel;
 
   public db: PGlite | undefined = undefined;
+  private server: net.Server | undefined = undefined;
+
   public async createPGServer(host: string = "127.0.0.1", port: number): Promise<net.Server> {
     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)",
         auth: { method: "trust" },
@@ -54,11 +58,11 @@
           const result = await db.execProtocolRaw(data);
           // Extended query patch removes the extra Ready for Query messages that
           // pglite wrongly sends.
           return extendedQueryPatch.filterResponse(data, result);
         },
       });
 
       const extendedQueryPatch: PGliteExtendedQueryPatch = new PGliteExtendedQueryPatch(connection);
 
       socket.on("end", () => {
         logger.debug("Postgres client disconnected");
@@ -67,6 +71,7 @@
         server.emit("error", err);
       });
     });
+    this.server = server;
 
     const listeningPromise = new Promise<void>((resolve) => {
       server.listen(port, host, () => {
@@ -81,12 +86,12 @@
     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,
+        debug: this.debug,
         extensions: {
           vector,
           uuidOssp,
@@ -132,11 +137,28 @@
     }
   }
 
-  constructor(database: string, username: string, dataDirectory?: string, importPath?: string) {
-    this.username = username;
-    this.database = database;
-    this.dataDirectory = dataDirectory;
-    this.importPath = importPath;
+  public async stop(): Promise<void> {
+    if (this.db) {
+      await this.db.close();
+    }
+    if (this.server) {
+      this.server.close();
+    }
+    return;
+  }
+
+  constructor(args: {
+    database: string;
+    username: string;
+    dataDirectory?: string;
+    importPath?: string;
+    debug?: boolean;
+  }) {
+    this.username = args.username;
+    this.database = args.database;
+    this.dataDirectory = args.dataDirectory;
+    this.importPath = args.importPath;
+    this.debug = args.debug ? 5 : 0;
   }
 }
 

src/emulator/dataconnectEmulator.ts

@@ -40,6 +40,7 @@ export interface DataConnectEmulatorArgs {
   enable_output_schema_extensions: boolean;
   enable_output_generated_sdk: boolean;
   importPath?: string;
+  debug?: boolean;
 }
 
 export interface DataConnectGenerateArgs {
@@ -116,7 +117,13 @@ export class DataConnectEmulator implements EmulatorInstance {
         const postgresDumpPath = this.args.importPath
           ? path.join(this.args.importPath, "postgres.tar.gz")
           : undefined;
-        this.postgresServer = new PostgresServer(dbId, "postgres", dataDirectory, postgresDumpPath);
+        this.postgresServer = new PostgresServer({
+          database: dbId,
+          username: "fdc",
+          dataDirectory,
+          importPath: postgresDumpPath,
+          debug: this.args.debug,
+        });
         const server = await this.postgresServer.createPGServer(pgHost, pgPort);
         const connectableHost = connectableHostname(pgHost);
         connStr = `postgres://${connectableHost}:${pgPort}/${dbId}?sslmode=disable`;
@@ -166,6 +173,9 @@ export class DataConnectEmulator implements EmulatorInstance {
       );
       return;
     }
+    if (this.postgresServer) {
+      await this.postgresServer.stop();
+    }
     return stop(Emulators.DATACONNECT);
   }