lobe-chat

1
import Dexie, { BulkError } from 'dexie';
2
import { ZodObject } from 'zod';
3

4
import { nanoid } from '@/utils/uuid';
5

6
import { BrowserDB, BrowserDBSchema, browserDB } from './db';
7
import { dataSync } from './sync';
8
import { DBBaseFieldsSchema } from './types/db';
9

10
export class BaseModel<N extends keyof BrowserDBSchema = any, T = BrowserDBSchema[N]['table']> {
11
  protected readonly db: BrowserDB;
12
  private readonly schema: ZodObject<any>;
13
  private readonly _tableName: keyof BrowserDBSchema;
14

15
  constructor(table: N, schema: ZodObject<any>, db = browserDB) {
16
    this.db = db;
17
    this.schema = schema;
18
    this._tableName = table;
19
  }
20

21
  get table() {
22
    return this.db[this._tableName] as Dexie.Table;
23
  }
24

25
  get yMap() {
26
    return dataSync.getYMap(this._tableName);
27
  }
28

29
  // **************** Create *************** //
30

31
  /**
32
   * create a new record
33
   */
34
  protected async _addWithSync<T = BrowserDBSchema[N]['model']>(
35
    data: T,
36
    id: string | number = nanoid(),
37
    primaryKey: string = 'id',
38
  ) {
39
    const result = this.schema.safeParse(data);
40

41
    if (!result.success) {
42
      const errorMsg = `[${this.db.name}][${this._tableName}] Failed to create new record. Error: ${result.error}`;
43

44
      const newError = new TypeError(errorMsg);
45
      // make this error show on console to help debug
46
      console.error(newError);
47
      throw newError;
48
    }
49

50
    const tableName = this._tableName;
51

52
    const record: any = {
53
      ...result.data,
54
      createdAt: Date.now(),
55
      [primaryKey]: id,
56
      updatedAt: Date.now(),
57
    };
58

59
    const newId = await this.db[tableName].add(record);
60

61
    // sync data to yjs data map
62
    this.updateYMapItem(newId);
63

64
    return { id: newId };
65
  }
66

67
  /**
68
   * Batch create new records
69
   * @param dataArray An array of data to be added
70
   * @param options
71
   * @param options.generateId
72
   * @param options.createWithNewId
73
   */
74
  protected async _batchAdd<T = BrowserDBSchema[N]['model']>(
75
    dataArray: T[],
76
    options: {
77
      /**
78
       * always create with a new id
79
       */
80
      createWithNewId?: boolean;
81
      idGenerator?: () => string;
82
      withSync?: boolean;
83
    } = {},
84
  ): Promise<{
85
    added: number;
86
    errors?: Error[];
87
    ids: string[];
88
    skips: string[];
89
    success: boolean;
90
  }> {
91
    const { idGenerator = nanoid, createWithNewId = false, withSync = true } = options;
92
    const validatedData: any[] = [];
93
    const errors = [];
94
    const skips: string[] = [];
95

96
    for (const data of dataArray) {
97
      const schemaWithId = this.schema.merge(DBBaseFieldsSchema.partial());
98

99
      const result = schemaWithId.safeParse(data);
100

101
      if (result.success) {
102
        const item = result.data;
103
        const autoId = idGenerator();
104

105
        const id = createWithNewId ? autoId : item.id ?? autoId;
106

107
        // skip if the id already exists
108
        if (await this.table.get(id)) {
109
          skips.push(id as string);
110
          continue;
111
        }
112

113
        const getTime = (time?: string | number) => {
114
          if (!time) return Date.now();
115
          if (typeof time === 'number') return time;
116

117
          return new Date(time).valueOf();
118
        };
119

120
        validatedData.push({
121
          ...item,
122
          createdAt: getTime(item.createdAt as string),
123
          id,
124
          updatedAt: getTime(item.updatedAt as string),
125
        });
126
      } else {
127
        errors.push(result.error);
128

129
        const errorMsg = `[${this.db.name}][${
130
          this._tableName
131
        }] Failed to create the record. Data: ${JSON.stringify(data)}. Errors: ${result.error}`;
132
        console.error(new TypeError(errorMsg));
133
      }
134
    }
135
    if (validatedData.length === 0) {
136
      // No valid data to add
137
      return { added: 0, errors, ids: [], skips, success: false };
138
    }
139

140
    // Using bulkAdd to insert validated data
141
    try {
142
      await this.table.bulkAdd(validatedData);
143

144
      if (withSync) {
145
        dataSync.transact(() => {
146
          const pools = validatedData.map(async (item) => {
147
            await this.updateYMapItem(item.id);
148
          });
149
          Promise.all(pools);
150
        });
151
      }
152

153
      return {
154
        added: validatedData.length,
155
        ids: validatedData.map((item) => item.id),
156
        skips,
157
        success: true,
158
      };
159
    } catch (error) {
160
      const bulkError = error as BulkError;
161
      // Handle bulkAdd errors here
162
      console.error(`[${this.db.name}][${this._tableName}] Bulk add error:`, bulkError);
163
      // Return the number of successfully added records and errors
164
      return {
165
        added: validatedData.length - skips.length - bulkError.failures.length,
166
        errors: bulkError.failures,
167
        ids: validatedData.map((item) => item.id),
168
        skips,
169
        success: false,
170
      };
171
    }
172
  }
173

174
  // **************** Delete *************** //
175

176
  protected async _deleteWithSync(id: string) {
177
    const result = await this.table.delete(id);
178
    // sync delete data to yjs data map
179
    this.yMap?.delete(id);
180
    return result;
181
  }
182

183
  protected async _bulkDeleteWithSync(keys: string[]) {
184
    await this.table.bulkDelete(keys);
185
    // sync delete data to yjs data map
186

187
    dataSync.transact(() => {
188
      keys.forEach((id) => {
189
        this.yMap?.delete(id);
190
      });
191
    });
192
  }
193

194
  protected async _clearWithSync() {
195
    const result = await this.table.clear();
196
    // sync clear data to yjs data map
197
    this.yMap?.clear();
198
    return result;
199
  }
200

201
  // **************** Update *************** //
202

203
  protected async _updateWithSync(id: string, data: Partial<T>) {
204
    // we need to check whether the data is valid
205
    // pick data related schema from the full schema
206
    const keys = Object.keys(data);
207
    const partialSchema = this.schema.pick(Object.fromEntries(keys.map((key) => [key, true])));
208

209
    const result = partialSchema.safeParse(data);
210
    if (!result.success) {
211
      const errorMsg = `[${this.db.name}][${this._tableName}] Failed to update the record:${id}. Error: ${result.error}`;
212

213
      const newError = new TypeError(errorMsg);
214
      // make this error show on console to help debug
215
      console.error(newError);
216
      throw newError;
217
    }
218

219
    const success = await this.table.update(id, { ...data, updatedAt: Date.now() });
220

221
    // sync data to yjs data map
222
    this.updateYMapItem(id);
223

224
    return { success };
225
  }
226

227
  protected async _putWithSync(data: any, id: string) {
228
    const result = await this.table.put(data, id);
229

230
    // sync data to yjs data map
231
    this.updateYMapItem(id);
232

233
    return result;
234
  }
235

236
  protected async _bulkPutWithSync(items: T[]) {
237
    await this.table.bulkPut(items);
238

239
    await dataSync.transact(() => {
240
      items.forEach((items) => {
241
        this.updateYMapItem((items as any).id);
242
      });
243
    });
244
  }
245

246
  // **************** Helper *************** //
247

248
  private updateYMapItem = async (id: string) => {
249
    const newData = await this.table.get(id);
250
    this.yMap?.set(id, newData);
251
  };
252
}
253