Skip to content

Commit e4764b7

Browse files
mtr002mtr002
authored
Enhance PostgreSQL schema handling in Data Connect (#10059)
- 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.
1 parent 68caa46 commit e4764b7

File tree

8 files changed

+138
-20
lines changed

8 files changed

+138
-20
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
- Add support for VPC direct connect in GCF 2nd gen (#10033)
22
- Added `--only` flag for `emulators:export` (#4033)
3+
- Added support for custom PostgreSQL schema names in Data Connect. (#9271)

‎schema/dataconnect-yaml.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@
1010
"type": "string",
1111
"description": "The name of the PostgreSQL database."
1212
},
13+
"schema": {
14+
"type": "string",
15+
"description": "The PostgreSQL schema name. Defaults to 'public' if not specified."
16+
},
1317
"cloudSql": {
1418
"additionalProperties": false,
1519
"type": "object",

‎src/commands/dataconnect-sql-setup.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import { requireAuth } from "../requireAuth";
66
import { requirePermissions } from "../requirePermissions";
77
import { ensureApis } from "../dataconnect/ensureApis";
88
import { setupSQLPermissions, getSchemaMetadata } from "../gcp/cloudsql/permissionsSetup";
9-
import { DEFAULT_SCHEMA } from "../gcp/cloudsql/permissions";
109
import { getIdentifiers, ensureServiceIsConnectedToCloudSql } from "../dataconnect/schemaMigration";
1110
import { setupIAMUsers } from "../gcp/cloudsql/connect";
1211
import { pickOneService } from "../dataconnect/load";
@@ -45,19 +44,20 @@ export const command = new Command("dataconnect:sql:setup")
4544
);
4645
}
4746

48-
const { serviceName, instanceName, databaseId } = getIdentifiers(
47+
const { serviceName, instanceName, databaseId, schemaName } = getIdentifiers(
4948
mainSchema(serviceInfo.schemas),
5049
);
5150
await ensureServiceIsConnectedToCloudSql(
5251
serviceName,
5352
instanceName,
5453
databaseId,
5554
/* linkIfNotConnected=*/ true,
55+
schemaName,
5656
);
5757

5858
// Setup the IAM user for the current identity.
5959
await setupIAMUsers(instanceId, options);
6060

61-
const schemaInfo = await getSchemaMetadata(instanceId, databaseId, DEFAULT_SCHEMA, options);
61+
const schemaInfo = await getSchemaMetadata(instanceId, databaseId, schemaName, options);
6262
await setupSQLPermissions(instanceId, databaseId, schemaInfo, options);
6363
});

‎src/commands/dataconnect-sql-shell.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ export const command = new Command("dataconnect:sql:shell")
104104
options.service,
105105
options.location,
106106
);
107-
const { instanceId, databaseId } = getIdentifiers(mainSchema(serviceInfo.schemas));
107+
const { instanceId, databaseId, schemaName } = getIdentifiers(mainSchema(serviceInfo.schemas));
108108
const { user: username } = await getIAMUser(options);
109109
const instance = await cloudSqlAdminClient.getInstance(projectId, instanceId);
110110

@@ -130,6 +130,9 @@ export const command = new Command("dataconnect:sql:shell")
130130
});
131131
const conn: pg.PoolClient = await pool.connect();
132132

133+
// Set search_path to the configured PostgreSQL schema so unqualified table names resolve correctly.
134+
await conn.query(`SET search_path TO "${schemaName}"`);
135+
133136
logger.info(`Logged in as ${username}`);
134137
logger.info(clc.cyan("Welcome to Data Connect Cloud SQL Shell"));
135138
logger.info(

‎src/dataconnect/schemaMigration.spec.ts‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { expect } from "chai";
2-
import { serviceNameFromSchema } from "./schemaMigration";
2+
import { serviceNameFromSchema, getIdentifiers } from "./schemaMigration";
33
import { Schema } from "./types";
44

55
describe("serviceNameFromSchema", () => {
@@ -40,3 +40,88 @@ describe("serviceNameFromSchema", () => {
4040
expect(serviceName).to.equal("projects/project-id/locations/us-central1/services/service-id");
4141
});
4242
});
43+
44+
describe("getIdentifiers", () => {
45+
it("should return custom schema name when specified", () => {
46+
const schema: Schema = {
47+
name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
48+
datasources: [
49+
{
50+
postgresql: {
51+
database: "fdcdb",
52+
schema: "movies",
53+
cloudSql: {
54+
instance: "projects/project-id/locations/us-east4/instances/my-instance",
55+
},
56+
},
57+
},
58+
],
59+
source: {},
60+
};
61+
const ids = getIdentifiers(schema);
62+
expect(ids.schemaName).to.equal("movies");
63+
expect(ids.databaseId).to.equal("fdcdb");
64+
expect(ids.instanceId).to.equal("my-instance");
65+
expect(ids.instanceName).to.equal(
66+
"projects/project-id/locations/us-east4/instances/my-instance",
67+
);
68+
expect(ids.serviceName).to.equal(
69+
"projects/project-id/locations/us-central1/services/service-id",
70+
);
71+
});
72+
73+
it("should default schemaName to 'public' when not specified", () => {
74+
const schema: Schema = {
75+
name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
76+
datasources: [
77+
{
78+
postgresql: {
79+
database: "fdcdb",
80+
cloudSql: {
81+
instance: "projects/project-id/locations/us-east4/instances/my-instance",
82+
},
83+
},
84+
},
85+
],
86+
source: {},
87+
};
88+
const ids = getIdentifiers(schema);
89+
expect(ids.schemaName).to.equal("public");
90+
});
91+
92+
it("should throw if no database is specified", () => {
93+
const schema: Schema = {
94+
name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
95+
datasources: [
96+
{
97+
postgresql: {
98+
cloudSql: {
99+
instance: "projects/project-id/locations/us-east4/instances/my-instance",
100+
},
101+
},
102+
},
103+
],
104+
source: {},
105+
};
106+
expect(() => getIdentifiers(schema)).to.throw(
107+
"Data Connect schema must have a postgres datasource with a database name.",
108+
);
109+
});
110+
111+
it("should throw if no CloudSQL instance is specified", () => {
112+
const schema: Schema = {
113+
name: "projects/project-id/locations/us-central1/services/service-id/schemas/main",
114+
datasources: [
115+
{
116+
postgresql: {
117+
database: "fdcdb",
118+
},
119+
},
120+
],
121+
source: {},
122+
};
123+
expect(() => getIdentifiers(schema)).to.throw(
124+
"Data Connect schema must have a postgres datasource with a CloudSQL instance.",
125+
);
126+
});
127+
});

‎src/dataconnect/schemaMigration.ts‎

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,12 @@ import { checkBillingEnabled } from "../gcp/cloudbilling";
3535
async function setupSchemaIfNecessary(
3636
instanceId: string,
3737
databaseId: string,
38+
schemaName: string,
3839
options: Options,
3940
): Promise<SchemaSetupStatus.GreenField | SchemaSetupStatus.BrownField> {
4041
try {
4142
await setupIAMUsers(instanceId, options);
42-
const schemaInfo = await getSchemaMetadata(instanceId, databaseId, DEFAULT_SCHEMA, options);
43+
const schemaInfo = await getSchemaMetadata(instanceId, databaseId, schemaName, options);
4344
switch (schemaInfo.setupStatus) {
4445
case SchemaSetupStatus.BrownField:
4546
case SchemaSetupStatus.GreenField:
@@ -77,12 +78,13 @@ export async function diffSchema(
7778
setSchemaValidationMode(schema, validationMode);
7879
displayStartSchemaDiff(validationMode);
7980

80-
const { serviceName, instanceName, databaseId, instanceId } = getIdentifiers(schema);
81+
const { serviceName, instanceName, databaseId, instanceId, schemaName } = getIdentifiers(schema);
8182
await ensureServiceIsConnectedToCloudSql(
8283
serviceName,
8384
instanceName,
8485
databaseId,
8586
/* linkIfNotConnected=*/ false,
87+
schemaName,
8688
);
8789

8890
let incompatible: IncompatibleSqlSchemaError | undefined = undefined;
@@ -162,12 +164,13 @@ export async function migrateSchema(args: {
162164
displayStartSchemaDiff(validationMode);
163165

164166
const projectId = needProjectId(options);
165-
const { serviceName, instanceId, instanceName, databaseId } = getIdentifiers(schema);
167+
const { serviceName, instanceId, instanceName, databaseId, schemaName } = getIdentifiers(schema);
166168
await ensureServiceIsConnectedToCloudSql(
167169
serviceName,
168170
instanceName,
169171
databaseId,
170172
/* linkIfNotConnected=*/ true,
173+
schemaName,
171174
);
172175

173176
// Check if Cloud SQL instance is still being created.
@@ -196,7 +199,7 @@ export async function migrateSchema(args: {
196199
}
197200

198201
// Make sure database is setup.
199-
await setupSchemaIfNecessary(instanceId, databaseId, options);
202+
await setupSchemaIfNecessary(instanceId, databaseId, schemaName, options);
200203

201204
let diffs: Diff[] = [];
202205
try {
@@ -247,6 +250,7 @@ export async function migrateSchema(args: {
247250
options,
248251
databaseId,
249252
instanceId,
253+
schemaName,
250254
incompatibleSchemaError: incompatible,
251255
choice: migrationMode,
252256
});
@@ -295,6 +299,7 @@ export async function migrateSchema(args: {
295299
options,
296300
databaseId,
297301
instanceId,
302+
schemaName,
298303
incompatibleSchemaError: incompatible,
299304
choice: migrationMode,
300305
});
@@ -349,17 +354,23 @@ export async function grantRoleToUserInSchema(options: Options, schema: Schema)
349354
const role = options.role as string;
350355
const email = options.email as string;
351356

352-
const { serviceName, instanceId, instanceName, databaseId } = getIdentifiers(schema);
357+
const { serviceName, instanceId, instanceName, databaseId, schemaName } = getIdentifiers(schema);
353358

354359
await ensureServiceIsConnectedToCloudSql(
355360
serviceName,
356361
instanceName,
357362
databaseId,
358363
/* linkIfNotConnected=*/ false,
364+
schemaName,
359365
);
360366

361367
// Make sure we have the right setup for the requested role grant.
362-
const schemaSetupStatus = await setupSchemaIfNecessary(instanceId, databaseId, options);
368+
const schemaSetupStatus = await setupSchemaIfNecessary(
369+
instanceId,
370+
databaseId,
371+
schemaName,
372+
options,
373+
);
363374

364375
// Edge case: we can't grant firebase owner unless database is greenfield.
365376
if (schemaSetupStatus !== SchemaSetupStatus.GreenField && role === "owner") {
@@ -369,7 +380,7 @@ export async function grantRoleToUserInSchema(options: Options, schema: Schema)
369380
}
370381

371382
// Grant the role to the user.
372-
await grantRoleTo(options, instanceId, databaseId, role, email);
383+
await grantRoleTo(options, instanceId, databaseId, role, email, schemaName);
373384
}
374385

375386
function diffsEqual(x: Diff[], y: Diff[]): boolean {
@@ -399,6 +410,7 @@ export function getIdentifiers(schema: Schema): {
399410
instanceName: string;
400411
instanceId: string;
401412
databaseId: string;
413+
schemaName: string;
402414
serviceName: string;
403415
} {
404416
const postgresDatasource = schema.datasources.find((d) => d.postgresql);
@@ -415,11 +427,13 @@ export function getIdentifiers(schema: Schema): {
415427
);
416428
}
417429
const instanceId = instanceName.split("/").pop()!;
430+
const schemaName = postgresDatasource?.postgresql?.schema || DEFAULT_SCHEMA;
418431
const serviceName = serviceNameFromSchema(schema);
419432
return {
420433
databaseId,
421434
instanceId,
422435
instanceName,
436+
schemaName,
423437
serviceName,
424438
};
425439
}
@@ -442,9 +456,10 @@ async function handleIncompatibleSchemaError(args: {
442456
options: Options;
443457
instanceId: string;
444458
databaseId: string;
459+
schemaName: string;
445460
choice: "all" | "safe" | "none";
446461
}): Promise<Diff[]> {
447-
const { incompatibleSchemaError, options, instanceId, databaseId, choice } = args;
462+
const { incompatibleSchemaError, options, instanceId, databaseId, schemaName, choice } = args;
448463
const commandsToExecute = incompatibleSchemaError.diffs.filter((d) => {
449464
switch (choice) {
450465
case "all":
@@ -467,7 +482,7 @@ async function handleIncompatibleSchemaError(args: {
467482
${diffsToString(commandsToExecuteBySuperUser)}`);
468483
}
469484

470-
const schemaInfo = await getSchemaMetadata(instanceId, databaseId, DEFAULT_SCHEMA, options);
485+
const schemaInfo = await getSchemaMetadata(instanceId, databaseId, schemaName, options);
471486
if (schemaInfo.setupStatus !== SchemaSetupStatus.GreenField) {
472487
throw new FirebaseError(
473488
`Brownfield database are protected from SQL changes by Data Connect.\n` +
@@ -482,7 +497,7 @@ async function handleIncompatibleSchemaError(args: {
482497
options,
483498
instanceId,
484499
databaseId,
485-
firebaseowner(databaseId),
500+
firebaseowner(databaseId, schemaName),
486501
(await getIAMUser(options)).user,
487502
))
488503
) {
@@ -493,7 +508,7 @@ async function handleIncompatibleSchemaError(args: {
493508
}
494509
const account = (await requireAuth(options))!;
495510
logLabeledBullet("dataconnect", `Granting firebaseowner role to myself ${account}...`);
496-
await grantRoleTo(options, instanceId, databaseId, "owner", account);
511+
await grantRoleTo(options, instanceId, databaseId, "owner", account, schemaName);
497512
}
498513

499514
if (commandsToExecuteBySuperUser.length) {
@@ -512,7 +527,10 @@ async function handleIncompatibleSchemaError(args: {
512527
options,
513528
instanceId,
514529
databaseId,
515-
[`SET ROLE "${firebaseowner(databaseId)}"`, ...commandsToExecuteByOwner.map((d) => d.sql)],
530+
[
531+
`SET ROLE "${firebaseowner(databaseId, schemaName)}"`,
532+
...commandsToExecuteByOwner.map((d) => d.sql),
533+
],
516534
/** silent=*/ false,
517535
);
518536
return incompatibleSchemaError.diffs;
@@ -645,6 +663,7 @@ export async function ensureServiceIsConnectedToCloudSql(
645663
instanceName: string,
646664
databaseId: string,
647665
linkIfNotConnected: boolean,
666+
schemaName?: string,
648667
): Promise<void> {
649668
let currentSchema = await getSchema(serviceName);
650669
let postgresql = currentSchema?.datasources?.find((d) => d.postgresql)?.postgresql;
@@ -720,6 +739,7 @@ export async function ensureServiceIsConnectedToCloudSql(
720739
try {
721740
postgresql.schemaValidation = "STRICT";
722741
postgresql.database = databaseId;
742+
postgresql.schema = schemaName;
723743
postgresql.cloudSql = { instance: instanceName };
724744
await upsertSchema(currentSchema, /** validateOnly=*/ false);
725745
} catch (err: any) {

‎src/dataconnect/types.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ export type SchemaValidation = "STRICT" | "COMPATIBLE";
3636
export interface PostgreSql {
3737
ephemeral?: boolean;
3838
database?: string;
39+
schema?: string;
3940
cloudSql?: CloudSqlInstance;
4041
schemaValidation?: SchemaValidation | "NONE" | "SQL_SCHEMA_VALIDATION_UNSPECIFIED";
4142
schemaMigration?: "MIGRATE_COMPATIBLE";
@@ -144,6 +145,7 @@ export interface SchemaYaml {
144145
export interface DatasourceYaml {
145146
postgresql?: {
146147
database: string;
148+
schema?: string;
147149
cloudSql: {
148150
instanceId: string;
149151
};
@@ -223,6 +225,7 @@ export function toDatasource(
223225
return {
224226
postgresql: {
225227
database: ds.postgresql.database,
228+
schema: ds.postgresql.schema,
226229
cloudSql: {
227230
instance: `projects/${projectId}/locations/${locationId}/instances/${ds.postgresql.cloudSql.instanceId}`,
228231
},

0 commit comments

Comments
 (0)