From eb52df2f931b6d4ac9a2d568192c81759c24404e Mon Sep 17 00:00:00 2001
From: Harold Shen <hlshen@google.com>
Date: Thu, 6 Jun 2024 17:31:52 -0400
Subject: [PATCH] fix behaviour on failed postgres connection (#7287)

Co-authored-by: joehan <joehanley@google.com>
---
 firebase-vscode/CHANGELOG.md                        |  2 ++
 firebase-vscode/src/data-connect/emulator-stream.ts | 13 ++++++++++---
 firebase-vscode/src/data-connect/emulator.ts        | 10 +++++++---
 firebase-vscode/src/data-connect/index.ts           |  6 +++++-
 4 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/firebase-vscode/CHANGELOG.md b/firebase-vscode/CHANGELOG.md
index d549548de98..ad50960fff4 100644
--- a/firebase-vscode/CHANGELOG.md
+++ b/firebase-vscode/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## NEXT
 
+- Fix behaviour on failed postgres connection
+
 ## 0.2.5
 
 - Icon fix
diff --git a/firebase-vscode/src/data-connect/emulator-stream.ts b/firebase-vscode/src/data-connect/emulator-stream.ts
index 5d9f3ba77d5..840bf8ea106 100644
--- a/firebase-vscode/src/data-connect/emulator-stream.ts
+++ b/firebase-vscode/src/data-connect/emulator-stream.ts
@@ -3,6 +3,7 @@ import fetch from "node-fetch";
 import { Observable, of } from "rxjs";
 import { backOff } from "exponential-backoff";
 import { ResolvedDataConnectConfigs } from "./config";
+import { Signal } from "@preact/signals-core";
 
 enum Kind {
   KIND_UNSPECIFIED = "KIND_UNSPECIFIED",
@@ -28,19 +29,20 @@ export const emulatorOutputChannel =
   vscode.window.createOutputChannel("Firebase Emulators");
 
 /**
- *
+ * TODO: convert to class
  * @param fdcEndpoint FDC Emulator endpoint
  */
 export async function runEmulatorIssuesStream(
   configs: ResolvedDataConnectConfigs,
   fdcEndpoint: string,
+  isPostgresEnabled: Signal<boolean>,
 ) {
   const obsErrors = await getEmulatorIssuesStream(configs, fdcEndpoint);
   const obsConverter = {
     next(nextResponse: EmulatorIssueResponse) {
       if (nextResponse.result?.issues?.length) {
         for (const issue of nextResponse.result.issues) {
-          displayIssue(issue);
+          displayAndHandleIssue(issue, isPostgresEnabled);
         }
       }
     },
@@ -57,7 +59,7 @@ export async function runEmulatorIssuesStream(
 /**
  * Based on the severity of the issue, either log, display notification, or display interactive popup to the user
  */
-export function displayIssue(issue: EmulatorIssue) {
+export function displayAndHandleIssue(issue: EmulatorIssue, isPostgresEnabled: Signal<boolean>) {
   const issueMessage = `Data Connect Emulator: ${issue.kind.toString()} - ${issue.message}`;
   if (issue.severity === Severity.ALERT) {
     vscode.window.showErrorMessage(issueMessage);
@@ -65,6 +67,11 @@ export function displayIssue(issue: EmulatorIssue) {
     vscode.window.showWarningMessage(issueMessage);
   }
   emulatorOutputChannel.appendLine(issueMessage);
+
+  // special handlings
+  if (issue.kind === Kind.SQL_CONNECTION) {
+    isPostgresEnabled.value = false;
+  }
 }
 
 /**
diff --git a/firebase-vscode/src/data-connect/emulator.ts b/firebase-vscode/src/data-connect/emulator.ts
index 6a2c5d48be7..6f03c3bb467 100644
--- a/firebase-vscode/src/data-connect/emulator.ts
+++ b/firebase-vscode/src/data-connect/emulator.ts
@@ -22,6 +22,11 @@ export class DataConnectEmulatorController implements vscode.Disposable {
 
       // Notify webviews when the emulator status changes
       effect(() => {
+        if (this.isPostgresEnabled) {
+          this.emulatorsController.emulatorStatusItem.show();
+        } else {
+          this.emulatorsController.emulatorStatusItem.hide();
+        }
         notifyIsConnectedToPostgres(this.isPostgresEnabled.value);
       }),
 
@@ -84,10 +89,9 @@ export class DataConnectEmulatorController implements vscode.Disposable {
 
     // configure the emulator to use the local psql string
     const emulatorClient = new DataConnectEmulatorClient();
-    emulatorClient.configureEmulator({ connectionString: newConnectionString });
-
     this.isPostgresEnabled.value = true;
-    this.emulatorsController.emulatorStatusItem.show();
+
+    emulatorClient.configureEmulator({ connectionString: newConnectionString });
   }
 
   dispose() {
diff --git a/firebase-vscode/src/data-connect/index.ts b/firebase-vscode/src/data-connect/index.ts
index 24b037301e7..63c9663d0df 100644
--- a/firebase-vscode/src/data-connect/index.ts
+++ b/firebase-vscode/src/data-connect/index.ts
@@ -191,7 +191,11 @@ export function registerFdc(
         vscode.commands.executeCommand(
           "firebase.dataConnect.executeIntrospection",
         );
-        runEmulatorIssuesStream(configs, emulatorController.getLocalEndpoint().value);
+        runEmulatorIssuesStream(
+          configs,
+          emulatorController.getLocalEndpoint().value,
+          fdcEmulatorsController.isPostgresEnabled,
+        );
         runDataConnectCompiler(configs, emulatorController.getLocalEndpoint().value);
       }
     }),