[FDC] Generate schema during init dataconnect flow

CHANGELOG.md

@@ -9,6 +9,7 @@
 - Fixed multi tenancy support for SSO users in the auth emulator (#8544)
 - Add SDK autoinit capabilities to App Hosting emulator. (#8582)
 - Provision App Hosting compute service account during init flow. (#8580)
+- `firebase init dataconnect` is now integrated with Gemini in Firebase API to generate Schema based on description. (#8596)
 - Updated the Firebase Data Connect local toolkit to v2.6.1, which includes the following changes: (#8598)
   - Fixed a bug where `@transaction` continued to execute after errors and produces more confusing errors.
   - Fixed a bug where rolled-back fields in `@transaction` returned bogus data that had been invalidated by the rollback. They now correctly return null with a corresponding error "(rolled back)".

src/dataconnect/ensureApis.ts

@@ -1,15 +1,20 @@
 import * as api from "../api";
 import { ensure } from "../ensureApiEnabled";
 
+const prefix = "dataconnect";
+
 export async function ensureApis(projectId: string): Promise<void> {
-  const prefix = "dataconnect";
   await ensure(projectId, api.dataconnectOrigin(), prefix);
   await ensure(projectId, api.cloudSQLAdminOrigin(), prefix);
   await ensure(projectId, api.computeOrigin(), prefix);
 }
 
 export async function ensureSparkApis(projectId: string): Promise<void> {
-  const prefix = "dataconnect";
   // These are the APIs that can be enabled without a billing account.
   await ensure(projectId, api.cloudSQLAdminOrigin(), prefix);
+  await ensure(projectId, api.cloudAiCompanionOrigin(), prefix);
+}
+
+export async function ensureGIFApis(projectId: string): Promise<void> {
+  await ensure(projectId, api.cloudAiCompanionOrigin(), prefix);
 }

src/gif/fdcExperience.spec.ts

@@ -0,0 +1,67 @@
+import { expect } from "chai";
+import { extractCodeBlock } from "./fdcExperience.js"; // Adjust the path to your file
+
+describe("extractCodeBlock", () => {
+  it("should extract a basic GraphQL query block", () => {
+    const text =
+      'Here is a GraphQL query:\n```graphql\nquery GetUser { user(id: "1") { name email } }\n```\nThanks!';
+    const expected = 'query GetUser { user(id: "1") { name email } }';
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract a multi-line GraphQL mutation block", () => {
+    const text = `
+    Some preamble.
+    \`\`\`graphql
+    mutation CreatePost($title: String!, $content: String!) {
+      createPost(title: $title, content: $content) {
+        id
+        title
+      }
+    }
+    \`\`\`
+    Followed by some description.
+    `;
+    const expected = `mutation CreatePost($title: String!, $content: String!) {
+      createPost(title: $title, content: $content) {
+        id
+        title
+      }
+    }`;
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract a GraphQL fragment block", () => {
+    const text = "```graphql\nfragment UserFields on User { id name }\n```";
+    const expected = "fragment UserFields on User { id name }";
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract an empty GraphQL code block", () => {
+    const text = "```graphql\n\n```";
+    const expected = "";
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract a GraphQL schema definition block", () => {
+    const text = `
+    \`\`\`graphql
+    type Query {
+      hello: String
+    }
+
+    schema {
+      query: Query
+    }
+    \`\`\`
+    `;
+    const expected = `type Query {
+      hello: String
+    }
+
+    schema {
+      query: Query
+    }`;
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+});

src/gif/fdcExperience.ts

@@ -1,5 +1,6 @@
 import { Client } from "../apiv2";
 import { cloudCompanionOrigin } from "../api";
+import { FirebaseError } from "../error";
 
 const apiClient = new Client({ urlPrefix: cloudCompanionOrigin(), auth: true });
 const schemaGeneratorExperience = "/appeco/firebase/fdc-schema-generator";
@@ -128,3 +129,17 @@
   );
   return res.body.output.messages[0].content;
 }
+
+/**
+ * extractCodeBlock extracts the code block from the generated response.
+ * @param text the generated response from the service.
+ * @return the code block from the generated response.
+ */
+export function extractCodeBlock(text: string): string {
+  const regex = /```(?:[a-z]+\n)?([\s\S]*?)```/m;
+  const match = text.match(regex);
+  if (match && match[1]) {
+    return match[1].trim();
+  }
+  throw new FirebaseError(`No code block found in the generated response: ${text}`);
+}

src/init/features/dataconnect/index.ts

@@ -8,7 +8,7 @@
 import { provisionCloudSql } from "../../../dataconnect/provisionCloudSql";
 import { checkFreeTrialInstanceUsed, upgradeInstructions } from "../../../dataconnect/freeTrial";
 import * as cloudsql from "../../../gcp/cloudsql/cloudsqladmin";
-import { ensureApis, ensureSparkApis } from "../../../dataconnect/ensureApis";
+import { ensureApis, ensureGIFApis, ensureSparkApis } from "../../../dataconnect/ensureApis";
 import {
   listLocations,
   listAllServices,
@@ -19,10 +19,11 @@
 import { parseCloudSQLInstanceName, parseServiceName } from "../../../dataconnect/names";
 import { logger } from "../../../logger";
 import { readTemplateSync } from "../../../templates";
-import { logBullet, envOverride } from "../../../utils";
+import { logBullet, envOverride, promiseWithSpinner } from "../../../utils";
 import { isBillingEnabled } from "../../../gcp/cloudbilling";
 import * as sdk from "./sdk";
 import { getPlatformFromFolder } from "../../../dataconnect/fileUtils";
+import { extractCodeBlock, generateSchema } from "../../../gif/fdcExperience";
 
 const DATACONNECT_YAML_TEMPLATE = readTemplateSync("init/dataconnect/dataconnect.yaml");
 const CONNECTOR_YAML_TEMPLATE = readTemplateSync("init/dataconnect/connector.yaml");
@@ -33,7 +34,7 @@
 // serviceEnvVar is used by Firebase Console to specify which service to import.
 // It should be in the form <location>/<serviceId>
 // It must be an existing service - if set to anything else, we'll ignore it.
 const serviceEnvVar = () => envOverride("FDC_SERVICE", "");
 
 export interface RequiredInfo {
   serviceId: string;
@@ -76,7 +77,7 @@
 
 // askQuestions prompts the user about the Data Connect service they want to init. Any prompting
 // logic should live here, and _no_ actuation logic should live here.
 export async function askQuestions(setup: Setup): Promise<void> {
   const hasBilling = await isBillingEnabled(setup);
   if (setup.projectId) {
     hasBilling ? await ensureApis(setup.projectId) : await ensureSparkApis(setup.projectId);
@@ -110,8 +111,7 @@
       default: true,
     }));
   if (shouldConfigureBackend) {
-    // TODO: Prompt for app idea and use GiF backend to generate them.
-    info = await promptForService(info);
+    info = await promptForSchema(setup, info);
     info = await promptForCloudSQL(setup, info);
 
     info.shouldProvisionCSQL = !!(
@@ -130,10 +130,10 @@
 
 // actuate writes product specific files and makes product specifc API calls.
 // It does not handle writing firebase.json and .firebaserc
 export async function actuate(setup: Setup, config: Config, options: any): Promise<void> {
   // Most users will want to persist data between emulator runs, so set this to a reasonable default.
   const dir: string = config.get("dataconnect.source", "dataconnect");
   const dataDir = config.get("emulators.dataconnect.dataDir", `${dir}/.dataconnect/pgliteData`);
   config.set("emulators.dataconnect.dataDir", dataDir);
 
   const info = setup.featureInfo?.dataconnect;
@@ -425,12 +425,33 @@
   return info;
 }
 
-async function promptForService(info: RequiredInfo): Promise<RequiredInfo> {
+async function promptForSchema(setup: Setup, info: RequiredInfo): Promise<RequiredInfo> {
   if (info.serviceId === "") {
     info.serviceId = await input({
       message: "What ID would you like to use for this service?",
       default: basename(process.cwd()),
     });
+    logBullet(
+      `Check out the terms of service of Gemini in Firebase https://firebase.google.com/docs/gemini-in-firebase/set-up-gemini#required-permissions`,
+    );
+    if (
+      await confirm({
+        message: `Do you want Gemini to help generate the schema?`,
+        default: true,
+      })
+    ) {
+      await ensureGIFApis(setup.projectId!);
+      const prompt = await input({
+        message: "Describe the app you are building:",
+        default: "movie rating app",
+      });
+      const schema = await promiseWithSpinner(
+        () => generateSchema(prompt, setup.projectId!),
+        "Generating the Data Connect Schema...",
+      );
+      info.schemaGql = [{ path: "schema.gql", content: extractCodeBlock(schema) }];
+      info.connectors = [emptyConnector];
+    }
   }
   return info;
 }