@credo-ts/drizzle-storage 0.6.0-pr-2324-20250708122805 → 0.6.0-pr-2324-20250708123651

package/src/adapter/BaseDrizzleRecordAdapter.ts

@@ -0,0 +1,319 @@
+/**
+ * Adapter between a specific Record class and the record Type
+ */
+
+import { AgentContext, BaseRecord, CredoError, Query, QueryOptions, RecordNotFoundError } from '@credo-ts/core'
+import { Simplify, and, eq } from 'drizzle-orm'
+import { PgTable, pgTable } from 'drizzle-orm/pg-core'
+import { SQLiteTable as _SQLiteTable, sqliteTable } from 'drizzle-orm/sqlite-core'
+import { DrizzleDatabase, isDrizzlePostgresDatabase, isDrizzleSqliteDatabase } from '../DrizzleDatabase'
+import { CredoDrizzleStorageError } from '../error'
+import { getPostgresBaseRecordTable } from '../postgres'
+import { getSqliteBaseRecordTable } from '../sqlite'
+import { DrizzleCustomTagKeyMapping, queryToDrizzlePostgres } from './queryToDrizzlePostgres'
+import { queryToDrizzleSqlite } from './queryToDrizzleSqlite'
+
+// biome-ignore lint/suspicious/noExplicitAny: <explanation>
+export type AnyDrizzleAdapter = BaseDrizzleRecordAdapter<any, any, any, any, any>
+
+export type DrizzleAdapterValues<Table extends _SQLiteTable> = Simplify<
+  Omit<
+    { [Key in keyof Table['$inferInsert']]: Table['$inferInsert'][Key] },
+    Exclude<keyof ReturnType<typeof getSqliteBaseRecordTable>, 'customTags'>
+  >
+>
+
+export type DrizzleAdapterRecordValues<Table extends _SQLiteTable> = Simplify<
+  Omit<{ [Key in keyof Table['$inferInsert']]: Table['$inferInsert'][Key] }, 'contextCorrelationId'>
+>
+
+export abstract class BaseDrizzleRecordAdapter<
+  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+  CredoRecord extends BaseRecord<any, any, any>,
+  PostgresTable extends ReturnType<typeof pgTable<string, ReturnType<typeof getPostgresBaseRecordTable>>>,
+  PostgresSchema extends Record<string, unknown>,
+  SQLiteTable extends ReturnType<typeof sqliteTable<string, ReturnType<typeof getSqliteBaseRecordTable>>>,
+  SQLiteSchema extends Record<string, unknown>,
+> {
+  public recordType: CredoRecord['type']
+
+  public table: {
+    postgres: PostgresTable
+    sqlite: SQLiteTable
+  }
+
+  /**
+   * Allows overriding top level tags (as used by Credo record classes)
+   * to the database structure. For example mapping from the tag `presentationAuthSession`
+   * to the nested database json structure `presentation.authSession`.
+   */
+  public tagKeyMapping?: DrizzleCustomTagKeyMapping
+
+  public constructor(
+    public database: DrizzleDatabase<PostgresSchema, SQLiteSchema>,
+    table: {
+      postgres: PostgresTable
+      sqlite: SQLiteTable
+    },
+    recordType: CredoRecord['type']
+  ) {
+    this.table = table
+    this.recordType = recordType
+  }
+
+  public abstract getValues(record: CredoRecord): DrizzleAdapterValues<SQLiteTable>
+  public getValuesWithBase(agentContext: AgentContext, record: CredoRecord) {
+    return {
+      ...this.getValues(record),
+
+      // Always store based on context correlation id
+      contextCorrelationId: agentContext.contextCorrelationId,
+
+      id: record.id,
+      metadata: record.metadata.data,
+      createdAt: record.createdAt,
+      updatedAt: record.updatedAt,
+    }
+  }
+
+  private _toRecord(values: DrizzleAdapterRecordValues<SQLiteTable>): CredoRecord {
+    const filteredValues = Object.fromEntries(
+      Object.entries(values).filter(([_key, value]) => value !== null)
+    ) as DrizzleAdapterRecordValues<SQLiteTable>
+
+    return this.toRecord(filteredValues)
+  }
+
+  public abstract toRecord(values: DrizzleAdapterRecordValues<SQLiteTable>): CredoRecord
+
+  public async query(agentContext: AgentContext, query?: Query<CredoRecord>, queryOptions?: QueryOptions) {
+    if (isDrizzlePostgresDatabase(this.database)) {
+      let queryResult = this.database.select().from(this.table.postgres as PgTable)
+
+      if (query) {
+        queryResult = queryResult.where(
+          and(
+            // Always filter based on context correlation id
+            eq(this.table.postgres.contextCorrelationId, agentContext.contextCorrelationId),
+            queryToDrizzlePostgres(query ?? {}, this.table.postgres, this.tagKeyMapping)
+          )
+        ) as typeof queryResult
+      }
+
+      if (queryOptions?.limit !== undefined) {
+        queryResult = queryResult.limit(queryOptions.limit) as typeof queryResult
+      }
+
+      if (queryOptions?.offset !== undefined) {
+        queryResult = queryResult.offset(queryOptions.offset ?? 0) as typeof queryResult
+      }
+
+      const result = await queryResult
+      return result.map(({ contextCorrelationId, ...item }) =>
+        this._toRecord(item as DrizzleAdapterRecordValues<SQLiteTable>)
+      )
+    }
+
+    if (isDrizzleSqliteDatabase(this.database)) {
+      let queryResult = this.database.select().from(this.table.sqlite as SQLiteTable)
+
+      if (query) {
+        queryResult = queryResult.where(
+          and(
+            // Always filter based on context correlation id
+            eq(this.table.sqlite.contextCorrelationId, agentContext.contextCorrelationId),
+            queryToDrizzleSqlite(query ?? {}, this.table.sqlite, this.tagKeyMapping)
+          )
+        ) as unknown as typeof queryResult
+      }
+
+      if (queryOptions?.limit !== undefined) {
+        queryResult = queryResult.limit(queryOptions.limit) as unknown as typeof queryResult
+      }
+
+      if (queryOptions?.offset !== undefined) {
+        queryResult = queryResult.offset(queryOptions.offset ?? 0) as unknown as typeof queryResult
+      }
+
+      const result = await queryResult
+      return result.map(({ contextCorrelationId, ...item }) =>
+        this._toRecord(item as DrizzleAdapterRecordValues<SQLiteTable>)
+      )
+    }
+
+    throw new CredoError('Unsupported database type')
+  }
+
+  public async getById(agentContext: AgentContext, id: string) {
+    if (isDrizzlePostgresDatabase(this.database)) {
+      const [result] = await this.database
+        .select()
+        .from(this.table.postgres as PgTable)
+        .where(
+          and(
+            eq(this.table.postgres.id, id),
+            eq(this.table.postgres.contextCorrelationId, agentContext.contextCorrelationId)
+          )
+        )
+        .limit(1)
+
+      if (!result) {
+        throw new RecordNotFoundError(`record with id ${id} not found.`, {
+          recordType: this.recordType,
+        })
+      }
+
+      const { contextCorrelationId, ...item } = result
+      return this._toRecord(item as DrizzleAdapterRecordValues<SQLiteTable>)
+    }
+
+    if (isDrizzleSqliteDatabase(this.database)) {
+      const [result] = await this.database
+        .select()
+        .from(this.table.sqlite)
+        .where(
+          and(
+            eq(this.table.sqlite.id, id),
+            eq(this.table.sqlite.contextCorrelationId, agentContext.contextCorrelationId)
+          )
+        )
+        .limit(1)
+
+      if (!result) {
+        throw new RecordNotFoundError(`record with id ${id} not found.`, {
+          recordType: this.recordType,
+        })
+      }
+
+      const { contextCorrelationId, ...item } = result
+      return this._toRecord(item as DrizzleAdapterRecordValues<SQLiteTable>)
+    }
+
+    throw new CredoError('Unsupported database type')
+  }
+
+  public async insert(agentContext: AgentContext, record: CredoRecord) {
+    if (isDrizzlePostgresDatabase(this.database)) {
+      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+      await this.database.insert(this.table.postgres).values(this.getValuesWithBase(agentContext, record) as any)
+      return
+    }
+
+    if (isDrizzleSqliteDatabase(this.database)) {
+      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+      await this.database.insert(this.table.sqlite).values(this.getValuesWithBase(agentContext, record) as any)
+      return
+    }
+
+    throw new CredoError('Unsupported database type')
+  }
+
+  public async update(agentContext: AgentContext, record: CredoRecord) {
+    // Although id should always be set, if for some reason it is not set it can be quite impactful
+    if (!record.id) {
+      throw new CredoDrizzleStorageError(`Record of type ${record.type}' is missing 'id' column.`)
+    }
+
+    if (isDrizzlePostgresDatabase(this.database)) {
+      const updated = await this.database
+        .update(this.table.postgres)
+        // biome-ignore lint/suspicious/noExplicitAny: generics really don't play well here
+        .set(this.getValuesWithBase(agentContext, record) as any)
+        .where(
+          and(
+            eq(this.table.postgres.id, record.id),
+            eq(this.table.postgres.contextCorrelationId, agentContext.contextCorrelationId)
+          )
+        )
+        .returning({ id: this.table.postgres.id })
+
+      if (updated.length === 0) {
+        throw new RecordNotFoundError(`record with id ${record.id} not found.`, {
+          recordType: this.recordType,
+        })
+      }
+
+      return
+    }
+
+    if (isDrizzleSqliteDatabase(this.database)) {
+      const updated = await this.database
+        .update(this.table.sqlite)
+        // biome-ignore lint/suspicious/noExplicitAny: generics really don't play well here
+        .set(this.getValuesWithBase(agentContext, record) as any)
+        .where(
+          and(
+            eq(this.table.sqlite.id, record.id),
+            eq(this.table.sqlite.contextCorrelationId, agentContext.contextCorrelationId)
+          )
+        )
+        .returning({
+          id: this.table.sqlite.id,
+        })
+
+      if (updated.length === 0) {
+        throw new RecordNotFoundError(`record with id ${record.id} not found.`, {
+          recordType: this.recordType,
+        })
+      }
+
+      return
+    }
+
+    throw new CredoError('Unsupported database type')
+  }
+
+  public async delete(agentContext: AgentContext, id: string) {
+    // Although id should always be set, if for some reason it is not set it can be quite impactful
+    if (!id) {
+      throw new CredoDrizzleStorageError(`Missing required 'id' for delete.`)
+    }
+
+    if (isDrizzlePostgresDatabase(this.database)) {
+      const deleted = await this.database
+        .delete(this.table.postgres)
+        .where(
+          and(
+            eq(this.table.postgres.id, id),
+            eq(this.table.postgres.contextCorrelationId, agentContext.contextCorrelationId)
+          )
+        )
+        .returning({
+          id: this.table.postgres.id,
+        })
+
+      if (deleted.length === 0) {
+        throw new RecordNotFoundError(`record with id ${id} not found.`, {
+          recordType: this.recordType,
+        })
+      }
+
+      return
+    }
+
+    if (isDrizzleSqliteDatabase(this.database)) {
+      const deleted = await this.database
+        .delete(this.table.postgres)
+        .where(
+          and(
+            eq(this.table.sqlite.id, id),
+            eq(this.table.sqlite.contextCorrelationId, agentContext.contextCorrelationId)
+          )
+        )
+        .returning({
+          id: this.table.sqlite.id,
+        })
+
+      if (deleted.length === 0) {
+        throw new RecordNotFoundError(`record with id ${id} not found.`, {
+          recordType: this.recordType,
+        })
+      }
+
+      return
+    }
+
+    // @ts-expect-error
+    throw new CredoError(`Unsupported database type '${database.type}'`)
+  }
+}

package/src/adapter/__tests__/queryToDrizzlePostgres.test.ts

@@ -0,0 +1,328 @@
+import { ConnectionRecord, DidExchangeRole, DidExchangeState } from '@credo-ts/didcomm'
+import { pushSchema } from 'drizzle-kit/api'
+import { PgDatabase } from 'drizzle-orm/pg-core'
+import { drizzle } from 'drizzle-orm/pglite'
+import { DrizzlePostgresDatabase } from '../../DrizzleDatabase'
+import * as coreContextSchema from '../../core/context/postgres'
+import * as didcommConnectionSchema from '../../didcomm/connection/postgres'
+import { queryToDrizzlePostgres } from '../queryToDrizzlePostgres'
+
+const { didcommConnection } = didcommConnectionSchema
+const { context } = coreContextSchema
+
+const db = drizzle('memory://', {
+  schema: { ...didcommConnectionSchema, ...coreContextSchema },
+}) as unknown as DrizzlePostgresDatabase<typeof didcommConnectionSchema>
+
+describe('queryToDrizzlePostgres', () => {
+  beforeAll(async () => {
+    const { apply } = await pushSchema(
+      { ...didcommConnectionSchema, ...coreContextSchema },
+      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+      db as unknown as PgDatabase<any>
+    )
+    await apply()
+
+    await db.insert(context).values({
+      contextCorrelationId: 'b2fc0867-d0d1-4182-ade7-813b695d43c2',
+    })
+    await db.insert(didcommConnection).values({
+      createdAt: new Date(),
+      updatedAt: new Date(),
+      contextCorrelationId: 'b2fc0867-d0d1-4182-ade7-813b695d43c2',
+      id: 'db0ddc8f-5339-431b-9675-749f8a2ac92f',
+      role: DidExchangeRole.Requester,
+      state: DidExchangeState.Abandoned,
+      connectionTypes: ['one', 'three', 'four'],
+      invitationDid: 'some string',
+      customTags: {
+        myCustomTag: ['First', 'Second'],
+        anotherCustomTag: true,
+      },
+    })
+  })
+
+  test('should correctly query column', async () => {
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              invitationDid: 'some other string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(0)
+  })
+
+  test('should correctly query array value column', async () => {
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one', 'three'],
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one', 'two', 'five'],
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(0)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one'],
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+  })
+
+  test('should correctly query custom tag column value', async () => {
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              anotherCustomTag: true,
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              anotherCustomTag: false,
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(0)
+  })
+
+  test('should correctly query custom tag column array value', async () => {
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              myCustomTag: ['First', 'Second'],
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              myCustomTag: ['Third'],
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(0)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              myCustomTag: ['First'],
+              invitationDid: 'some string',
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+  })
+
+  test('should correctly query with and/or', async () => {
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one'],
+              $or: [
+                {
+                  connectionTypes: ['two'],
+                },
+                {
+                  connectionTypes: ['three'],
+                },
+              ],
+              $and: [
+                {
+                  connectionTypes: ['four'],
+                },
+              ],
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one'],
+              $or: [
+                {
+                  connectionTypes: ['two'],
+                },
+                {
+                  connectionTypes: ['three'],
+                },
+              ],
+              $and: [
+                {
+                  connectionTypes: ['four'],
+                },
+                {
+                  connectionTypes: ['two'],
+                },
+              ],
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(0)
+  })
+
+  test('should correctly query with not', async () => {
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one'],
+              $not: {
+                $and: [
+                  {
+                    connectionTypes: ['two'],
+                  },
+                  {
+                    connectionTypes: ['five'],
+                  },
+                ],
+                $or: [
+                  {
+                    connectionTypes: ['two'],
+                  },
+                  {
+                    connectionTypes: ['three'],
+                  },
+                ],
+              },
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(1)
+
+    expect(
+      await db
+        .select()
+        .from(didcommConnection)
+        .where(
+          queryToDrizzlePostgres<ConnectionRecord>(
+            {
+              connectionTypes: ['one'],
+              $not: {
+                $and: [
+                  {
+                    connectionTypes: ['two'],
+                  },
+                  {
+                    connectionTypes: ['five'],
+                  },
+                  {
+                    connectionTypes: ['three'],
+                  },
+                ],
+                $or: [
+                  {
+                    connectionTypes: ['two'],
+                  },
+                  {
+                    connectionTypes: ['three'],
+                  },
+                ],
+              },
+            },
+            didcommConnection
+          )
+        )
+    ).toHaveLength(0)
+  })
+})