SQLite as a Graph Database

DB.test.ts

import DB from './DB'

describe('DB', () => {
  let db: DB
  const data = { name: 'Tom Cook', type: ['person', 'CEO'] }

  beforeEach(() => {
    db = new DB()

    db.addNode('1', { name: 'Apple Computer Company', type: ['company', 'start-up'], founded: 'April 1, 1976' })
    db.addNode('2', { name: 'Steve Wozniak', type: ['person', 'engineer', 'founder'] })
    db.addNode('3', { name: 'Steve Jobs', type: ['person', 'designer', 'founder'] })
    db.addNode('4', { name: 'Ronald Wayne', type: ['person', 'administrator', 'founder'] })
    db.addNode('5', { name: 'Mike Markkula', type: ['person', 'investor'] })

    db.connectNodes('2', '1', { action: 'founded' })
    db.connectNodes('3', '1', { action: 'founded' })
    db.connectNodes('4', '1', { action: 'founded' })
    db.connectNodes('5', '1', { action: 'invested', equity: 80000, debt: 170000 })
    db.connectNodes('1', '4', { action: 'divested', amount: 800, date: 'April 12, 1976' })
    db.connectNodes('2', '3')
  })

  describe('#addNode()', () => {
    it('should add a node', () => {
      const before = db.findNode('0')
      expect(before).toEqual({})

      db.addNode('0', data)

      const after = db.findNode('0')
      expect(after).toEqual({
        id: '0',
        ...data
      })
    })
  })

  describe('#upsertNode()', () => {
    it('should update an existing node', () => {
      const before = db.findNode('2')
      expect(before.nickname).toEqual(undefined)

      db.upsertNode('2', { nickname: 'Woz' })

      const after = db.findNode('2')
      expect(after.nickname).toEqual('Woz')
    })

    it('should create a new node if there\'s no existing node', () => {
      const before = db.findNode('6')
      expect(before).toEqual({})

      db.upsertNode('6', data)

      const after = db.findNode('6')
      expect(after).toEqual({
        id: '6',
        ...data
      })
    })
  })

  describe('#removeNode()', () => {
    it('should remove an existing node', () => {
      const before = db.findNode('1')
      expect(before.id).toEqual('1')

      db.removeNode('1')

      const after = db.findNode('1')
      expect(after).toEqual({})
    })

    it('should do nothing when removing a non-existent node', () => {
      const before = db.findNode('0')
      expect(before).toEqual({})

      db.removeNode('0')

      const after = db.findNode('0')
      expect(after).toEqual({})
    })
  })

  describe('#connectNodes()', () => {
    it('should connect 2 nodes', () => {
      db.addNode('0', data)

      const before = db.getConnections('0')
      expect(before).toEqual([])

      db.connectNodes('0', '1')

      const after = db.getConnections('0')
      expect(after).toEqual([{ source: '0', target: '1', properties: {} }])
    })
  })

  describe('#getConnectionsOneWay()', () => {
    beforeEach(() => {
      db.addNode('0', data)
      db.connectNodes('0', '1')
      db.connectNodes('1', '0')
    })

    describe('#getConnectionsIn()', () => {
      it('should return the inbound connections', () => {
        const result = db.getConnectionsOneWay('0', db.getConnectionsIn)
        expect(result).toEqual([{ source: '0', target: '1', properties: {} }])
      })
    })

    describe('#getConnectionsOut()', () => {
      it('should return the outbound connections', () => {
        const result = db.getConnectionsOneWay('0', db.getConnectionsOut)
        expect(result).toEqual([{ source: '1', target: '0', properties: {} }])
      })
    })
  })

  describe('#findNode()', () => {
    it('should find a node by id', () => {
      const result = db.findNode('1')
      expect(result).toEqual({
        id: '1',
        name: 'Apple Computer Company',
        type: ['company', 'start-up'],
        founded: 'April 1, 1976'
      })
    })
  })

  describe('#findNodes()', () => {
    it('should find nodes via the where- and search functions', () => {
      let result

      result = db.findNodes({ name: 'Apple Computer Company', founded: 'April 1, 1976' })
      expect(result).toEqual([{
        id: '1',
        name: 'Apple Computer Company',
        type: ['company', 'start-up'],
        founded: 'April 1, 1976'
      }])

      result = db.findNodes({ name: 'Steve' }, db.searchLike, db.searchStartsWith)
      expect(result).toEqual([
        { id: '2', name: 'Steve Wozniak', type: ['person', 'engineer', 'founder'] },
        { id: '3', name: 'Steve Jobs', type: ['person', 'designer', 'founder'] }
      ])
    })
  })

  describe('#traverse()', () => {
    it('should traverse the graph', () => {
      function normalize (result: any[]): string[] {
        return result.reduce<string[]>((accum, item) => {
          if (item.type === '()') accum.push(item.id)
          return accum
        }, [])
      }

      let result

      result = normalize(db.traverse('2', '3'))
      expect(result).toEqual(['2', '1', '3'])

      result = normalize(db.traverse('4', '5'))
      expect(result).toEqual(['4', '1', '2', '3', '5'])

      result = normalize(db.traverse('5', undefined, db.findInboundNeighbors))
      expect(result).toEqual(['5'])

      result = normalize(db.traverse('5', undefined, db.findOutboundNeighbors))
      expect(result).toEqual(['5', '1', '4'])

      result = normalize(db.traverse('5', undefined, db.findNeighbors))
      expect(result).toEqual(['5', '1', '2', '3', '4'])
    })
  })
})

DB.ts

import fs from 'fs'
import path from 'path'
import Database, { Database as BetterSqlite3 } from 'better-sqlite3'

// Adaption of:
//  https://github.com/dpapathanasiou/simple-graph
//  https://pythonawesome.com/a-simple-graph-database-in-sqlite/

class DB {
  private readonly _db: BetterSqlite3

  constructor (filePath: string = ':memory:') {
    const db = new Database(filePath)
    db.pragma('foreign_keys = TRUE;')
    db.transaction(() => {
      db.exec(read('schema.sql'))
    })()
    this._db = db
  }

  addNode (id: string, data: Record<string, unknown>): void {
    const json = JSON.stringify({ ...data, id })
    const stmt = this._db.prepare(read('insert-node.sql'))
    this._db.transaction(() => stmt.run(json))()
  }

  upsertNode (id: string, data: Record<string, unknown>): void {
    const currentData = this.findNode(id)
    if (Object.keys(currentData).length === 0) {
      this.addNode(id, data)
    } else {
      const json = JSON.stringify({ ...currentData, ...data, id })
      const stmt = this._db.prepare(read('update-node.sql'))
      this._db.transaction(() => stmt.run(json, id))()
    }
  }

  removeNode (id: string): void {
    const stmtNode = this._db.prepare(read('delete-node.sql'))
    const stmtEdge = this._db.prepare(read('delete-edge.sql'))
    this._db.transaction(() => {
      stmtEdge.run(id, id)
      stmtNode.run(id)
    })()
  }

  connectNodes (sourceId: string, targetId: string, data: Record<string, unknown> = {}): void {
    const json = JSON.stringify(data)
    const stmt = this._db.prepare(read('insert-edge.sql'))
    this._db.transaction(() => stmt.run(sourceId, targetId, json))()
  }

  getConnections (id: string): Connection[] {
    const stmt = this._db.prepare(read('search-edges.sql'))
    return this.parseConnectionsResult(stmt.all([id, id]))
  }

  getConnectionsOneWay (id: string, directionFn = this.getConnectionsIn): Connection[] {
    const stmt = this._db.prepare(directionFn())
    return this.parseConnectionsResult(stmt.all(id))
  }

  parseConnectionsResult (connections: Array<Record<string, unknown>>): Connection[] {
    return connections.map(item => ({
      source: String(item.source),
      target: String(item.target),
      properties: JSON.parse(String(item.properties))
    }))
  }

  findNode (id: string): Record<string, unknown> {
    const stmt = this._db.prepare(read('search-node-by-id.sql'))
    const row: Record<string, unknown> | undefined = stmt.get(id)
    const json = (row != null) ? row.body as any : '{}'
    return JSON.parse(json)
  }

  findNodes (data: Record<string, unknown>, whereFn = this.searchWhere, searchFn = this.searchEquals): Node[] {
    const stmt = this._db.prepare(read('search-node.sql').replace('<replace>', whereFn.bind(this)(data)))
    const rows = stmt.all(searchFn(data))
    return this.parseSearchResult(rows)
  }

  parseSearchResult (results: Array<Record<string, unknown>>, idx = 0): Node[] {
    return results.map((result) => JSON.parse(result[Object.keys(result)[idx]] as string))
  }

  traverse (source: string, target?: string, neighborsFn = this.findNeighbors): Object[] {
    const path: Object[] = []
    const stmt = this._db.prepare(neighborsFn())
    let header = null
    for (const row of stmt.iterate({ source })) {
      if (header === null) {
        header = row
        continue
      }
      const { x: id, y: type, obj: data } = row
      path.push({
        id,
        type,
        data: JSON.parse(data)
      })
      if (id === target && type === '()') {
        break
      }
    }
    return path
  }

  // ---------- The following functions are only used as Higher-Order-Functions ----------
  getConnectionsIn (): string {
    return read('search-edges-inbound.sql')
  }

  getConnectionsOut (): string {
    return read('search-edges-outbound.sql')
  }

  findNeighbors (): string {
    return read('traverse.sql')
  }

  findOutboundNeighbors (): string {
    return read('traverse-outbound.sql')
  }

  findInboundNeighbors (): string {
    return read('traverse-inbound.sql')
  }

  searchWhere (props: Record<string, unknown>, predicate: string = '='): string {
    return Object.keys(props).map((key) => `json_extract(body, '$.${key}') ${predicate} ?`).join(' AND ')
  }

  searchLike (props: Record<string, unknown>): string {
    return this.searchWhere(props, 'LIKE')
  }

  searchEquals (props: Record<string, unknown>): string[] {
    return Object.values(props).map((value: unknown) => String(value))
  }

  searchStartsWith (props: Record<string, unknown>): string[] {
    return Object.values(props).map((value: unknown) => String(value) + '%')
  }

  searchContains (props: Record<string, unknown>): string[] {
    return Object.values(props).map((value: unknown) => '%' + String(value) + '%')
  }
}

function read (filePath: string): string {
  return fs.readFileSync(path.join(__dirname, filePath), 'utf-8')
}

type Node = Record<string, unknown> & { id: string }

interface Connection {
  source: string
  target: string
  properties: Record<string, unknown>
}

interface Object {
  id: string
  type: string
  data: Record<string, unknown>
}

export default DB

delete-edge.sql

DELETE FROM edges WHERE source = ? OR target = ?;

delete-node.sql

DELETE FROM nodes WHERE id = ?;

insert-edge.sql

INSERT INTO edges VALUES(?, ?, json(?));

insert-node.sql

INSERT INTO nodes VALUES(json(?));

schema.sql

CREATE TABLE IF NOT EXISTS nodes (
  body TEXT,
  id   TEXT GENERATED ALWAYS AS (json_extract(body, '$.id')) VIRTUAL NOT NULL UNIQUE
);

CREATE INDEX IF NOT EXISTS id_idx ON nodes(id);

CREATE TABLE IF NOT EXISTS edges (
  source     TEXT,
  target     TEXT,
  properties TEXT,
  UNIQUE(source, target, properties) ON CONFLICT REPLACE,
  FOREIGN KEY(source) REFERENCES nodes(id),
  FOREIGN KEY(target) REFERENCES nodes(id)
);

CREATE INDEX IF NOT EXISTS source_idx ON edges(source);
CREATE INDEX IF NOT EXISTS target_idx ON edges(target);

search-edges-inbound.sql

SELECT * FROM edges WHERE source = ?;

search-edges-outbound.sql

SELECT * FROM edges WHERE target = ?;

search-edges.sql

SELECT * FROM edges WHERE source = ?
UNION
SELECT * FROM edges WHERE target = ?;

search-node-by-id.sql

SELECT body FROM nodes WHERE id = ?;

search-node.sql

SELECT body FROM nodes WHERE <replace>

traverse-inbound.sql

WITH RECURSIVE traverse(x, y, obj) AS (
  SELECT :source, '()', '{}'
  UNION
  SELECT id, '()', body FROM nodes JOIN traverse ON id = x
  UNION
  SELECT source, '<-', properties FROM edges JOIN traverse ON target = x
) SELECT x, y, obj FROM traverse;

traverse-outbound.sql

WITH RECURSIVE traverse(x, y, obj) AS (
  SELECT :source, '()', '{}'
  UNION
  SELECT id, '()', body FROM nodes JOIN traverse ON id = x
  UNION
  SELECT target, '->', properties FROM edges JOIN traverse ON source = x
) SELECT x, y, obj FROM traverse;

traverse.sql

WITH RECURSIVE traverse(x, y, obj) AS (
  SELECT :source, '()', '{}'
  UNION
  SELECT id, '()', body FROM nodes JOIN traverse ON id = x
  UNION
  SELECT source, '<-', properties FROM edges JOIN traverse ON target = x
  UNION
  SELECT target, '->', properties FROM edges JOIN traverse ON source = x
) SELECT x, y, obj FROM traverse;

update-node.sql

UPDATE nodes SET body = json(?) WHERE id = ?;