Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[db] Load seed files with vite dev server #10941

Merged
merged 2 commits into from
May 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 79 additions & 3 deletions packages/db/src/core/integration/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,14 @@ import { fileURLToPath } from 'url';
import type { AstroConfig, AstroIntegration } from 'astro';
import { mkdir, writeFile } from 'fs/promises';
import { blue, yellow } from 'kleur/colors';
import { loadEnv } from 'vite';
import {
createServer,
loadEnv,
mergeConfig,
type HMRPayload,
type UserConfig,
type ViteDevServer,
} from 'vite';
import parseArgs from 'yargs-parser';
import { SEED_DEV_FILE_NAME } from '../queries.js';
import { AstroDbError } from '../../runtime/utils.js';
Expand All @@ -18,15 +25,20 @@ import {
type LateSeedFiles,
type LateTables,
vitePluginDb,
RESOLVED_VIRTUAL_MODULE_ID,
type SeedHandler,
resolved,
} from './vite-plugin-db.js';
import { vitePluginInjectEnvTs } from './vite-plugin-inject-env-ts.js';
import { LibsqlError } from '@libsql/client';
import { EXEC_DEFAULT_EXPORT_ERROR, EXEC_ERROR } from '../errors.js';

function astroDBIntegration(): AstroIntegration {
let connectToStudio = false;
let configFileDependencies: string[] = [];
let root: URL;
let appToken: ManagedAppToken | undefined;
// Used during production builds to load seed files.
let tempViteServer: ViteDevServer | undefined;

// Make table loading "late" to pass to plugins from `config:setup`,
// but load during `config:done` to wait for integrations to settle.
Expand All @@ -40,6 +52,13 @@ function astroDBIntegration(): AstroIntegration {
throw new Error('[astro:db] INTERNAL Seed files not loaded yet');
},
};
let seedHandler: SeedHandler = {
execute: () => {
throw new Error('[astro:db] INTERNAL Seed handler not loaded yet');
},
inProgress: false,
};

let command: 'dev' | 'build' | 'preview';
let output: AstroConfig['output'] = 'server';
return {
Expand All @@ -65,6 +84,7 @@ function astroDBIntegration(): AstroIntegration {
root: config.root,
srcDir: config.srcDir,
output: config.output,
seedHandler,
});
} else {
dbPlugin = vitePluginDb({
Expand All @@ -75,6 +95,7 @@ function astroDBIntegration(): AstroIntegration {
srcDir: config.srcDir,
output: config.output,
logger,
seedHandler,
});
}

Expand Down Expand Up @@ -104,6 +125,9 @@ function astroDBIntegration(): AstroIntegration {
await typegenInternal({ tables: tables.get() ?? {}, root: config.root });
},
'astro:server:setup': async ({ server, logger }) => {
seedHandler.execute = async (fileUrl) => {
await executeSeedFile({ fileUrl, viteServer: server });
};
const filesToWatch = [
...CONFIG_FILE_NAMES.map((c) => new URL(c, getDbDirectoryUrl(root))),
...configFileDependencies.map((c) => new URL(c, root)),
Expand All @@ -126,7 +150,7 @@ function astroDBIntegration(): AstroIntegration {
);
// Eager load astro:db module on startup
if (seedFiles.get().length || localSeedPaths.find((path) => existsSync(path))) {
server.ssrLoadModule(RESOLVED_VIRTUAL_MODULE_ID).catch((e) => {
server.ssrLoadModule(resolved.module).catch((e) => {
logger.error(e instanceof Error ? e.message : String(e));
});
}
Expand All @@ -146,8 +170,15 @@ function astroDBIntegration(): AstroIntegration {

logger.info('database: ' + (connectToStudio ? yellow('remote') : blue('local database.')));
},
'astro:build:setup': async ({ vite }) => {
tempViteServer = await getTempViteServer({ viteConfig: vite });
seedHandler.execute = async (fileUrl) => {
await executeSeedFile({ fileUrl, viteServer: tempViteServer! });
};
},
'astro:build:done': async ({}) => {
await appToken?.destroy();
await tempViteServer?.close();
},
},
};
Expand All @@ -161,3 +192,48 @@ function databaseFileEnvDefined() {
export function integration(): AstroIntegration[] {
return [astroDBIntegration(), fileURLIntegration()];
}

async function executeSeedFile({
fileUrl,
viteServer,
}: {
fileUrl: URL;
viteServer: ViteDevServer;
}) {
const mod = await viteServer.ssrLoadModule(fileUrl.pathname);
if (typeof mod.default !== 'function') {
throw new AstroDbError(EXEC_DEFAULT_EXPORT_ERROR(fileURLToPath(fileUrl)));
}
try {
await mod.default();
} catch (e) {
if (e instanceof LibsqlError) {
throw new AstroDbError(EXEC_ERROR(e.message));
}
throw e;
}
}

/**
* Inspired by Astro content collection config loader.
*/
async function getTempViteServer({ viteConfig }: { viteConfig: UserConfig }) {
const tempViteServer = await createServer(
mergeConfig(viteConfig, {
server: { middlewareMode: true, hmr: false, watch: null },
optimizeDeps: { noDiscovery: true },
ssr: { external: [] },
logLevel: 'silent',
})
);

const hotSend = tempViteServer.hot.send;
tempViteServer.hot.send = (payload: HMRPayload) => {
if (payload.type === 'error') {
throw payload.err;
}
return hotSend(payload);
};

return tempViteServer;
}
75 changes: 32 additions & 43 deletions packages/db/src/core/integration/vite-plugin-db.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,24 @@ import { createLocalDatabaseClient } from '../../runtime/db-client.js';
import { type SQL, sql } from 'drizzle-orm';
import { existsSync } from 'node:fs';
import { normalizeDatabaseUrl } from '../../runtime/index.js';
import { bundleFile, getResolvedFileUrl, importBundledFile } from '../load-file.js';
import { getResolvedFileUrl } from '../load-file.js';
import { SQLiteAsyncDialect } from 'drizzle-orm/sqlite-core';
import { EXEC_DEFAULT_EXPORT_ERROR, EXEC_ERROR } from '../errors.js';
import { LibsqlError } from '@libsql/client';
import { AstroDbError } from '../../runtime/utils.js';

export const RESOLVED_VIRTUAL_MODULE_ID = '\0' + VIRTUAL_MODULE_ID;
export const resolved = {
module: '\0' + VIRTUAL_MODULE_ID,
importedFromSeedFile: '\0' + VIRTUAL_MODULE_ID + ':seed',
};

export type LateTables = {
get: () => DBTables;
};
export type LateSeedFiles = {
get: () => Array<string | URL>;
};
export type SeedHandler = {
inProgress: boolean;
execute: (fileUrl: URL) => Promise<void>;
};

type VitePluginDBParams =
| {
Expand All @@ -32,6 +36,7 @@ type VitePluginDBParams =
root: URL;
logger?: AstroIntegrationLogger;
output: AstroConfig['output'];
seedHandler: SeedHandler;
}
| {
connectToStudio: true;
Expand All @@ -40,6 +45,7 @@ type VitePluginDBParams =
srcDir: URL;
root: URL;
output: AstroConfig['output'];
seedHandler: SeedHandler;
};

export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
Expand All @@ -51,10 +57,14 @@ export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
command = resolvedConfig.command;
},
async resolveId(id) {
if (id === VIRTUAL_MODULE_ID) return RESOLVED_VIRTUAL_MODULE_ID;
if (id !== VIRTUAL_MODULE_ID) return;
if (params.seedHandler.inProgress) {
return resolved.importedFromSeedFile;
}
return resolved.module;
},
async load(id) {
if (id !== RESOLVED_VIRTUAL_MODULE_ID) return;
if (id !== resolved.module && id !== resolved.importedFromSeedFile) return;

if (params.connectToStudio) {
return getStudioVirtualModContents({
Expand All @@ -64,24 +74,31 @@ export function vitePluginDb(params: VitePluginDBParams): VitePlugin {
output: params.output,
});
}

// When seeding, we resolved to a different virtual module.
// this prevents an infinite loop attempting to rerun seed files.
// Short circuit with the module contents in this case.
if (id === resolved.importedFromSeedFile) {
return getLocalVirtualModContents({
root: params.root,
tables: params.tables.get(),
});
}

await recreateTables(params);
const seedFiles = getResolvedSeedFiles(params);
let hasSeeded = false;
for await (const seedFile of seedFiles) {
// Use `addWatchFile()` to invalidate the `astro:db` module
// when a seed file changes.
this.addWatchFile(fileURLToPath(seedFile));
if (existsSync(seedFile)) {
hasSeeded = true;
await executeSeedFile({
tables: params.tables.get() ?? {},
fileUrl: seedFile,
root: params.root,
});
params.seedHandler.inProgress = true;
await params.seedHandler.execute(seedFile);
}
}
if (hasSeeded) {
if (params.seedHandler.inProgress) {
(params.logger ?? console).info('Seeded database.');
params.seedHandler.inProgress = false;
}
return getLocalVirtualModContents({
root: params.root,
Expand Down Expand Up @@ -197,31 +214,3 @@ function getResolvedSeedFiles({
const integrationSeedFiles = seedFiles.get().map((s) => getResolvedFileUrl(root, s));
return [...integrationSeedFiles, ...localSeedFiles];
}

async function executeSeedFile({
tables,
root,
fileUrl,
}: {
tables: DBTables;
root: URL;
fileUrl: URL;
}) {
const virtualModContents = getLocalVirtualModContents({
tables: tables ?? {},
root,
});
const { code } = await bundleFile({ virtualModContents, root, fileUrl });
const mod = await importBundledFile({ code, root });
if (typeof mod.default !== 'function') {
throw new AstroDbError(EXEC_DEFAULT_EXPORT_ERROR(fileURLToPath(fileUrl)));
}
try {
await mod.default();
} catch (e) {
if (e instanceof LibsqlError) {
throw new AstroDbError(EXEC_ERROR(e.message));
}
throw e;
}
}
Loading