1
// Copyright (c) 2023 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
export class D1MockDO {
6
  constructor(state, env) {
7
    this.state = state;
8
    this.sql = this.state.storage.sql;
9
  }
10
 
11
  async fetch(request) {
12
    const { pathname, searchParams } = new URL(request.url);
13
    const is_query = pathname === '/query';
14
    const is_execute = pathname === '/execute';
15
    if (request.method === 'POST' && (is_query || is_execute)) {
16
      const body = await request.json();
17
      const resultsFormatParam = searchParams.get('resultsFormat');
18
      const resultsFormat =
19
        resultsFormatParam === 'ROWS_AND_COLUMNS'
20
          ? 'ROWS_AND_COLUMNS'
21
          : resultsFormatParam === 'NONE'
22
            ? 'NONE'
23
            : 'ARRAY_OF_OBJECTS';
24
      const safeRunQuery = (query) => {
25
        try {
26
          return this.runQuery(query, resultsFormat);
27
        } catch (e) {
28
          // Reproduce the production behavior by catching any error and returning a V4Failure
29
          return { success: false, error: String(e.message) };
30
        }
31
      };
32
      return Response.json(
33
        Array.isArray(body)
34
          ? body.map((query) => safeRunQuery(query))
35
          : safeRunQuery(body)
36
      );
37
    } else {
38
      return Response.json({ error: 'Not found' }, { status: 404 });
39
    }
40
  }
41
 
42
  runQuery(query, resultsFormat) {
43
    const { sql, params = [] } = query;
44
 
45
    const changes_stmt = this.sql.prepare(
46
      `SELECT total_changes() as changes, last_insert_rowid() as last_row_id`
47
    );
48
    const size_before = this.sql.databaseSize;
49
    const [[changes_before, last_row_id_before]] = Array.from(
50
      changes_stmt().raw()
51
    );
52
 
53
    const stmt = this.sql.prepare(sql)(...params);
54
    const columnNames = stmt.columnNames;
55
    const rawResults = Array.from(stmt.raw());
56
 
57
    // Convert to object-style results if necessary (for backwards compatibility)
58
    // .run() previously returned results. Folks relied on that, and we broke their running Workers, we shouldn't have done that.
59
    // To make the existing workers (those that didn't update to fix their .run to be .all) work again, we're hardcoding ResultFormat.NONE to do the same thing as what .all used to do (send back the array of results)
60
    const results =
61
      resultsFormat === 'NONE'
62
        ? rawResults.map((row) =>
63
            Object.fromEntries(columnNames.map((c, i) => [c, row[i]]))
64
          )
65
        : resultsFormat === 'ROWS_AND_COLUMNS'
66
          ? { columns: columnNames, rows: rawResults }
67
          : rawResults.map((row) =>
68
              Object.fromEntries(columnNames.map((c, i) => [c, row[i]]))
69
            );
70
 
71
    const [[changes_after, last_row_id_after]] = Array.from(
72
      changes_stmt().raw()
73
    );
74
 
75
    const size_after = this.sql.databaseSize;
76
    const num_changes = changes_after - changes_before;
77
    const has_changes = num_changes !== 0;
78
    const last_row_changed = last_row_id_after !== last_row_id_before;
79
 
80
    const db_size_different = size_after != size_before;
81
 
82
    // `changed_db` includes multiple ways the DB might be altered
83
    const changed_db = has_changes || last_row_changed || db_size_different;
84
 
85
    const { rowsRead: rows_read, rowsWritten: rows_written } = stmt;
86
 
87
    return {
88
      success: true,
89
      results,
90
      meta: {
91
        duration: Math.random() * 0.01,
92
        served_by: 'd1-mock',
93
        served_by_colo: 'DFW',
94
        changes: num_changes,
95
        last_row_id: last_row_id_after,
96
        changed_db,
97
        size_after,
98
        rows_read,
99
        rows_written,
100
      },
101
    };
102
  }
103
}
104
 
105
export default {
106
  commitTokenNum: 0,
107
  commitTokensReceived: [],
108
  commitTokensReturned: [],
109
  nextTokenExpected: null,
110
 
111
  async fetch(request, env, ctx) {
112
    if (request.url.startsWith('http://d1-api-test/commitTokens')) {
113
      return this.handleD1ApiTestRoutes(request);
114
    }
115
 
116
    // For our testing purposes, record any commit token passed through.
117
    const reqCommitToken = request.headers.get('x-cf-d1-session-commit-token');
118
    this.commitTokensReceived.push(reqCommitToken);
119
 
120
    try {
121
      const stub = env.db.get(env.db.idFromName('test'));
122
 
123
      // Add a commitToken to all responses.
124
      return stub
125
        .fetch(request)
126
        .then((resp) => this.buildResponseWithCommitToken(resp));
127
    } catch (err) {
128
      return Response.json(
129
        { error: err.message, stack: err.stack },
130
        { status: 500 }
131
      );
132
    }
133
  },
134
 
135
  async buildResponseWithCommitToken(resp) {
136
    let newToken = `token-${(++this.commitTokenNum).toLocaleString('en-US', {
137
      minimumIntegerDigits: 4,
138
      // no commas
139
      useGrouping: false,
140
    })}`;
141
    if (this.nextTokenExpected) {
142
      newToken = this.nextTokenExpected;
143
      this.nextTokenExpected = null;
144
    }
145
    this.commitTokensReturned.push(newToken);
146
    // Append an ever increasing commit token to the response.
147
    // Simulating the D1 eyeball worker.
148
    const newHeaders = new Headers(resp.headers);
149
    newHeaders.set('x-cf-d1-session-commit-token', newToken);
150
    return Response.json(await resp.json(), {
151
      status: resp.status,
152
      statusText: resp?.statusText,
153
      headers: newHeaders,
154
    });
155
  },
156
 
157
  async handleD1ApiTestRoutes(request) {
158
    const respondTokens = () =>
159
      Response.json(
160
        {
161
          commitTokensReceived: this.commitTokensReceived,
162
          commitTokensReturned: this.commitTokensReturned,
163
        },
164
        { status: 200 }
165
      );
166
 
167
    switch (new URL(request.url).pathname) {
168
      case '/commitTokens':
169
        // Special endpoints to accommodate our tests.
170
        return respondTokens();
171
 
172
      case '/commitTokens/nextToken':
173
        this.nextTokenExpected = new URL(request.url).searchParams.get('t');
174
        return respondTokens();
175
 
176
      case '/commitTokens/reset':
177
        this.commitTokensReceived = [];
178
        this.commitTokensReturned = [];
179
        this.commitTokenNum = 0;
180
        this.nextTokenExpected = null;
181
        return respondTokens();
182
 
183
      default:
184
        return Response.json({ error: 'invalid test route' }, { status: 404 });
185
    }
186
  },
187
};