feat(drizzle-kit): add --omitSuffixForSchema to control schema suffix in introspected outputs

drizzle-kit/src/cli/commands/introspect.ts

@@ -61,6 +61,7 @@ export const introspectPostgres = async (
 	schemasFilter: string[],
 	prefix: Prefix,
 	entities: Entities,
+	omitSuffixForSchema: string,
 ) => {
 	const { preparePostgresDB } = await import('../connections');
 	const db = await preparePostgresDB(credentials);
@@ -108,7 +109,7 @@ export const introspectPostgres = async (
 	);
 
 	const schema = { id: originUUID, prevId: '', ...res } as PgSchema;
-	const ts = postgresSchemaToTypeScript(schema, casing);
+	const ts = postgresSchemaToTypeScript(schema, casing, omitSuffixForSchema);
 	const relationsTs = relationsToTypeScript(schema, casing);
 	const { internal, ...schemaWithoutInternals } = schema;
 
@@ -187,6 +188,7 @@ export const introspectGel = async (
 	schemasFilter: string[],
 	prefix: Prefix,
 	entities: Entities,
+	omitSuffixForSchema: string,
 ) => {
 	const { prepareGelDB } = await import('../connections');
 	const db = await prepareGelDB(credentials);
@@ -234,7 +236,7 @@ export const introspectGel = async (
 	);
 
 	const schema = { id: originUUID, prevId: '', ...res } as GelSchema;
-	const ts = gelSchemaToTypeScript(schema, casing);
+	const ts = gelSchemaToTypeScript(schema, casing, omitSuffixForSchema);
 	const relationsTs = relationsToTypeScript(schema, casing);
 	const { internal, ...schemaWithoutInternals } = schema;
 

drizzle-kit/src/cli/commands/utils.ts

@@ -461,6 +461,7 @@ export const preparePullConfig = async (
 		| {
 			dialect: 'postgresql';
 			credentials: PostgresCredentials;
+			omitSuffixForSchema: string;
 		}
 		| {
 			dialect: 'sqlite';
@@ -477,6 +478,7 @@ export const preparePullConfig = async (
 		| {
 			dialect: 'gel';
 			credentials?: GelCredentials;
+			omitSuffixForSchema: string;
 		}
 	) & {
 		out: string;

drizzle-kit/src/cli/schema.ts

@@ -47,6 +47,10 @@ const optionDriver = string()
 
 const optionCasing = string().enum('camelCase', 'snake_case').desc('Casing for serialization');
 
+const optionOmitSuffixForSchema = string().desc(
+	`Schema name to treat as suffixless; when schema equals this value, the In<Schema> suffix is not appended. Defaults to 'public'.`
+);
+
 export const generate = command({
 	name: 'generate',
 	options: {
@@ -476,6 +480,7 @@ export const pull = command({
 		out: optionOut,
 		breakpoints: optionBreakpoints,
 		casing: string('introspect-casing').enum('camel', 'preserve'),
+		omitSuffixForSchema: optionOmitSuffixForSchema,
 		...optionsFilters,
 		...optionsDatabaseCredentials,
 	},
@@ -502,6 +507,7 @@ export const pull = command({
 				'schemaFilters',
 				'extensionsFilters',
 				'tlsSecurity',
+				'omitSuffixForSchema',
 			],
 		);
 		return preparePullConfig(opts, from);
@@ -567,6 +573,7 @@ export const pull = command({
 					schemasFilter,
 					prefix,
 					entities,
+					config.omitSuffixForSchema
 				);
 			} else if (dialect === 'mysql') {
 				const { introspectMysql } = await import('./commands/introspect');
@@ -619,6 +626,7 @@ export const pull = command({
 					schemasFilter,
 					prefix,
 					entities,
+					config.omitSuffixForSchema
 				);
 			} else {
 				assertUnreachable(dialect);

drizzle-kit/src/cli/validations/common.ts

@@ -128,6 +128,7 @@ export const introspectParams = object({
 	introspect: object({
 		casing,
 	}).default({ casing: 'camel' }),
+	omitSuffixForSchema: string().optional().default('public'),
 });
 
 export type IntrospectParams = TypeOf<typeof introspectParams>;
@@ -142,6 +143,7 @@ export const configIntrospectCliSchema = object({
 	introspectCasing: union([literal('camel'), literal('preserve')]).default(
 		'camel',
 	),
+	omitSuffixForSchema: string().optional().default('public')
 });
 
 export const configGenerateSchema = object({

drizzle-kit/src/introspect-gel.ts

@@ -203,12 +203,12 @@ function generateIdentityParams(identity: Column['identity']) {
 	return `.generatedByDefaultAsIdentity(${paramsObj})`;
 }
 
-export const paramNameFor = (name: string, schema?: string) => {
-	const schemaSuffix = schema && schema !== 'public' ? `In${schema.capitalise()}` : '';
+export const paramNameFor = (name: string, schema?: string, omitSuffixForSchema = 'public') => {
+	const schemaSuffix = schema && schema !== omitSuffixForSchema ? `In${schema.capitalise()}` : '';
 	return `${name}${schemaSuffix}`;
 };
 
-export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) => {
+export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing, omitSuffixForSchema: string) => {
 	// collectFKs
 	Object.values(schema.tables).forEach((table) => {
 		Object.values(table.foreignKeys).forEach((fk) => {
@@ -314,7 +314,7 @@ export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) =>
 	// 	.map((it) => {
 	// 		const enumSchema = schemas[it.schema];
 	// 		// const func = schema || schema === "public" ? "gelTable" : schema;
-	// 		const paramName = paramNameFor(it.name, enumSchema);
+	// 		const paramName = paramNameFor(it.name, enumSchema, omitSuffixForSchema);
 
 	// 		const func = enumSchema ? `${enumSchema}.enum` : 'gelEnum';
 
@@ -329,7 +329,7 @@ export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) =>
 	// const sequencesStatements = Object.values(schema.sequences)
 	// 	.map((it) => {
 	// 		const seqSchema = schemas[it.schema];
-	// 		const paramName = paramNameFor(it.name, seqSchema);
+	// 		const paramName = paramNameFor(it.name, seqSchema, omitSuffixForSchema);
 
 	// 		const func = seqSchema ? `${seqSchema}.sequence` : 'gelSequence';
 
@@ -390,7 +390,7 @@ export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) =>
 
 	const tableStatements = Object.values(schema.tables).map((table) => {
 		const tableSchema = schemas[table.schema];
-		const paramName = paramNameFor(table.name, tableSchema);
+		const paramName = paramNameFor(table.name, tableSchema, omitSuffixForSchema);
 
 		const func = tableSchema ? `${tableSchema}.table` : 'gelTable';
 		let statement = `export const ${withCasing(paramName, casing)} = ${func}("${table.name}", {\n`;
@@ -403,6 +403,7 @@ export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) =>
 			schemas,
 			casing,
 			schema.internal,
+			omitSuffixForSchema,
 		);
 		statement += '}';
 
@@ -423,7 +424,7 @@ export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) =>
 			statement += ', ';
 			statement += '(table) => [';
 			statement += createTableIndexes(table.name, Object.values(table.indexes), casing);
-			statement += createTableFKs(Object.values(table.foreignKeys), schemas, casing);
+			statement += createTableFKs(Object.values(table.foreignKeys), schemas, casing, omitSuffixForSchema);
 			statement += createTablePKs(
 				Object.values(table.compositePrimaryKeys),
 				casing,
@@ -452,7 +453,7 @@ export const schemaToTypeScript = (schema: GelSchemaInternal, casing: Casing) =>
 	// 	.map((it) => {
 	// 		const viewSchema = schemas[it.schema];
 
-	// 		const paramName = paramNameFor(it.name, viewSchema);
+	// 		const paramName = paramNameFor(it.name, viewSchema, omitSuffixForSchema);
 
 	// 		const func = viewSchema
 	// 			? (it.materialized ? `${viewSchema}.materializedView` : `${viewSchema}.view`)
@@ -704,12 +705,13 @@ const column = (
 	casing: Casing,
 	defaultValue?: any,
 	internals?: GelKitInternals,
+	omitSuffixForSchema?: string,
 ) => {
 	const isExpression = internals?.tables[tableName]?.columns[name]?.isDefaultAnExpression ?? false;
 	const lowered = type.toLowerCase().replace('[]', '');
 
 	if (enumTypes.has(`${typeSchema}.${type.replace('[]', '')}`)) {
-		let out = `${withCasing(name, casing)}: ${withCasing(paramNameFor(type.replace('[]', ''), typeSchema), casing)}(${
+		let out = `${withCasing(name, casing)}: ${withCasing(paramNameFor(type.replace('[]', ''), typeSchema, omitSuffixForSchema), casing)}(${
 			dbColumnName({ name, casing })
 		})`;
 		return out;
@@ -839,6 +841,7 @@ const createTableColumns = (
 	schemas: Record<string, string>,
 	casing: Casing,
 	internals: GelKitInternals,
+	omitSuffixForSchema?: string,
 ): string => {
 	let statement = '';
 
@@ -866,6 +869,7 @@ const createTableColumns = (
 			casing,
 			it.default,
 			internals,
+			omitSuffixForSchema,
 		);
 		statement += '\t';
 		statement += columnStatement;
@@ -882,7 +886,7 @@ const createTableColumns = (
 		statement += it.generated ? `.generatedAlwaysAs(sql\`${it.generated.as}\`)` : '';
 
 		// const fks = fkByColumnName[it.name];
-		// Andrii: I switched it off until we will get a custom naem setting in references
+		// Andrii: I switched it off until we will get a custom name setting in references
 		// if (fks) {
 		// 	const fksStatement = fks
 		// 		.map((it) => {
@@ -894,7 +898,7 @@ const createTableColumns = (
 
 		// 			const paramsStr = objToStatement2(params);
 		// 			const tableSchema = schemas[it.schemaTo || ''];
-		// 			const paramName = paramNameFor(it.tableTo, tableSchema);
+		// 			const paramName = paramNameFor(it.tableTo, tableSchema, omitSuffixForSchema);
 		// 			if (paramsStr) {
 		// 				return `.references(()${typeSuffix} => ${
 		// 					withCasing(
@@ -1062,12 +1066,17 @@ const createTableChecks = (
 	return statement;
 };
 
-const createTableFKs = (fks: ForeignKey[], schemas: Record<string, string>, casing: Casing): string => {
+const createTableFKs = (
+	fks: ForeignKey[],
+	schemas: Record<string, string>,
+	casing: Casing,
+	omitSuffixForSchema?: string,
+): string => {
 	let statement = '';
 
 	fks.forEach((it) => {
 		const tableSchema = schemas[it.schemaTo || ''];
-		const paramName = paramNameFor(it.tableTo, tableSchema);
+		const paramName = paramNameFor(it.tableTo, tableSchema, omitSuffixForSchema);
 
 		const isSelf = it.tableTo === it.tableFrom;
 		const tableTo = isSelf ? 'table' : `${withCasing(paramName, casing)}`;

drizzle-kit/src/introspect-pg.ts

@@ -301,12 +301,12 @@ function generateIdentityParams(identity: Column['identity']) {
 	return `.generatedByDefaultAsIdentity(${paramsObj})`;
 }
 
-export const paramNameFor = (name: string, schema?: string) => {
-	const schemaSuffix = schema && schema !== 'public' ? `In${schema.capitalise()}` : '';
+export const paramNameFor = (name: string, schema?: string, omitSuffixForSchema = 'public') => {
+	const schemaSuffix = schema && schema !== omitSuffixForSchema ? `In${schema.capitalise()}` : '';
 	return `${name}${schemaSuffix}`;
 };
 
-export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) => {
+export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing, omitSuffixForSchema: string) => {
 	// collectFKs
 	Object.values(schema.tables).forEach((table) => {
 		Object.values(table.foreignKeys).forEach((fk) => {
@@ -431,7 +431,7 @@ export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) =>
 		.map((it) => {
 			const enumSchema = schemas[it.schema];
 			// const func = schema || schema === "public" ? "pgTable" : schema;
-			const paramName = paramNameFor(it.name, enumSchema);
+			const paramName = paramNameFor(it.name, enumSchema, omitSuffixForSchema);
 
 			const func = enumSchema ? `${enumSchema}.enum` : 'pgEnum';
 
@@ -446,7 +446,7 @@ export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) =>
 	const sequencesStatements = Object.values(schema.sequences)
 		.map((it) => {
 			const seqSchema = schemas[it.schema];
-			const paramName = paramNameFor(it.name, seqSchema);
+			const paramName = paramNameFor(it.name, seqSchema, omitSuffixForSchema);
 
 			const func = seqSchema ? `${seqSchema}.sequence` : 'pgSequence';
 
@@ -507,7 +507,7 @@ export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) =>
 
 	const tableStatements = Object.values(schema.tables).map((table) => {
 		const tableSchema = schemas[table.schema];
-		const paramName = paramNameFor(table.name, tableSchema);
+		const paramName = paramNameFor(table.name, tableSchema, omitSuffixForSchema);
 
 		const func = tableSchema ? `${tableSchema}.table` : 'pgTable';
 		let statement = `export const ${withCasing(paramName, casing)} = ${func}("${table.name}", {\n`;
@@ -519,6 +519,7 @@ export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) =>
 			schemas,
 			casing,
 			schema.internal,
+			omitSuffixForSchema,
 		);
 		statement += '}';
 
@@ -568,7 +569,7 @@ export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) =>
 		.map((it) => {
 			const viewSchema = schemas[it.schema];
 
-			const paramName = paramNameFor(it.name, viewSchema);
+			const paramName = paramNameFor(it.name, viewSchema, omitSuffixForSchema);
 
 			const func = viewSchema
 				? (it.materialized ? `${viewSchema}.materializedView` : `${viewSchema}.view`)
@@ -590,6 +591,7 @@ export const schemaToTypeScript = (schema: PgSchemaInternal, casing: Casing) =>
 				schemas,
 				casing,
 				schema.internal,
+				omitSuffixForSchema,
 			);
 
 			let statement = `export const ${withCasing(paramName, casing)} = ${func}("${it.name}", {${columns}})`;
@@ -844,12 +846,13 @@ const column = (
 	casing: Casing,
 	defaultValue?: any,
 	internals?: PgKitInternals,
+	omitSuffixForSchema?: string,
 ) => {
 	const isExpression = internals?.tables[tableName]?.columns[name]?.isDefaultAnExpression ?? false;
 	const lowered = type.toLowerCase().replace('[]', '');
 
 	if (enumTypes.has(`${typeSchema}.${type.replace('[]', '')}`)) {
-		let out = `${withCasing(name, casing)}: ${withCasing(paramNameFor(type.replace('[]', ''), typeSchema), casing)}(${
+		let out = `${withCasing(name, casing)}: ${withCasing(paramNameFor(type.replace('[]', ''), typeSchema, omitSuffixForSchema), casing)}(${
 			dbColumnName({ name, casing })
 		})`;
 		return out;
@@ -1117,6 +1120,7 @@ const createTableColumns = (
 	schemas: Record<string, string>,
 	casing: Casing,
 	internals: PgKitInternals,
+	omitSuffixForSchema?: string,
 ): string => {
 	let statement = '';
 
@@ -1144,6 +1148,7 @@ const createTableColumns = (
 			casing,
 			it.default,
 			internals,
+			omitSuffixForSchema,
 		);
 		statement += '\t';
 		statement += columnStatement;
@@ -1172,7 +1177,7 @@ const createTableColumns = (
 
 		// 			const paramsStr = objToStatement2(params);
 		// 			const tableSchema = schemas[it.schemaTo || ''];
-		// 			const paramName = paramNameFor(it.tableTo, tableSchema);
+		// 			const paramName = paramNameFor(it.tableTo, tableSchema, omitSuffixForSchema);
 		// 			if (paramsStr) {
 		// 				return `.references(()${typeSuffix} => ${
 		// 					withCasing(
@@ -1340,12 +1345,17 @@ const createTableChecks = (
 	return statement;
 };
 
-const createTableFKs = (fks: ForeignKey[], schemas: Record<string, string>, casing: Casing): string => {
+const createTableFKs = (
+	fks: ForeignKey[],
+	schemas: Record<string, string>,
+	casing: Casing,
+	omitSuffixForSchema?: string
+): string => {
 	let statement = '';
 
 	fks.forEach((it) => {
 		const tableSchema = schemas[it.schemaTo || ''];
-		const paramName = paramNameFor(it.tableTo, tableSchema);
+		const paramName = paramNameFor(it.tableTo, tableSchema, omitSuffixForSchema);
 
 		const isSelf = it.tableTo === it.tableFrom;
 		const tableTo = isSelf ? 'table' : `${withCasing(paramName, casing)}`;