fix: catch up with feature updates

- improve code styling
- enable schema support for db pull

Author: svetch

packages/cli/package.json

@@ -31,6 +31,7 @@
     "dependencies": {
         "@dotenvx/dotenvx": "^1.51.0",
         "@zenstackhq/common-helpers": "workspace:*",
+        "@zenstackhq/schema": "workspace:*",
         "@zenstackhq/language": "workspace:*",
         "@zenstackhq/sdk": "workspace:*",
         "prisma": "catalog:",

packages/cli/src/actions/action-utils.ts

@@ -1,5 +1,5 @@
-import { loadDocument } from '@zenstackhq/language';
-import { isDataSource } from '@zenstackhq/language/ast';
+import { type ZModelServices, loadDocument } from '@zenstackhq/language';
+import { type Model, isDataSource } from '@zenstackhq/language/ast';
 import { PrismaSchemaGenerator } from '@zenstackhq/sdk';
 import colors from 'colors';
 import fs from 'node:fs';
@@ -41,22 +41,22 @@ export function getSchemaFile(file?: string) {
     }
 }
 
-export async function loadSchemaDocument(schemaFile: string) {
-    const loadResult = await loadDocument(schemaFile);
-    if (!loadResult.success) {
-        loadResult.errors.forEach((err) => {
-            console.error(colors.red(err));
-        });
-        throw new CliError('Schema contains errors. See above for details.');
-    }
-    loadResult.warnings.forEach((warn) => {
-        console.warn(colors.yellow(warn));
-    });
-    return loadResult.model;
-}
+export async function loadSchemaDocument(
+    schemaFile: string,
+    opts?: { keepImports?: boolean; returnServices?: false },
+): Promise<Model>;
+export async function loadSchemaDocument(
+    schemaFile: string,
+    opts: { returnServices: true; keepImports?: boolean },
+): Promise<{ model: Model; services: ZModelServices }>;
+export async function loadSchemaDocument(
+    schemaFile: string,
+    opts: { returnServices?: boolean; keepImports?: boolean } = {},
+) {
+    const returnServices = opts.returnServices || false;
+    const keepImports = opts.keepImports || false;
 
-export async function loadSchemaDocumentWithServices(schemaFile: string) {
-    const loadResult = await loadDocument(schemaFile, [], true);
+    const loadResult = await loadDocument(schemaFile, [], keepImports);
     if (!loadResult.success) {
         loadResult.errors.forEach((err) => {
             console.error(colors.red(err));
@@ -66,7 +66,10 @@ export async function loadSchemaDocumentWithServices(schemaFile: string) {
     loadResult.warnings.forEach((warn) => {
         console.warn(colors.yellow(warn));
     });
-    return { services: loadResult.services, model: loadResult.model };
+
+    if (returnServices) return { model: loadResult.model, services: loadResult.services };
+
+    return loadResult.model;
 }
 
 export function handleSubProcessError(err: unknown) {

packages/cli/src/actions/db.ts

@@ -1,13 +1,13 @@
-import { Model, Enum, DataModel, DataField } from '@zenstackhq/language/ast';
-import { ZModelCodeGenerator } from '@zenstackhq/sdk';
+import { config } from '@dotenvx/dotenvx';
+import { ZModelCodeGenerator } from '@zenstackhq/language';
+import { type DataField, DataModel, Enum, type Model } from '@zenstackhq/language/ast';
 import fs from 'node:fs';
 import path from 'node:path';
 import { execPrisma } from '../utils/exec-utils';
 import { generateTempPrismaSchema, getSchemaFile, handleSubProcessError, requireDataSourceUrl, loadSchemaDocumentWithServices } from './action-utils';
 import { syncEnums, syncRelation, syncTable, type Relation } from './pull';
 import { providers } from './pull/provider';
 import { getDatasource, getDbName, getRelationFkName } from './pull/utils';
-import { config } from '@dotenvx/dotenvx';
 
 type PushOptions = {
     schema?: string;
@@ -17,7 +17,6 @@ type PushOptions = {
 
 export type PullOptions = {
     schema?: string;
-    excludeSchemas?: string[];
     out?: string;
     modelCasing: 'pascal' | 'camel' | 'snake' | 'kebab' | 'none';
     fieldCasing: 'pascal' | 'camel' | 'snake' | 'kebab' | 'none';
@@ -74,7 +73,7 @@ async function runPush(options: PushOptions) {
 async function runPull(options: PullOptions) {
     try {
         const schemaFile = getSchemaFile(options.schema);
-        const { model, services } = await loadSchemaDocumentWithServices(schemaFile);
+        const { model, services } = await loadSchemaDocument(schemaFile, { returnServices: true });
         config();
         const SUPPORTED_PROVIDERS = ['sqlite', 'postgresql'];
         const datasource = getDatasource(model);
@@ -94,8 +93,8 @@ async function runPull(options: PullOptions) {
         }
 
         const { enums: allEnums, tables: allTables } = await provider.introspect(datasource.url);
-        const enums = allEnums.filter((e) => !options.excludeSchemas?.includes(e.schema_name));
-        const tables = allTables.filter((t) => !options.excludeSchemas?.includes(t.schema));
+        const enums = allEnums.filter((e) => datasource.schemas.includes(e.schema_name));
+        const tables = allTables.filter((t) => datasource.schemas.includes(t.schema));
 
         const newModel: Model = {
             $type: 'Model',
@@ -106,11 +105,18 @@ async function runPull(options: PullOptions) {
             imports: [],
         };
 
-        syncEnums({ dbEnums: enums, model: newModel, services, options });
+        syncEnums({ dbEnums: enums, model: newModel, services, options, defaultSchema: datasource.defaultSchema });
 
         const resolvedRelations: Relation[] = [];
         for (const table of tables) {
-            const relations = syncTable({ table, model: newModel, provider, services, options });
+            const relations = syncTable({
+                table,
+                model: newModel,
+                provider,
+                services,
+                options,
+                defaultSchema: datasource.defaultSchema,
+            });
             resolvedRelations.push(...relations);
         }
 

packages/cli/src/actions/pull/index.ts

@@ -1,7 +1,7 @@
 import type { ZModelServices } from '@zenstackhq/language';
 import {
-    Attribute,
     isEnum,
+    type Attribute,
     type DataField,
     type DataModel,
     type Enum,
@@ -23,11 +23,13 @@ export function syncEnums({
     model,
     options,
     services,
+    defaultSchema,
 }: {
     dbEnums: IntrospectedEnum[];
     model: Model;
     services: ZModelServices;
     options: PullOptions;
+    defaultSchema: string;
 }) {
     for (const dbEnum of dbEnums) {
         const { modified, name } = resolveNameCasing(options.modelCasing, dbEnum.enum_type);
@@ -55,18 +57,12 @@ export function syncEnums({
             });
         });
 
-        try {
-            if (dbEnum.schema_name && dbEnum.schema_name !== '' && dbEnum.schema_name !== 'public') {
-                factory.addAttribute((b) =>
-                    b
-                        .setDecl(getAttributeRef('@@schema', services))
-                        .addArg((a) => a.StringLiteral.setValue(dbEnum.schema_name)),
-                );
-            }
-        } catch (error: any) {
-            if (error?.message !== `Declaration not found: @@schema`) throw error;
-            //Waiting to support multi-schema
-            //TODO: remove catch after multi-schema support is implemented
+        if (dbEnum.schema_name && dbEnum.schema_name !== '' && dbEnum.schema_name !== defaultSchema) {
+            factory.addAttribute((b) =>
+                b
+                    .setDecl(getAttributeRef('@@schema', services))
+                    .addArg((a) => a.StringLiteral.setValue(dbEnum.schema_name)),
+            );
         }
 
         model.declarations.push(factory.get({ $container: model }));
@@ -143,12 +139,14 @@ export function syncTable({
     table,
     services,
     options,
+    defaultSchema,
 }: {
     table: IntrospectedTable;
     model: Model;
     provider: IntrospectionProvider;
     services: ZModelServices;
     options: PullOptions;
+    defaultSchema: string;
 }) {
     const idAttribute = getAttributeRef('@id', services);
     const modelIdAttribute = getAttributeRef('@@id', services);
@@ -372,16 +370,10 @@ export function syncTable({
         );
     });
 
-    try {
-        if (table.schema && table.schema !== '' && table.schema !== 'public') {
-            modelFactory.addAttribute((b) =>
-                b.setDecl(getAttributeRef('@@schema', services)).addArg((a) => a.StringLiteral.setValue(table.schema)),
-            );
-        }
-    } catch (error: any) {
-        if (error?.message !== `Declaration not found: @@schema`) throw error;
-        //Waiting to support multi-schema
-        //TODO: remove catch after multi-schema support is implemented
+    if (table.schema && table.schema !== '' && table.schema !== defaultSchema) {
+        modelFactory.addAttribute((b) =>
+            b.setDecl(getAttributeRef('@@schema', services)).addArg((a) => a.StringLiteral.setValue(table.schema)),
+        );
     }
 
     model.declarations.push(modelFactory.node);

packages/cli/src/actions/pull/provider/index.ts

@@ -1,9 +1,11 @@
+import type { DataSourceProviderType } from '@zenstackhq/schema';
 export * from './provider';
 
 import { postgresql } from './postgresql';
+import type { IntrospectionProvider } from './provider';
 import { sqlite } from './sqlite';
 
-export const providers = {
+export const providers: Record<DataSourceProviderType, IntrospectionProvider> = {
     postgresql,
     sqlite,
 };

packages/cli/src/actions/pull/utils.ts

@@ -11,8 +11,8 @@ import {
     type Model,
     type StringLiteral,
 } from '@zenstackhq/language/ast';
-import { getStringLiteral } from '@zenstackhq/language/utils';
-import type { DataSourceProviderType } from '@zenstackhq/sdk/schema';
+import { getLiteralArray, getStringLiteral } from '@zenstackhq/language/utils';
+import type { DataSourceProviderType } from '@zenstackhq/schema';
 import type { Reference } from 'langium';
 
 export function getAttribute(model: Model, attrName: string) {
@@ -53,12 +53,26 @@ export function getDatasource(model: Model) {
         throw new Error('The url field must be a string literal or an env().');
     }
 
+    const defaultSchemaField = datasource.fields.find((f) => f.name === 'defaultSchema');
+    const defaultSchema = (defaultSchemaField && getStringLiteral(defaultSchemaField.value)) || 'public';
+
+    const schemasField = datasource.fields.find((f) => f.name === 'schemas');
+    const schemas =
+        (schemasField &&
+            getLiteralArray(schemasField.value)
+                ?.map(getStringLiteral)
+                .filter((s) => s !== undefined)) ||
+        [];
+
     return {
         name: datasource.name,
         provider: getStringLiteral(
             datasource.fields.find((f) => f.name === 'provider')?.value,
         ) as DataSourceProviderType,
         url,
+        defaultSchema,
+        schemas,
+        allSchemas: [defaultSchema, ...schemas],
     };
 }
 

packages/cli/src/index.ts

@@ -142,7 +142,6 @@ function createProgram() {
         .description('Introspect your database.')
         .addOption(schemaOption)
         .addOption(noVersionCheckOption)
-        .addOption(new Option('-e, --exclude-schemas <schemas...>', 'exclude specific schemas from introspection'))
         .addOption(new Option('-o, --out <path>', 'add custom output path for the introspected schema'))
         .addOption(
             new Option('--model-casing <pascal|camel|snake|kebab>', 'set the casing of generated models').default(