drizzle-driver: fix rebased beta rewrite

Author: ziahamza

packages/drizzle-driver/src/sqlite/PowerSyncSQLitePreparedQuery.ts

@@ -15,23 +15,19 @@ import {
 
 type PreparedQueryConfig = Omit<PreparedQueryConfigBase, 'statement' | 'run'>;
 
-/**
- * Callback which uses a LockContext for database operations.
- */
 export type LockCallback<T> = (ctx: LockContext) => Promise<T>;
 
-/**
- * Provider for specific database contexts.
- * Handlers are provided a context to the provided callback.
- * This does not necessarily need to acquire a database lock for each call.
- * Calls might use the same lock context for multiple operations.
- * The read/write context may relate to a single read OR write context.
- */
 export type ContextProvider = {
   useReadContext: <T>(fn: LockCallback<T>) => Promise<T>;
   useWriteContext: <T>(fn: LockCallback<T>) => Promise<T>;
 };
 
+type ResultMapper = (rows: unknown[][], mapColumnValue?: (value: unknown) => unknown) => unknown;
+type RelationalResultMapper = (
+  rows: Record<string, unknown>[],
+  mapColumnValue?: (value: unknown) => unknown
+) => unknown;
+
 export class PowerSyncSQLitePreparedQuery<
   T extends PreparedQueryConfig = PreparedQueryConfig
 > extends SQLitePreparedQuery<{
@@ -46,25 +42,62 @@ export class PowerSyncSQLitePreparedQuery<
 
   private readOnly = false;
 
+  constructor(
+    contextProvider: ContextProvider,
+    query: Query,
+    logger: Logger,
+    fields: SelectedFieldsOrdered | undefined,
+    executeMethod: SQLiteExecuteMethod,
+    isResponseInArrayMode: boolean,
+    customResultMapper?: ResultMapper,
+    cache?: Cache | undefined,
+    queryMetadata?:
+      | {
+          type: 'select' | 'update' | 'delete' | 'insert';
+          tables: string[];
+        }
+      | undefined,
+    cacheConfig?: WithCacheConfig | undefined,
+    relationalQueryMode?: false
+  );
+  constructor(
+    contextProvider: ContextProvider,
+    query: Query,
+    logger: Logger,
+    fields: SelectedFieldsOrdered | undefined,
+    executeMethod: SQLiteExecuteMethod,
+    isResponseInArrayMode: boolean,
+    customResultMapper: RelationalResultMapper,
+    cache: Cache | undefined,
+    queryMetadata:
+      | {
+          type: 'select' | 'update' | 'delete' | 'insert';
+          tables: string[];
+        }
+      | undefined,
+    cacheConfig: WithCacheConfig | undefined,
+    relationalQueryMode: true
+  );
   constructor(
     private contextProvider: ContextProvider,
     query: Query,
     private logger: Logger,
     private fields: SelectedFieldsOrdered | undefined,
     executeMethod: SQLiteExecuteMethod,
     private _isResponseInArrayMode: boolean,
-    private customResultMapper?: (rows: unknown[][]) => unknown,
+    private customResultMapper?: ResultMapper | RelationalResultMapper,
     cache?: Cache | undefined,
     queryMetadata?:
       | {
           type: 'select' | 'update' | 'delete' | 'insert';
           tables: string[];
         }
       | undefined,
-    cacheConfig?: WithCacheConfig | undefined
+    cacheConfig?: WithCacheConfig | undefined,
+    private relationalQueryMode = false
   ) {
     super('async', executeMethod, query, cache, queryMetadata, cacheConfig);
-    this.readOnly = queryMetadata?.type == 'select';
+    this.readOnly = queryMetadata?.type == 'select' || relationalQueryMode;
   }
 
   async run(placeholderValues?: Record<string, unknown>): Promise<QueryResult> {
@@ -85,10 +118,20 @@ export class PowerSyncSQLitePreparedQuery<
       });
     }
 
+    if (customResultMapper && this.relationalQueryMode) {
+      const params = fillPlaceholders(query.params, placeholderValues ?? {});
+      logger.logQuery(query.sql, params);
+      const relationalResultMapper = customResultMapper as RelationalResultMapper;
+      return await this.useContext(async (ctx) => {
+        const rows = (await ctx.getAll(this.query.sql, params)) as Record<string, unknown>[];
+        return relationalResultMapper(rows) as T['all'];
+      });
+    }
+
     const rows = (await this.values(placeholderValues)) as unknown[][];
     if (customResultMapper) {
-      const mapped = customResultMapper(rows) as T['all'];
-      return mapped;
+      const resultMapper = customResultMapper as ResultMapper;
+      return resultMapper(rows) as T['all'];
     }
     return rows.map((row) => mapResultRow(fields!, row, (this as any).joinsNotNullableMap));
   }
@@ -105,6 +148,17 @@ export class PowerSyncSQLitePreparedQuery<
       });
     }
 
+    if (customResultMapper && this.relationalQueryMode) {
+      const relationalResultMapper = customResultMapper as RelationalResultMapper;
+      return this.useContext(async (ctx) => {
+        const row = (await ctx.get(this.query.sql, params)) as Record<string, unknown> | undefined;
+        if (!row) {
+          return undefined as T['get'];
+        }
+        return relationalResultMapper([row]) as T['get'];
+      });
+    }
+
     const rows = (await this.values(placeholderValues)) as unknown[][];
     const row = rows[0];
 
@@ -113,7 +167,8 @@ export class PowerSyncSQLitePreparedQuery<
     }
 
     if (customResultMapper) {
-      return customResultMapper(rows) as T['get'];
+      const resultMapper = customResultMapper as ResultMapper;
+      return resultMapper(rows) as T['get'];
     }
 
     return mapResultRow(fields!, row, joinsNotNullableMap);
@@ -141,17 +196,11 @@ export class PowerSyncSQLitePreparedQuery<
   }
 }
 
-/**
- * Maps a flat array of database row values to a result object based on the provided column definitions.
- * It reconstructs the hierarchical structure of the result by following the specified paths for each field.
- * It also handles nullification of nested objects when joined tables are nullable.
- */
 export function mapResultRow<TResult>(
   columns: SelectedFieldsOrdered,
   row: unknown[],
   joinsNotNullableMap: Record<string, boolean> | undefined
 ): TResult {
-  // Key -> nested object key, value -> table name if all fields in the nested object are from the same table, false otherwise
   const nullifyMap: Record<string, string | false> = {};
 
   const result = columns.reduce<Record<string, any>>((result, { path, field }, columnIndex) => {
@@ -178,10 +227,7 @@ export function mapResultRow<TResult>(
   return result as TResult;
 }
 
-/**
- * Determines the appropriate decoder for a given field.
- */
-function getDecoder(field: SQLiteColumn | SQL<unknown> | SQL.Aliased): DriverValueDecoder<unknown, unknown> {
+function getDecoder(field: any): DriverValueDecoder<unknown, unknown> {
   if (is(field, Column)) {
     return field;
   } else if (is(field, SQL)) {
@@ -204,15 +250,15 @@ function updateNullifyMap(
 
   const objectName = path[0]!;
   if (!(objectName in nullifyMap)) {
-    nullifyMap[objectName] = value === null ? getTableName(field.table) : false;
-  } else if (typeof nullifyMap[objectName] === 'string' && nullifyMap[objectName] !== getTableName(field.table)) {
+    nullifyMap[objectName] = value === null ? getTableName((field as any).table) : false;
+  } else if (
+    typeof nullifyMap[objectName] === 'string' &&
+    nullifyMap[objectName] !== getTableName((field as any).table)
+  ) {
     nullifyMap[objectName] = false;
   }
 }
 
-/**
- * Nullify all nested objects from nullifyMap that are nullable
- */
 function applyNullifyMap(
   result: Record<string, any>,
   nullifyMap: Record<string, string | false>,

packages/drizzle-driver/src/utils/schema.ts

@@ -23,13 +23,14 @@ import {
   type TableConfig
 } from 'drizzle-orm/sqlite-core';
 
-type PowerSyncColumnValue<T> = null extends T ? PowerSyncNonNullColumnValue<Exclude<T, null>> | null : PowerSyncNonNullColumnValue<T>;
+type PowerSyncColumnValue<T> = null extends T
+  ? PowerSyncNonNullColumnValue<Exclude<T, null>> | null
+  : PowerSyncNonNullColumnValue<T>;
 
 type PowerSyncNonNullColumnValue<T> = T extends number | string ? T : T extends boolean | Date ? number : string;
 
-type DrizzleTableColumns<T extends SQLiteTableWithColumns<any>> = T extends SQLiteTableWithColumns<infer TConfig>
-  ? TConfig['columns']
-  : never;
+type DrizzleTableColumns<T extends SQLiteTableWithColumns<any>> =
+  T extends SQLiteTableWithColumns<infer TConfig> ? TConfig['columns'] : never;
 
 type DrizzleColumnData<T> = T extends { _: { data: infer TData } } ? TData : never;
 
@@ -115,7 +116,6 @@ export type DrizzleTableWithPowerSyncOptions = {
   options?: DrizzleTablePowerSyncOptions;
 };
 
-<<<<<<< HEAD
 type DrizzleSchemaEntry = SQLiteTableWithColumns<any> | DrizzleTableWithPowerSyncOptions | Record<string, unknown>;
 
 export type TableName<T> = T extends { _: { name: infer TName extends string } }
@@ -125,33 +125,32 @@ export type TableName<T> = T extends { _: { name: infer TName extends string } }
     : never;
 
 export type TablesFromSchemaEntries<T> = {
-  [K in keyof T as T[K] extends Relations
-    ? never
-    : T[K] extends SQLiteTableWithColumns<any> | DrizzleTableWithPowerSyncOptions
-      ? TableName<T[K]>
-      : never]: T[K] extends SQLiteTableWithColumns<any>
+  [K in keyof T as T[K] extends SQLiteTableWithColumns<any> | DrizzleTableWithPowerSyncOptions
+    ? TableName<T[K]>
+    : never]: T[K] extends SQLiteTableWithColumns<any>
     ? Table<Expand<ExtractPowerSyncColumns<T[K]>>>
     : T[K] extends DrizzleTableWithPowerSyncOptions
       ? Table<Expand<ExtractPowerSyncColumns<T[K]['tableDefinition']>>>
       : never;
 };
 
-function toPowerSyncTables<
-  T extends Record<string, SQLiteTableWithColumns<any> | Relations | DrizzleTableWithPowerSyncOptions>
->(schemaEntries: T, options?: DrizzleAppSchemaOptions) {
+function toPowerSyncTables<T extends Record<string, DrizzleSchemaEntry>>(
+  schemaEntries: T,
+  options?: DrizzleAppSchemaOptions
+) {
   const casingCache = options?.casing ? new CasingCache(options?.casing) : undefined;
 
   const tables: Record<string, Table> = {};
   for (const schemaEntry of Object.values(schemaEntries)) {
-    let maybeTable: SQLiteTableWithColumns<any> | Relations | undefined = undefined;
+    let maybeTable: SQLiteTableWithColumns<any> | undefined = undefined;
     let maybeOptions: DrizzleTablePowerSyncOptions | undefined = undefined;
 
     if (typeof schemaEntry === 'object' && 'tableDefinition' in schemaEntry) {
       const tableWithOptions = schemaEntry as DrizzleTableWithPowerSyncOptions;
       maybeTable = tableWithOptions.tableDefinition;
       maybeOptions = tableWithOptions.options;
     } else {
-      maybeTable = schemaEntry;
+      maybeTable = schemaEntry as SQLiteTableWithColumns<any> | undefined;
     }
 
     if (isTable(maybeTable)) {
@@ -169,9 +168,7 @@ function toPowerSyncTables<
 export type DrizzleAppSchemaOptions = {
   casing?: Casing;
 };
-export class DrizzleAppSchema<
-  T extends Record<string, SQLiteTableWithColumns<any> | Relations | DrizzleTableWithPowerSyncOptions>
-> extends Schema {
+export class DrizzleAppSchema<T extends Record<string, DrizzleSchemaEntry>> extends Schema {
   constructor(drizzleSchema: T, options?: DrizzleAppSchemaOptions) {
     super(toPowerSyncTables(drizzleSchema, options));
     // This is just used for typing