Enhance PostgreSQL schema handling in Data Connect

mtr002 · mtr002 · commit 7ebe8dc27d8a · 2026-03-10T13:07:49.000-07:00
- Added support for specifying a custom PostgreSQL schema name in the configuration.
- Updated the `getIdentifiers` function to return the schema name, defaulting to 'public' if not specified.
- Modified relevant functions to utilize the schema name for database operations.
- Added tests to validate schema name handling in various scenarios.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1 +1,2 @@
 - Added `--only` flag for `emulators:export` (#4033)
+- Added support for custom PostgreSQL schema names in Data Connect. The `schema` field in `dataconnect.yaml` under `datasource.postgresql` is now respected during deploy, migration, SQL setup, and SQL shell commands. (#10059)
diff --git a/schema/dataconnect-yaml.json b/schema/dataconnect-yaml.json
@@ -10,6 +10,10 @@
           "type": "string",
           "description": "The name of the PostgreSQL database."
         },
+        "schema": {
+          "type": "string",
+          "description": "The PostgreSQL schema name. Defaults to 'public' if not specified."
+        },
         "cloudSql": {
           "additionalProperties": false,
           "type": "object",
diff --git a/src/commands/dataconnect-sql-setup.ts b/src/commands/dataconnect-sql-setup.ts
@@ -6,7 +6,6 @@ import { requireAuth } from "../requireAuth";
 import { requirePermissions } from "../requirePermissions";
 import { ensureApis } from "../dataconnect/ensureApis";
 import { setupSQLPermissions, getSchemaMetadata } from "../gcp/cloudsql/permissionsSetup";
-import { DEFAULT_SCHEMA } from "../gcp/cloudsql/permissions";
 import { getIdentifiers, ensureServiceIsConnectedToCloudSql } from "../dataconnect/schemaMigration";
 import { setupIAMUsers } from "../gcp/cloudsql/connect";
 import { pickOneService } from "../dataconnect/load";
@@ -45,19 +44,20 @@ export const command = new Command("dataconnect:sql:setup")
       );
     }
 
-    const { serviceName, instanceName, databaseId } = getIdentifiers(
+    const { serviceName, instanceName, databaseId, schemaName } = getIdentifiers(
       mainSchema(serviceInfo.schemas),
     );
     await ensureServiceIsConnectedToCloudSql(
       serviceName,
       instanceName,
       databaseId,
       /* linkIfNotConnected=*/ true,
+      schemaName,
     );
 
     // Setup the IAM user for the current identity.
     await setupIAMUsers(instanceId, options);
 
-    const schemaInfo = await getSchemaMetadata(instanceId, databaseId, DEFAULT_SCHEMA, options);
+    const schemaInfo = await getSchemaMetadata(instanceId, databaseId, schemaName, options);
     await setupSQLPermissions(instanceId, databaseId, schemaInfo, options);
   });
diff --git a/src/commands/dataconnect-sql-shell.ts b/src/commands/dataconnect-sql-shell.ts
@@ -104,7 +104,7 @@ export const command = new Command("dataconnect:sql:shell")
       options.service,
       options.location,
     );
-    const { instanceId, databaseId } = getIdentifiers(mainSchema(serviceInfo.schemas));
+    const { instanceId, databaseId, schemaName } = getIdentifiers(mainSchema(serviceInfo.schemas));
     const { user: username } = await getIAMUser(options);
     const instance = await cloudSqlAdminClient.getInstance(projectId, instanceId);
 
@@ -130,6 +130,9 @@ export const command = new Command("dataconnect:sql:shell")
     });
     const conn: pg.PoolClient = await pool.connect();
 
+    // Set search_path to the configured PostgreSQL schema so unqualified table names resolve correctly.
+    await conn.query(`SET search_path TO "${schemaName}"`);
+
     logger.info(`Logged in as ${username}`);
     logger.info(clc.cyan("Welcome to Data Connect Cloud SQL Shell"));
     logger.info(
diff --git a/src/dataconnect/schemaMigration.spec.ts b/src/dataconnect/schemaMigration.spec.ts
@@ -1,5 +1,5 @@
 import { expect } from "chai";
-import { serviceNameFromSchema } from "./schemaMigration";
+import { serviceNameFromSchema, getIdentifiers } from "./schemaMigration";
 import { Schema } from "./types";
 
 describe("serviceNameFromSchema", () => {
@@ -40,3 +40,88 @@ describe("serviceNameFromSchema", () => {
     expect(serviceName).to.equal("projects/project-id/locations/us-central1/services/service-id");
   });
 });
+
+describe("getIdentifiers", () => {
+  it("should return custom schema name when specified", () => {
+    const schema: Schema = {
+      name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
+      datasources: [
+        {
+          postgresql: {
+            database: "fdcdb",
+            schema: "movies",
+            cloudSql: {
+              instance: "projects/project-id/locations/us-east4/instances/my-instance",
+            },
+          },
+        },
+      ],
+      source: {},
+    };
+    const ids = getIdentifiers(schema);
+    expect(ids.schemaName).to.equal("movies");
+    expect(ids.databaseId).to.equal("fdcdb");
+    expect(ids.instanceId).to.equal("my-instance");
+    expect(ids.instanceName).to.equal(
+      "projects/project-id/locations/us-east4/instances/my-instance",
+    );
+    expect(ids.serviceName).to.equal(
+      "projects/project-id/locations/us-central1/services/service-id",
+    );
+  });
+
+  it("should default schemaName to 'public' when not specified", () => {
+    const schema: Schema = {
+      name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
+      datasources: [
+        {
+          postgresql: {
+            database: "fdcdb",
+            cloudSql: {
+              instance: "projects/project-id/locations/us-east4/instances/my-instance",
+            },
+          },
+        },
+      ],
+      source: {},
+    };
+    const ids = getIdentifiers(schema);
+    expect(ids.schemaName).to.equal("public");
+  });
+
+  it("should throw if no database is specified", () => {
+    const schema: Schema = {
+      name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
+      datasources: [
+        {
+          postgresql: {
+            cloudSql: {
+              instance: "projects/project-id/locations/us-east4/instances/my-instance",
+            },
+          },
+        },
+      ],
+      source: {},
+    };
+    expect(() => getIdentifiers(schema)).to.throw(
+      "Data Connect schema must have a postgres datasource with a database name.",
+    );
+  });
+
+  it("should throw if no CloudSQL instance is specified", () => {
+    const schema: Schema = {
+      name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
+      datasources: [
+        {
+          postgresql: {
+            database: "fdcdb",
+          },
+        },
+      ],
+      source: {},
+    };
+    expect(() => getIdentifiers(schema)).to.throw(
+      "Data Connect schema must have a postgres datasource with a CloudSQL instance.",
+    );
+  });
+});
diff --git a/src/dataconnect/schemaMigration.ts b/src/dataconnect/schemaMigration.ts
@@ -35,11 +35,12 @@ import { checkBillingEnabled } from "../gcp/cloudbilling";
 async function setupSchemaIfNecessary(
   instanceId: string,
   databaseId: string,
+  schemaName: string,
   options: Options,
 ): Promise<SchemaSetupStatus.GreenField | SchemaSetupStatus.BrownField> {
   try {
     await setupIAMUsers(instanceId, options);
-    const schemaInfo = await getSchemaMetadata(instanceId, databaseId, DEFAULT_SCHEMA, options);
+    const schemaInfo = await getSchemaMetadata(instanceId, databaseId, schemaName, options);
     switch (schemaInfo.setupStatus) {
       case SchemaSetupStatus.BrownField:
       case SchemaSetupStatus.GreenField:
@@ -77,12 +78,13 @@ export async function diffSchema(
   setSchemaValidationMode(schema, validationMode);
   displayStartSchemaDiff(validationMode);
 
-  const { serviceName, instanceName, databaseId, instanceId } = getIdentifiers(schema);
+  const { serviceName, instanceName, databaseId, instanceId, schemaName } = getIdentifiers(schema);
   await ensureServiceIsConnectedToCloudSql(
     serviceName,
     instanceName,
     databaseId,
     /* linkIfNotConnected=*/ false,
+    schemaName,
   );
 
   let incompatible: IncompatibleSqlSchemaError | undefined = undefined;
@@ -162,12 +164,13 @@ export async function migrateSchema(args: {
   displayStartSchemaDiff(validationMode);
 
   const projectId = needProjectId(options);
-  const { serviceName, instanceId, instanceName, databaseId } = getIdentifiers(schema);
+  const { serviceName, instanceId, instanceName, databaseId, schemaName } = getIdentifiers(schema);
   await ensureServiceIsConnectedToCloudSql(
     serviceName,
     instanceName,
     databaseId,
     /* linkIfNotConnected=*/ true,
+    schemaName,
   );
 
   // Check if Cloud SQL instance is still being created.
@@ -196,7 +199,7 @@ export async function migrateSchema(args: {
   }
 
   // Make sure database is setup.
-  await setupSchemaIfNecessary(instanceId, databaseId, options);
+  await setupSchemaIfNecessary(instanceId, databaseId, schemaName, options);
 
   let diffs: Diff[] = [];
   try {
@@ -247,6 +250,7 @@ export async function migrateSchema(args: {
         options,
         databaseId,
         instanceId,
+        schemaName,
         incompatibleSchemaError: incompatible,
         choice: migrationMode,
       });
@@ -295,6 +299,7 @@ export async function migrateSchema(args: {
           options,
           databaseId,
           instanceId,
+          schemaName,
           incompatibleSchemaError: incompatible,
           choice: migrationMode,
         });
@@ -349,17 +354,23 @@ export async function grantRoleToUserInSchema(options: Options, schema: Schema)
   const role = options.role as string;
   const email = options.email as string;
 
-  const { serviceName, instanceId, instanceName, databaseId } = getIdentifiers(schema);
+  const { serviceName, instanceId, instanceName, databaseId, schemaName } = getIdentifiers(schema);
 
   await ensureServiceIsConnectedToCloudSql(
     serviceName,
     instanceName,
     databaseId,
     /* linkIfNotConnected=*/ false,
+    schemaName,
   );
 
   // Make sure we have the right setup for the requested role grant.
-  const schemaSetupStatus = await setupSchemaIfNecessary(instanceId, databaseId, options);
+  const schemaSetupStatus = await setupSchemaIfNecessary(
+    instanceId,
+    databaseId,
+    schemaName,
+    options,
+  );
 
   // Edge case: we can't grant firebase owner unless database is greenfield.
   if (schemaSetupStatus !== SchemaSetupStatus.GreenField && role === "owner") {
@@ -369,7 +380,7 @@ export async function grantRoleToUserInSchema(options: Options, schema: Schema)
   }
 
   // Grant the role to the user.
-  await grantRoleTo(options, instanceId, databaseId, role, email);
+  await grantRoleTo(options, instanceId, databaseId, role, email, schemaName);
 }
 
 function diffsEqual(x: Diff[], y: Diff[]): boolean {
@@ -399,6 +410,7 @@ export function getIdentifiers(schema: Schema): {
   instanceName: string;
   instanceId: string;
   databaseId: string;
+  schemaName: string;
   serviceName: string;
 } {
   const postgresDatasource = schema.datasources.find((d) => d.postgresql);
@@ -415,11 +427,13 @@ export function getIdentifiers(schema: Schema): {
     );
   }
   const instanceId = instanceName.split("/").pop()!;
+  const schemaName = postgresDatasource?.postgresql?.schema || DEFAULT_SCHEMA;
   const serviceName = serviceNameFromSchema(schema);
   return {
     databaseId,
     instanceId,
     instanceName,
+    schemaName,
     serviceName,
   };
 }
@@ -442,9 +456,10 @@ async function handleIncompatibleSchemaError(args: {
   options: Options;
   instanceId: string;
   databaseId: string;
+  schemaName: string;
   choice: "all" | "safe" | "none";
 }): Promise<Diff[]> {
-  const { incompatibleSchemaError, options, instanceId, databaseId, choice } = args;
+  const { incompatibleSchemaError, options, instanceId, databaseId, schemaName, choice } = args;
   const commandsToExecute = incompatibleSchemaError.diffs.filter((d) => {
     switch (choice) {
       case "all":
@@ -467,7 +482,7 @@ async function handleIncompatibleSchemaError(args: {
         ${diffsToString(commandsToExecuteBySuperUser)}`);
     }
 
-    const schemaInfo = await getSchemaMetadata(instanceId, databaseId, DEFAULT_SCHEMA, options);
+    const schemaInfo = await getSchemaMetadata(instanceId, databaseId, schemaName, options);
     if (schemaInfo.setupStatus !== SchemaSetupStatus.GreenField) {
       throw new FirebaseError(
         `Brownfield database are protected from SQL changes by Data Connect.\n` +
@@ -482,7 +497,7 @@ async function handleIncompatibleSchemaError(args: {
         options,
         instanceId,
         databaseId,
-        firebaseowner(databaseId),
+        firebaseowner(databaseId, schemaName),
         (await getIAMUser(options)).user,
       ))
     ) {
@@ -493,7 +508,7 @@ async function handleIncompatibleSchemaError(args: {
       }
       const account = (await requireAuth(options))!;
       logLabeledBullet("dataconnect", `Granting firebaseowner role to myself ${account}...`);
-      await grantRoleTo(options, instanceId, databaseId, "owner", account);
+      await grantRoleTo(options, instanceId, databaseId, "owner", account, schemaName);
     }
 
     if (commandsToExecuteBySuperUser.length) {
@@ -512,7 +527,10 @@ async function handleIncompatibleSchemaError(args: {
         options,
         instanceId,
         databaseId,
-        [`SET ROLE "${firebaseowner(databaseId)}"`, ...commandsToExecuteByOwner.map((d) => d.sql)],
+        [
+          `SET ROLE "${firebaseowner(databaseId, schemaName)}"`,
+          ...commandsToExecuteByOwner.map((d) => d.sql),
+        ],
         /** silent=*/ false,
       );
       return incompatibleSchemaError.diffs;
@@ -645,6 +663,7 @@ export async function ensureServiceIsConnectedToCloudSql(
   instanceName: string,
   databaseId: string,
   linkIfNotConnected: boolean,
+  schemaName?: string,
 ): Promise<void> {
   let currentSchema = await getSchema(serviceName);
   let postgresql = currentSchema?.datasources?.find((d) => d.postgresql)?.postgresql;
@@ -720,6 +739,7 @@ export async function ensureServiceIsConnectedToCloudSql(
   try {
     postgresql.schemaValidation = "STRICT";
     postgresql.database = databaseId;
+    postgresql.schema = schemaName;
     postgresql.cloudSql = { instance: instanceName };
     await upsertSchema(currentSchema, /** validateOnly=*/ false);
   } catch (err: any) {
diff --git a/src/dataconnect/types.ts b/src/dataconnect/types.ts
@@ -36,6 +36,7 @@ export type SchemaValidation = "STRICT" | "COMPATIBLE";
 export interface PostgreSql {
   ephemeral?: boolean;
   database?: string;
+  schema?: string;
   cloudSql?: CloudSqlInstance;
   schemaValidation?: SchemaValidation | "NONE" | "SQL_SCHEMA_VALIDATION_UNSPECIFIED";
   schemaMigration?: "MIGRATE_COMPATIBLE";
@@ -144,6 +145,7 @@ export interface SchemaYaml {
 export interface DatasourceYaml {
   postgresql?: {
     database: string;
+    schema?: string;
     cloudSql: {
       instanceId: string;
     };
@@ -223,6 +225,7 @@ export function toDatasource(
     return {
       postgresql: {
         database: ds.postgresql.database,
+        schema: ds.postgresql.schema,
         cloudSql: {
           instance: `projects/${projectId}/locations/${locationId}/instances/${ds.postgresql.cloudSql.instanceId}`,
         },
diff --git a/src/gcp/cloudsql/permissionsSetup.ts b/src/gcp/cloudsql/permissionsSetup.ts

Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`	`1`	- Added `--only` flag for `emulators:export` (#4033)
	`2`	+- Added support for custom PostgreSQL schema names in Data Connect. The `schema` field in `dataconnect.yaml` under `datasource.postgresql` is now respected during deploy, migration, SQL setup, and SQL shell commands. (#10059)