From 89f1c8a86c223947564dd85b6ea4d67c71d9ac15 Mon Sep 17 00:00:00 2001
From: joehan <joehanley@google.com>
Date: Mon, 9 Dec 2024 11:38:10 -0800
Subject: [PATCH] Fixing an issue with --import and dataDir at the same time
 (#8048)

---
 CHANGELOG.md                             |  1 +
 src/emulator/controller.ts               | 22 +++++++++++++++++++++-
 src/emulator/dataconnect/pgliteServer.ts | 18 +++++++++++++++++-
 src/emulator/downloadableEmulators.ts    |  3 +--
 src/error.ts                             |  5 +++++
 5 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0d51e30f741..b91ad91680e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1 +1,2 @@
 - Changes default CF3 runtime to nodejs22 (#8037)
+- Fixed an issue where `--import` would error for the Data Connect emulator if `dataDir` was also set.
diff --git a/src/emulator/controller.ts b/src/emulator/controller.ts
index 5668381c7b8..a6d574a754f 100755
--- a/src/emulator/controller.ts
+++ b/src/emulator/controller.ts
@@ -28,7 +28,7 @@ import { LoggingEmulator } from "./loggingEmulator";
 import * as dbRulesConfig from "../database/rulesConfig";
 import { EmulatorLogger, Verbosity } from "./emulatorLogger";
 import { EmulatorHubClient } from "./hubClient";
-import { confirm } from "../prompt";
+import { confirm, promptOnce } from "../prompt";
 import {
   FLAG_EXPORT_ON_EXIT_NAME,
   JAVA_DEPRECATION_WARNING,
@@ -893,6 +893,26 @@ export async function startAll(
         importDirAbsPath,
         exportMetadata.dataconnect.path,
       );
+      const dataDirectory = options.config.get("emulators.dataconnect.dataDir");
+      if (exportMetadataFilePath && dataDirectory) {
+        EmulatorLogger.forEmulator(Emulators.DATACONNECT).logLabeled(
+          "WARN",
+          "dataconnect",
+          "'firebase.json#emulators.dataconnect.dataDir' is set and `--import` flag was passed. " +
+            "This will overwrite any data saved from previous runs.",
+        );
+        if (
+          !options.nonInteractive &&
+          !(await promptOnce({
+            type: "confirm",
+            message: `Do you wish to continue and overwrite data in ${dataDirectory}?`,
+            default: false,
+          }))
+        ) {
+          await cleanShutdown();
+          return { deprecationNotices: [] };
+        }
+      }
 
       EmulatorLogger.forEmulator(Emulators.DATACONNECT).logLabeled(
         "BULLET",
diff --git a/src/emulator/dataconnect/pgliteServer.ts b/src/emulator/dataconnect/pgliteServer.ts
index 3550ff859ad..1aee1c2272c 100644
--- a/src/emulator/dataconnect/pgliteServer.ts
+++ b/src/emulator/dataconnect/pgliteServer.ts
@@ -17,6 +17,7 @@ import {
 } from "./pg-gateway/index";
 import { fromNodeSocket } from "./pg-gateway/platforms/node";
 import { logger } from "../../logger";
+import { hasMessage } from "../../error";
 export const TRUNCATE_TABLES_SQL = `
 DO $do$
 BEGIN
@@ -98,7 +99,7 @@ export class PostgresServer {
         const file = new File([rf], this.importPath);
         pgliteArgs.loadDataDir = file;
       }
-      this.db = await PGlite.create(pgliteArgs);
+      this.db = await this.forceCreateDB(pgliteArgs);
       await this.db.waitReady;
     }
     return this.db;
@@ -116,6 +117,21 @@ export class PostgresServer {
     fs.writeFileSync(exportPath, new Uint8Array(arrayBuff));
   }
 
+  async forceCreateDB(pgliteArgs: PGliteOptions): Promise<PGlite> {
+    try {
+      const db = await PGlite.create(pgliteArgs);
+      return db;
+    } catch (err: unknown) {
+      if (pgliteArgs.dataDir && hasMessage(err) && /Database already exists/.test(err.message)) {
+        // Clear out the current pglite data
+        fs.rmSync(pgliteArgs.dataDir, { force: true, recursive: true });
+        const db = await PGlite.create(pgliteArgs);
+        return db;
+      }
+      throw err;
+    }
+  }
+
   constructor(database: string, username: string, dataDirectory?: string, importPath?: string) {
     this.username = username;
     this.database = database;
diff --git a/src/emulator/downloadableEmulators.ts b/src/emulator/downloadableEmulators.ts
index 8df3b721bfd..51431926b66 100755
--- a/src/emulator/downloadableEmulators.ts
+++ b/src/emulator/downloadableEmulators.ts
@@ -9,7 +9,7 @@ import {
 } from "./types";
 import { Constants } from "./constants";
 
-import { FirebaseError } from "../error";
+import { FirebaseError, hasMessage } from "../error";
 import * as childProcess from "child_process";
 import * as utils from "../utils";
 import { EmulatorLogger } from "./emulatorLogger";
@@ -655,7 +655,6 @@ export async function start(
 }
 
 export function isIncomaptibleArchError(err: unknown): boolean {
-  const hasMessage = (e: any): e is { message: string } => !!e?.message;
   return (
     hasMessage(err) &&
     /Unknown system error/.test(err.message ?? "") &&
diff --git a/src/error.ts b/src/error.ts
index 93a5ac62f03..e85ddeaef71 100644
--- a/src/error.ts
+++ b/src/error.ts
@@ -121,3 +121,8 @@ export function isBillingError(e: {
     );
   });
 }
+
+/**
+ * Checks whether an unknown object (such as an error) has a message field
+ */
+export const hasMessage = (e: any): e is { message: string } => !!e?.message;