Skip to content

Commit dafd14f

Browse files
AndriiShermanAndriiSherman
committed
cache fixes
1 parent bbf78ae commit dafd14f

3 files changed

Lines changed: 93 additions & 48 deletions

File tree

‎drizzle-orm/src/cache/upstash/cache.ts‎

Lines changed: 57 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -32,46 +32,76 @@ export class UpstashCache extends Cache {
3232
: undefined;
3333
}
3434

35-
override async get(key: string, _tables: string[], _isTag: boolean) {
36-
const res = await this.redis.get<any[]>(key) ?? undefined;
37-
return res;
35+
override async get(key: string, tables: string[], isTag: boolean = false): Promise<any[] | undefined> {
36+
const compositeKey = tables.sort().join(','); // Generate the composite key for the query
37+
38+
if (isTag) {
39+
// Handle cache lookup for tags
40+
const tagCompositeKey = await this.redis.hget<string>('tagsMap', key); // Retrieve composite key associated with the tag
41+
if (tagCompositeKey) {
42+
return await this.redis.hget(tagCompositeKey, key) as any[]; // Retrieve the cached result for the tag
43+
}
44+
return undefined;
45+
}
46+
47+
// Normal cache lookup for the composite key
48+
return await this.redis.hget(compositeKey, key) ?? undefined; // Retrieve result for normal query
3849
}
3950

40-
override async put(key: string, response: any, tables: string[], isTag: boolean, config?: CacheConfig) {
41-
await this.redis.set(key, response, config ? this.toInternalConfig(config) : this.internalConfig);
42-
for (const table of tables) {
43-
await this.redis.sadd(table, key);
51+
override async put(
52+
key: string,
53+
response: any,
54+
tables: string[],
55+
isTag: boolean = false,
56+
_config?: CacheConfig,
57+
): Promise<void> {
58+
if (isTag) {
59+
// When it's a tag, store the response in the composite key
60+
const compositeKey = tables.sort().join(',');
61+
await this.redis.hset(compositeKey, { [key]: response }); // Store the result with the tag under the composite key
62+
await this.redis.hset('tagsMap', { [key]: compositeKey }); // Store the tag and its composite key in the map
63+
for (const table of tables) {
64+
await this.redis.sadd(`prefix_${table}`, compositeKey);
65+
}
66+
} else {
67+
// Normal cache store
68+
const compositeKey = tables.sort().join(',');
69+
await this.redis.hset(compositeKey, { [key]: response }); // Store the result with the composite key
4470
}
4571
}
4672

4773
override async onMutate(params: MutationOption) {
48-
const tagsArray = params.tags ? Array.isArray(params.tags) ? params.tags : [params.tags] : [];
49-
const tablesArray = params.tables ? Array.isArray(params.tables) ? params.tables : [params.tables] : [];
74+
const tags = Array.isArray(params.tags) ? params.tags : params.tags ? [params.tags] : [];
75+
const tables = Array.isArray(params.tables) ? params.tables : params.tables ? [params.tables] : [];
76+
const mappedTables = tables.map((table) => is(table, Table) ? table[OriginalName] : table as string);
77+
const prefixedMappedTables = tables.map((table) => `prefix_${table}`);
5078

5179
const keysToDelete = new Set<string>();
5280

53-
for (const table of tablesArray) {
54-
const tableName = is(table, Table) ? table[OriginalName] : table as string;
55-
const keys = await this.redis.smembers(tableName);
56-
for (const key of keys) keysToDelete.add(key); // Add to the set
57-
}
58-
59-
if (keysToDelete.size > 0 || tagsArray.length > 0) {
60-
const pipeline = this.redis.pipeline();
61-
62-
for (const tag of tagsArray) {
63-
pipeline.del(tag);
81+
// Invalidate by table
82+
if (tables.length > 0) {
83+
// @ts-expect-error
84+
const compositeKeys: string[] = await this.redis.sunion(...prefixedMappedTables);
85+
for (const composite of compositeKeys) {
86+
const keys = await this.redis.hkeys(composite);
87+
for (const key of keys) keysToDelete.add(key);
88+
await this.redis.del(composite); // Remove composite entry
6489
}
90+
await this.redis.del(...prefixedMappedTables, ...mappedTables); // Remove table mappings
91+
}
6592

66-
for (const key of keysToDelete) {
67-
pipeline.del(key);
68-
for (const table of tablesArray) {
69-
const tableName = is(table, Table) ? table[OriginalName] : table as string;
70-
pipeline.srem(tableName, key);
71-
}
93+
// Invalidate by tag (WITHOUT invalidating the entire table)
94+
for (const tag of tags) {
95+
const compositeKey = await this.redis.hget<string>('tagsMap', tag);
96+
if (compositeKey) {
97+
await this.redis.hdel(compositeKey, tag); // Only remove the tag-related entry
98+
await this.redis.hdel('tagsMap', tag); // Remove tag reference
7299
}
100+
}
73101

74-
await pipeline.exec();
102+
// Delete affected cache entries
103+
if (keysToDelete.size > 0) {
104+
await this.redis.del(...keysToDelete);
75105
}
76106
}
77107
}

‎integration-tests/tests/pg/pglite.test.ts‎

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
import { PGlite } from '@electric-sql/pglite';
22
import { Name, sql } from 'drizzle-orm';
3+
import { upstashCache } from 'drizzle-orm/cache/upstash';
34
import { drizzle, type PgliteDatabase } from 'drizzle-orm/pglite';
45
import { migrate } from 'drizzle-orm/pglite/migrator';
56
import { afterAll, beforeAll, beforeEach, expect, test } from 'vitest';
67
import { skipTests } from '~/common';
7-
import { tests, usersMigratorTable, usersTable } from './pg-common';
8-
import { TestCache, TestGlobalCache, tests as cacheTests } from './pg-common-cache';
8+
import { usersMigratorTable, usersTable } from './pg-common';
9+
import { TestGlobalCache, tests as cacheTests } from './pg-common-cache';
910

1011
const ENABLE_LOGGING = false;
1112

@@ -17,8 +18,21 @@ let client: PGlite;
1718
beforeAll(async () => {
1819
client = new PGlite();
1920
db = drizzle(client, { logger: ENABLE_LOGGING });
20-
cachedDb = drizzle(client, { logger: ENABLE_LOGGING, cache: new TestCache() });
21-
dbGlobalCached = drizzle(client, { logger: ENABLE_LOGGING, cache: new TestGlobalCache() });
21+
cachedDb = drizzle(client, {
22+
logger: ENABLE_LOGGING,
23+
cache: upstashCache({
24+
url: 'https://healthy-deer-37505.upstash.io',
25+
token: 'AZKBAAIjcDFmYWYwMTA0YTVmNGE0NWZjODU2NWUzZmZkZTRhN2U0MnAxMA',
26+
}),
27+
});
28+
dbGlobalCached = drizzle(client, {
29+
logger: ENABLE_LOGGING,
30+
cache: upstashCache({
31+
url: 'https://healthy-deer-37505.upstash.io',
32+
token: 'AZKBAAIjcDFmYWYwMTA0YTVmNGE0NWZjODU2NWUzZmZkZTRhN2U0MnAxMA',
33+
global: true,
34+
}),
35+
});
2236
});
2337

2438
afterAll(async () => {
@@ -103,7 +117,7 @@ skipTests([
103117
'mySchema :: select with group by as column + sql',
104118
]);
105119

106-
tests();
120+
// tests();
107121
cacheTests();
108122

109123
beforeEach(async () => {

‎integration-tests/vitest.config.ts‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,22 +5,23 @@ import { defineConfig } from 'vitest/config';
55
export default defineConfig({
66
test: {
77
include: [
8-
'tests/seeder/**/*.test.ts',
9-
'tests/extensions/postgis/**/*',
10-
'tests/relational/**/*.test.ts',
11-
'tests/pg/**/*.test.ts',
12-
'tests/mysql/**/*.test.ts',
13-
'tests/singlestore/**/*.test.ts',
14-
'tests/sqlite/**/*.test.ts',
15-
'tests/replicas/**/*',
16-
'tests/imports/**/*',
17-
'tests/extensions/vectors/**/*',
18-
'tests/version.test.ts',
19-
'tests/pg/node-postgres.test.ts',
20-
'tests/utils/is-config.test.ts',
21-
'js-tests/driver-init/commonjs/*.test.cjs',
22-
'js-tests/driver-init/module/*.test.mjs',
23-
'tests/gel/**/*.test.ts',
8+
'tests/pg/pglite.test.ts',
9+
// 'tests/seeder/**/*.test.ts',
10+
// 'tests/extensions/postgis/**/*',
11+
// 'tests/relational/**/*.test.ts',
12+
// 'tests/pg/**/*.test.ts',
13+
// 'tests/mysql/**/*.test.ts',
14+
// 'tests/singlestore/**/*.test.ts',
15+
// 'tests/sqlite/**/*.test.ts',
16+
// 'tests/replicas/**/*',
17+
// 'tests/imports/**/*',
18+
// 'tests/extensions/vectors/**/*',
19+
// 'tests/version.test.ts',
20+
// 'tests/pg/node-postgres.test.ts',
21+
// 'tests/utils/is-config.test.ts',
22+
// 'js-tests/driver-init/commonjs/*.test.cjs',
23+
// 'js-tests/driver-init/module/*.test.mjs',
24+
// 'tests/gel/**/*.test.ts',
2425
],
2526
exclude: [
2627
...(process.env.SKIP_EXTERNAL_DB_TESTS

0 commit comments

Comments
 (0)