Skip to content

Commit

Permalink
[db] Load seed files with vite dev server (#10941)
Browse files Browse the repository at this point in the history 
* feat: load seed files with full vite dev server

* chore: remove unused export
  • Loading branch information
@bholmesdev
bholmesdev authored May 3, 2024
1 parent 1b41be4 commit 3b014fd
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 46 deletions.
82 changes: 79 additions & 3 deletions packages/db/src/core/integration/index.ts
View file
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
View file
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;
}
}

0 comments on commit 3b014fd

Please sign in to comment.