feat: sqlite

Author: jog1t

packages/core/src/client/worker-common.ts

@@ -1,11 +1,7 @@
-import type { AnyWorkerDefinition, WorkerDefinition } from "@/worker/definition";
-import type * as protoHttpResolve from "@/worker/protocol/http/resolve";
-import type { Encoding } from "@/worker/protocol/serde";
-import type { WorkerQuery } from "@/manager/protocol/query";
-import { logger } from "./log";
-import * as errors from "./errors";
-import { sendHttpRequest } from "./utils";
-import { HEADER_WORKER_QUERY, HEADER_ENCODING } from "@/worker/router-endpoints";
+import type {
+	AnyWorkerDefinition,
+	WorkerDefinition,
+} from "@/worker/definition";
 
 /**
  * Action function returned by Worker connections and handles.
@@ -27,11 +23,10 @@ export type WorkerActionFunction<
  * Maps action methods from worker definition to typed function signatures.
  */
 export type WorkerDefinitionActions<AD extends AnyWorkerDefinition> =
-	AD extends WorkerDefinition<any, any, any, any, infer R>
+	AD extends WorkerDefinition<any, any, any, any, infer R, any>
 		? {
 				[K in keyof R]: R[K] extends (...args: infer Args) => infer Return
 					? WorkerActionFunction<Args, Return>
 					: never;
 			}
 		: never;
-

packages/core/src/registry/config.ts

@@ -43,7 +43,7 @@ export type WorkerPeerConfig = z.infer<typeof WorkerPeerConfigSchema>;
 
 export const WorkersSchema = z.record(
 	z.string(),
-	z.custom<WorkerDefinition<any, any, any, any, any>>(),
+	z.custom<WorkerDefinition<any, any, any, any, any, any>>(),
 );
 export type Workers = z.infer<typeof WorkersSchema>;
 
@@ -75,10 +75,11 @@ export const RegistryConfigSchema = z.object({
 	inspector: InspectorConfigSchema.optional().default({ enabled: false }),
 
 	// TODO: Find a better way of passing around the test config
-	/** 
+	/**
 	 * Test configuration.
 	 *
 	 * DO NOT MANUALLY ENABLE. THIS IS USED INTERNALLY.
+	 * @internal
 	 **/
 	test: TestConfigSchema.optional().default({ enabled: false }),
 });

packages/core/src/worker/action.ts

@@ -13,8 +13,8 @@ import { WorkerContext } from "./context";
  *
  * @typeParam A Worker this action belongs to
  */
-export class ActionContext<S, CP, CS, V> {
-	#workerContext: WorkerContext<S, CP, CS, V>;
+export class ActionContext<S, CP, CS, V, DB> {
+	#workerContext: WorkerContext<S, CP, CS, V, DB>;
 
 	/**
 	 * Should not be called directly.
@@ -23,8 +23,8 @@ export class ActionContext<S, CP, CS, V> {
 	 * @param conn - The connection associated with the action
 	 */
 	constructor(
-		workerContext: WorkerContext<S, CP, CS, V>,
-		public readonly conn: Conn<S, CP, CS, V>,
+		workerContext: WorkerContext<S, CP, CS, V, DB>,
+		public readonly conn: Conn<S, CP, CS, V, DB>,
 	) {
 		this.#workerContext = workerContext;
 	}
@@ -95,10 +95,17 @@ export class ActionContext<S, CP, CS, V> {
 	/**
 	 * Gets the map of connections.
 	 */
-	get conns(): Map<ConnId, Conn<S, CP, CS, V>> {
+	get conns(): Map<ConnId, Conn<S, CP, CS, V, DB>> {
 		return this.#workerContext.conns;
 	}
 
+	/**
+	 * @experimental
+	 */
+	get db(): DB {
+		return this.#workerContext.db;
+	}
+
 	/**
 	 * Forces the state to get saved.
 	 */

packages/core/src/worker/config.ts

@@ -23,6 +23,7 @@ export const WorkerConfigSchema = z
 		connState: z.any().optional(),
 		createConnState: z.function().optional(),
 		vars: z.any().optional(),
+		db: z.any().optional(),
 		createVars: z.function().optional(),
 		options: z
 			.object({
@@ -95,11 +96,11 @@ export interface OnConnectOptions<CP> {
 // Creates state config
 //
 // This must have only one or the other or else S will not be able to be inferred
-type CreateState<S, CP, CS, V> =
+type CreateState<S, CP, CS, V, DB> =
 	| { state: S }
 	| {
 			createState: (
-				c: WorkerContext<undefined, undefined, undefined, undefined>,
+				c: WorkerContext<undefined, undefined, undefined, undefined, undefined>,
 				opts: CreateStateOptions,
 			) => S | Promise<S>;
 	  }
@@ -108,11 +109,11 @@ type CreateState<S, CP, CS, V> =
 // Creates connection state config
 //
 // This must have only one or the other or else S will not be able to be inferred
-type CreateConnState<S, CP, CS, V> =
+type CreateConnState<S, CP, CS, V, DB> =
 	| { connState: CS }
 	| {
 			createConnState: (
-				c: WorkerContext<undefined, undefined, undefined, undefined>,
+				c: WorkerContext<undefined, undefined, undefined, undefined, undefined>,
 				opts: OnConnectOptions<CP>,
 			) => CS | Promise<CS>;
 	  }
@@ -124,7 +125,7 @@ type CreateConnState<S, CP, CS, V> =
 /**
  * @experimental
  */
-type CreateVars<S, CP, CS, V> =
+type CreateVars<S, CP, CS, V, DB> =
 	| {
 			/**
 			 * @experimental
@@ -136,30 +137,47 @@ type CreateVars<S, CP, CS, V> =
 			 * @experimental
 			 */
 			createVars: (
-				c: WorkerContext<undefined, undefined, undefined, undefined>,
+				c: WorkerContext<undefined, undefined, undefined, undefined, undefined>,
 				driverCtx: unknown,
 			) => V | Promise<V>;
 	  }
 	| Record<never, never>;
 
-export interface Actions<S, CP, CS, V> {
-	[Action: string]: (c: ActionContext<S, CP, CS, V>, ...args: any[]) => any;
+/**
+ * @experimental
+ */
+type DatabaseConfig<DB> = {
+	/**
+	 * @experimental
+	 */
+	db?: DB;
+};
+
+export interface Actions<S, CP, CS, V, DB> {
+	[Action: string]: (c: ActionContext<S, CP, CS, V, DB>, ...args: any[]) => any;
 }
 
 //export type WorkerConfig<S, CP, CS, V> = BaseWorkerConfig<S, CP, CS, V> &
 //	WorkerConfigLifecycle<S, CP, CS, V> &
 //	CreateState<S, CP, CS, V> &
 //	CreateConnState<S, CP, CS, V>;
 
-interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
+interface BaseWorkerConfig<
+	S,
+	CP,
+	CS,
+	V,
+	DB,
+	R extends Actions<S, CP, CS, V, DB>,
+> {
 	/**
 	 * Called when the worker is first initialized.
 	 *
 	 * Use this hook to initialize your worker's state.
 	 * This is called before any other lifecycle hooks.
 	 */
 	onCreate?: (
-		c: WorkerContext<S, CP, CS, V>,
+		c: WorkerContext<S, CP, CS, V, DB>,
 		opts: OnCreateOptions,
 	) => void | Promise<void>;
 
@@ -171,7 +189,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 	 *
 	 * @returns Void or a Promise that resolves when startup is complete
 	 */
-	onStart?: (c: WorkerContext<S, CP, CS, V>) => void | Promise<void>;
+	onStart?: (c: WorkerContext<S, CP, CS, V, DB>) => void | Promise<void>;
 
 	/**
 	 * Called when the worker's state changes.
@@ -181,7 +199,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 	 *
 	 * @param newState The updated state
 	 */
-	onStateChange?: (c: WorkerContext<S, CP, CS, V>, newState: S) => void;
+	onStateChange?: (c: WorkerContext<S, CP, CS, V, DB>, newState: S) => void;
 
 	/**
 	 * Called before a client connects to the worker.
@@ -194,7 +212,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 	 * @throws Throw an error to reject the connection
 	 */
 	onBeforeConnect?: (
-		c: WorkerContext<S, CP, CS, V>,
+		c: WorkerContext<S, CP, CS, V, DB>,
 		opts: OnConnectOptions<CP>,
 	) => void | Promise<void>;
 
@@ -208,8 +226,8 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 	 * @returns Void or a Promise that resolves when connection handling is complete
 	 */
 	onConnect?: (
-		c: WorkerContext<S, CP, CS, V>,
-		conn: Conn<S, CP, CS, V>,
+		c: WorkerContext<S, CP, CS, V, DB>,
+		conn: Conn<S, CP, CS, V, DB>,
 	) => void | Promise<void>;
 
 	/**
@@ -222,8 +240,8 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 	 * @returns Void or a Promise that resolves when disconnect handling is complete
 	 */
 	onDisconnect?: (
-		c: WorkerContext<S, CP, CS, V>,
-		conn: Conn<S, CP, CS, V>,
+		c: WorkerContext<S, CP, CS, V, DB>,
+		conn: Conn<S, CP, CS, V, DB>,
 	) => void | Promise<void>;
 
 	/**
@@ -239,7 +257,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 	 * @returns The modified output to send to the client
 	 */
 	onBeforeActionResponse?: <Out>(
-		c: WorkerContext<S, CP, CS, V>,
+		c: WorkerContext<S, CP, CS, V, DB>,
 		name: string,
 		args: unknown[],
 		output: Out,
@@ -251,7 +269,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
 // 1. Infer schema
 // 2. Omit keys that we'll manually define (because of generics)
 // 3. Define our own types that have generic constraints
-export type WorkerConfig<S, CP, CS, V> = Omit<
+export type WorkerConfig<S, CP, CS, V, DB> = Omit<
 	z.infer<typeof WorkerConfigSchema>,
 	| "actions"
 	| "onCreate"
@@ -268,18 +286,20 @@ export type WorkerConfig<S, CP, CS, V> = Omit<
 	| "vars"
 	| "createVars"
 > &
-	BaseWorkerConfig<S, CP, CS, V, Actions<S, CP, CS, V>> &
-	CreateState<S, CP, CS, V> &
-	CreateConnState<S, CP, CS, V> &
-	CreateVars<S, CP, CS, V>;
+	BaseWorkerConfig<S, CP, CS, V, DB, Actions<S, CP, CS, V, DB>> &
+	CreateState<S, CP, CS, V, DB> &
+	CreateConnState<S, CP, CS, V, DB> &
+	CreateVars<S, CP, CS, V, DB> &
+	DatabaseConfig<DB>;
 
 // See description on `WorkerConfig`
 export type WorkerConfigInput<
 	S,
 	CP,
 	CS,
 	V,
-	R extends Actions<S, CP, CS, V>,
+	DB,
+	R extends Actions<S, CP, CS, V, DB>,
 > = Omit<
 	z.input<typeof WorkerConfigSchema>,
 	| "actions"
@@ -297,16 +317,23 @@ export type WorkerConfigInput<
 	| "vars"
 	| "createVars"
 > &
-	BaseWorkerConfig<S, CP, CS, V, R> &
-	CreateState<S, CP, CS, V> &
-	CreateConnState<S, CP, CS, V> &
-	CreateVars<S, CP, CS, V>;
+	BaseWorkerConfig<S, CP, CS, V, DB, R> &
+	CreateState<S, CP, CS, V, DB> &
+	CreateConnState<S, CP, CS, V, DB> &
+	CreateVars<S, CP, CS, V, DB> &
+	DatabaseConfig<DB>;
 
 // For testing type definitions:
-export function test<S, CP, CS, V, R extends Actions<S, CP, CS, V>>(
-	input: WorkerConfigInput<S, CP, CS, V, R>,
-): WorkerConfig<S, CP, CS, V> {
-	const config = WorkerConfigSchema.parse(input) as WorkerConfig<S, CP, CS, V>;
+export function test<S, CP, CS, V, DB, R extends Actions<S, CP, CS, V, DB>>(
+	input: WorkerConfigInput<S, CP, CS, V, DB, R>,
+): WorkerConfig<S, CP, CS, V, DB> {
+	const config = WorkerConfigSchema.parse(input) as WorkerConfig<
+		S,
+		CP,
+		CS,
+		V,
+		DB
+	>;
 	return config;
 }
 

packages/core/src/worker/connection.ts

@@ -17,7 +17,7 @@ export function generateConnToken(): string {
 
 export type ConnId = string;
 
-export type AnyConn = Conn<any, any, any, any>;
+export type AnyConn = Conn<any, any, any, any, any>;
 
 /**
  * Represents a client connection to a worker.
@@ -26,13 +26,13 @@ export type AnyConn = Conn<any, any, any, any>;
  *
  * @see {@link https://rivet.gg/docs/connections|Connection Documentation}
  */
-export class Conn<S, CP, CS, V> {
+export class Conn<S, CP, CS, V, DB> {
 	subscriptions: Set<string> = new Set<string>();
 
 	#stateEnabled: boolean;
 
 	// TODO: Remove this cyclical reference
-	#worker: WorkerInstance<S, CP, CS, V>;
+	#worker: WorkerInstance<S, CP, CS, V, DB>;
 
 	/**
 	 * The proxied state that notifies of changes automatically.
@@ -99,7 +99,7 @@ export class Conn<S, CP, CS, V> {
 	 * @protected
 	 */
 	public constructor(
-		worker: WorkerInstance<S, CP, CS, V>,
+		worker: WorkerInstance<S, CP, CS, V, DB>,
 		persist: PersistedConn<CP, CS>,
 		driver: ConnDriver,
 		stateEnabled: boolean,
@@ -153,6 +153,11 @@ export class Conn<S, CP, CS, V> {
 	 * @param reason - The reason for disconnection.
 	 */
 	public async disconnect(reason?: string) {
-		await this.#driver.disconnect(this.#worker, this, this.__persist.ds, reason);
+		await this.#driver.disconnect(
+			this.#worker,
+			this,
+			this.__persist.ds,
+			reason,
+		);
 	}
 }

packages/core/src/worker/context.ts

@@ -8,10 +8,10 @@ import { Schedule } from "./schedule";
 /**
  * WorkerContext class that provides access to worker methods and state
  */
-export class WorkerContext<S, CP, CS, V> {
-	#worker: WorkerInstance<S, CP, CS, V>;
+export class WorkerContext<S, CP, CS, V, DB> {
+	#worker: WorkerInstance<S, CP, CS, V, DB>;
 
-	constructor(worker: WorkerInstance<S, CP, CS, V>) {
+	constructor(worker: WorkerInstance<S, CP, CS, V, DB>) {
 		this.#worker = worker;
 	}
 
@@ -84,10 +84,17 @@ export class WorkerContext<S, CP, CS, V> {
 	/**
 	 * Gets the map of connections.
 	 */
-	get conns(): Map<ConnId, Conn<S, CP, CS, V>> {
+	get conns(): Map<ConnId, Conn<S, CP, CS, V, DB>> {
 		return this.#worker.conns;
 	}
 
+	/**
+	 * Gets the database.
+	 */
+	get db(): DB {
+		return this.#worker.db;
+	}
+
 	/**
 	 * Forces the state to get saved.
 	 *