API based Schema Migration behind experimental flag.

src/dataconnect/client.spec.ts

@@ -7,6 +7,7 @@ import * as client from "./client";
 import { FirebaseError } from "../error";
 import * as types from "./types";
 
+// eslint-disable-next-line @typescript-eslint/no-var-requires, @typescript-eslint/no-unsafe-argument
 chai.use(require("chai-as-promised"));
 
 describe("client", () => {
@@ -165,6 +166,23 @@ describe("client", () => {
       );
     });
 
+    it("executeSchemaMigration", async () => {
+      postStub.resolves({ body: { name: "op-name" } });
+      pollOperationStub.resolves({ done: true });
+      await client.executeSchemaMigration("projects/p/locations/l/services/s", [
+        { sql: "ALTER TABLE...", description: "test", destructive: false },
+      ]);
+      expect(postStub).to.be.calledWith("projects/p/locations/l/services/s/schemas/main:migrate", {
+        diffs: [{ sql: "ALTER TABLE...", description: "test", destructive: false }],
+      });
+      expect(pollOperationStub).to.be.calledWith({
+        apiOrigin: "https://firebasedataconnect.googleapis.com",
+        apiVersion: "v1",
+        operationResourceName: "op-name",
+        masterTimeout: 300000,
+      });
+    });
+
     it("deleteSchema", async () => {
       deleteStub.resolves({ body: { name: "op-name" } });
       pollOperationStub.resolves();

src/dataconnect/client.ts

@@ -6,14 +6,14 @@
 const DATACONNECT_API_VERSION = "v1";
 const PAGE_SIZE_MAX = 100;
 
 const dataconnectClient = () =>
   new Client({
     urlPrefix: dataconnectOrigin(),
     apiVersion: DATACONNECT_API_VERSION,
     auth: true,
   });
 
 export async function listLocations(projectId: string): Promise<string[]> {
   const res = await dataconnectClient().get<{
     locations: {
       name: string;
@@ -30,14 +30,14 @@
   return res.body;
 }
 
 export async function listAllServices(projectId: string): Promise<types.Service[]> {
   const res = await dataconnectClient().get<{ services: types.Service[] }>(
     `/projects/${projectId}/locations/-/services`,
   );
   return res.body.services ?? [];
 }
 
 export async function createService(
   projectId: string,
   locationId: string,
   serviceId: string,
@@ -60,15 +60,15 @@
       operationResourceName: op.body.name,
     });
     return pollRes;
   } catch (err: any) {
     if (err.status !== 409) {
       throw err;
     }
     return undefined; // Service already exists
   }
 }
 
 export async function deleteService(serviceName: string): Promise<types.Service> {
   // Note that we need to force delete in order to delete child resources too.
   const op = await dataconnectClient().delete<types.Service>(serviceName, {
     queryParams: { force: "true" },
@@ -83,15 +83,15 @@
 
 /** Schema methods */
 
 export async function getSchema(
   serviceName: string,
   schemaId: string = types.MAIN_SCHEMA_ID,
 ): Promise<types.Schema | undefined> {
   try {
     const res = await dataconnectClient().get<types.Schema>(`${serviceName}/schemas/${schemaId}`);
     return res.body;
   } catch (err: any) {
     if (err.status !== 404) {
       throw err;
     }
     return undefined;
@@ -207,3 +207,21 @@
   });
   return pollRes;
 }
+
+export async function executeSchemaMigration(
+  serviceName: string,
+  diffs: types.Diff[],
+): Promise<void> {
+  const client = dataconnectClient();
+  const op = await client.post<{ diffs: types.Diff[] }, { name: string }>(
+    `${serviceName}/schemas/main:migrate`,
+    { diffs },
+  );
+
+  await operationPoller.pollOperation<void>({
+    apiOrigin: dataconnectOrigin(),
+    apiVersion: DATACONNECT_API_VERSION,
+    operationResourceName: op.body.name,
+    masterTimeout: 300000, // Migrations might take longer than 60s
+  });
+}

src/dataconnect/schemaMigration.spec.ts

@@ -1,6 +1,15 @@
+import * as experiments from "../experiments";
+import * as sinon from "sinon";
+import * as client from "./client";
+import * as connect from "../gcp/cloudsql/connect";
+import * as cloudsqladmin from "../gcp/cloudsql/cloudsqladmin";
+import * as permissionsSetup from "../gcp/cloudsql/permissionsSetup";
+import { SchemaSetupStatus } from "../gcp/cloudsql/permissionsSetup";
+import { handleIncompatibleSchemaError } from "./schemaMigration";
 import { expect } from "chai";
 import { serviceNameFromSchema, getIdentifiers } from "./schemaMigration";
-import { Schema } from "./types";
+import { Schema, IncompatibleSqlSchemaError } from "./types";
+import { Options } from "../options";
 
 describe("serviceNameFromSchema", () => {
   it("main schema", () => {
@@ -125,3 +134,94 @@ describe("getIdentifiers", () => {
     );
   });
 });
+
+describe("handleIncompatibleSchemaError", () => {
+  let executeSchemaMigrationStub: sinon.SinonStub;
+  let executeSqlCmdsAsIamUserStub: sinon.SinonStub;
+  let isEnabledStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    executeSchemaMigrationStub = sinon.stub(client, "executeSchemaMigration").resolves();
+    executeSqlCmdsAsIamUserStub = sinon.stub(connect, "executeSqlCmdsAsIamUser").resolves();
+    sinon.stub(cloudsqladmin, "iamUserIsCSQLAdmin").resolves(true);
+    sinon
+      .stub(permissionsSetup, "getSchemaMetadata")
+      .resolves({ setupStatus: SchemaSetupStatus.GreenField } as permissionsSetup.SchemaMetadata);
+    sinon.stub(permissionsSetup, "checkSQLRoleIsGranted").resolves(true);
+    sinon.stub(connect, "getIAMUser").resolves({ user: "test-user", mode: "CLOUD_IAM_USER" });
+    isEnabledStub = sinon.stub(experiments, "isEnabled").returns(false);
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  const schema = {
+    name: "projects/p/locations/l/services/s/schemas/main",
+  } as Schema;
+
+  const incompatibleSchemaError = {
+    diffs: [
+      { sql: "CREATE TABLE a", destructive: false },
+      { sql: "DROP TABLE b", destructive: true },
+    ],
+  } as IncompatibleSqlSchemaError;
+
+  it("should execute all commands via IAM user when fdcapimigration experiment is not enabled", async () => {
+    isEnabledStub.withArgs("fdcapimigration").returns(false);
+    await handleIncompatibleSchemaError({
+      schema,
+      incompatibleSchemaError,
+      options: {} as Options,
+      instanceId: "instance",
+      databaseId: "db",
+      schemaName: "public",
+      choice: "all",
+    });
+
+    expect(executeSqlCmdsAsIamUserStub).to.be.calledOnce;
+    expect(executeSchemaMigrationStub).to.not.be.called;
+
+    const args = executeSqlCmdsAsIamUserStub.firstCall.args;
+    expect(args[3]).to.deep.equal([
+      'SET ROLE "firebaseowner_db_public"',
+      "CREATE TABLE a",
+      "DROP TABLE b",
+    ]);
+  });
+
+  it("should execute only safe commands via FDC API when fdcapimigration experiment is enabled and choice is safe", async () => {
+    isEnabledStub.withArgs("fdcapimigration").returns(true);
+    await handleIncompatibleSchemaError({
+      schema,
+      incompatibleSchemaError,
+      options: {} as Options,
+      instanceId: "instance",
+      databaseId: "db",
+      schemaName: "public",
+      choice: "safe",
+    });
+
+    expect(executeSchemaMigrationStub).to.be.calledOnce;
+    expect(executeSqlCmdsAsIamUserStub).to.not.be.called;
+
+    const args = executeSchemaMigrationStub.firstCall.args;
+    expect(args[0]).to.equal("projects/p/locations/l/services/s");
+    expect(args[1]).to.deep.equal([{ sql: "CREATE TABLE a", destructive: false }]);
+  });
+
+  it("should not execute any commands when choice is none", async () => {
+    await handleIncompatibleSchemaError({
+      schema,
+      incompatibleSchemaError,
+      options: {} as Options,
+      instanceId: "instance",
+      databaseId: "db",
+      schemaName: "public",
+      choice: "none",
+    });
+
+    expect(executeSchemaMigrationStub).to.not.be.called;
+    expect(executeSqlCmdsAsIamUserStub).to.not.be.called;
+  });
+});

src/dataconnect/schemaMigration.ts

@@ -1,8 +1,9 @@
+import * as experiments from "../experiments";
 import * as clc from "colorette";
 import { format } from "sql-formatter";
 
 import { IncompatibleSqlSchemaError, Diff, MAIN_SCHEMA_ID, SchemaValidation } from "./types";
-import { getSchema, upsertSchema, deleteConnector } from "./client";
+import { getSchema, upsertSchema, deleteConnector, executeSchemaMigration } from "./client";
 import {
   getIAMUser,
   executeSqlCmdsAsIamUser,
@@ -251,6 +252,7 @@ export async function migrateSchema(args: {
         databaseId,
         instanceId,
         schemaName,
+        schema,
         incompatibleSchemaError: incompatible,
         choice: migrationMode,
       });
@@ -300,6 +302,7 @@ export async function migrateSchema(args: {
           databaseId,
           instanceId,
           schemaName,
+          schema,
           incompatibleSchemaError: incompatible,
           choice: migrationMode,
         });
@@ -451,15 +454,17 @@ function suggestedCommand(serviceName: string, invalidConnectorNames: string[]):
   return `firebase deploy --only ${onlys}`;
 }
 
-async function handleIncompatibleSchemaError(args: {
+export async function handleIncompatibleSchemaError(args: {
+  schema: Schema;
   incompatibleSchemaError: IncompatibleSqlSchemaError;
   options: Options;
   instanceId: string;
   databaseId: string;
   schemaName: string;
   choice: "all" | "safe" | "none";
 }): Promise<Diff[]> {
-  const { incompatibleSchemaError, options, instanceId, databaseId, schemaName, choice } = args;
+  const { schema, incompatibleSchemaError, options, instanceId, databaseId, schemaName, choice } =
+    args;
   const commandsToExecute = incompatibleSchemaError.diffs.filter((d) => {
     switch (choice) {
       case "all":
@@ -523,16 +528,26 @@ async function handleIncompatibleSchemaError(args: {
     }
 
     if (commandsToExecuteByOwner.length) {
-      await executeSqlCmdsAsIamUser(
-        options,
-        instanceId,
-        databaseId,
-        [
-          `SET ROLE "${firebaseowner(databaseId, schemaName)}"`,
-          ...commandsToExecuteByOwner.map((d) => d.sql),
-        ],
-        /** silent=*/ false,
-      );
+      if (experiments.isEnabled("fdcapimigration")) {
+        logLabeledBullet(
+          "dataconnect",
+          `[EXPERIMENTAL] Delegating SQL execution to FDC Backend...`,
+        );
+
+        const serviceName = serviceNameFromSchema(schema);
+        await executeSchemaMigration(serviceName, commandsToExecuteByOwner);
+      } else {
+        await executeSqlCmdsAsIamUser(
+          options,
+          instanceId,
+          databaseId,
+          [
+            `SET ROLE "${firebaseowner(databaseId, schemaName)}"`,
+            ...commandsToExecuteByOwner.map((d) => d.sql),
+          ],
+          /** silent=*/ false,
+        );
+      }
       return incompatibleSchemaError.diffs;
     }
   }

src/experiments.ts

@@ -255,6 +255,12 @@ export const ALL_EXPERIMENTS = experiments({
     default: false,
     public: false,
   },
+  fdcapimigration: {
+    shortDescription: "Enable the FDC API schema migration path.",
+    fullDescription: "API based Schema Migration behind experimental flag.",
+    default: true,
+    public: false,
+  },
 });
 
 export type ExperimentName = keyof typeof ALL_EXPERIMENTS;