piazzaoptical.com/server/piazza-lmdb.lua

lightningmdb_lib=require "lightningmdb"

lightningmdb = _VERSION>="Lua 5.2" and lightningmdb_lib or lightningmdb
MDB = setmetatable({}, {__index = function(t, k)
  return lightningmdb["MDB_" .. k]
end})

--local binser = require'binser'
local bitser = require'bitser/bitser'
bitser.includeMetatables(false)

local lrandom = require'random'  -- 'lrandom' luarocks package
local lrandom_sequence = lrandom.new(os.time())

-- helper function to generate random integer
local randomid = function()
  local min = 0
  local max = 256*256*256*256-1  -- 32-bits exactly;  TODO: confirm this isn't larger than 32-bits
  return lrandom.value(lrandom_sequence, min, max)
end

local piazzadb = {}
local dbmt = {}
dbmt.__index = dbmt
local recordmt = {}
recordmt.__index = recordmt

piazzadb.init = function()
  local env = lightningmdb.env_create()
  env:set_mapsize(4096*1000000000) -- 4GB max DB size
  env:open('./piazza.lmdb', MDB.NOSUBDIR+MDB.NOLOCK, 420)
  
  local transaction = env:txn_begin(nil, 0)
  local dbi = transaction:dbi_open(nil, 0) -- only need one db handle; can be reused in later transactions
  transaction:commit()
  
  local db = {}
  db.env = env
  db.dbi = dbi
  setmetatable(db, dbmt)
  
  return db
end

dbmt.serializerecord = function(db, record)
  record.db = nil
  --return binser.serialize(record)
  return bitser.dumps(record)
end

dbmt.deserializerecord = function(db, record_string)
  --local record = binser.deserializeN(record_string, 1)
  local record = bitser.loads(record_string)
  record.db = db
  setmetatable(record, recordmt)
  return record
end


-- helper function to get a new available id
dbmt.newid = function(db)
  local transaction = db.env:txn_begin(nil, 0)
  -- generate new ids and check for pre-existing entry in DB until an available ID is found:
  local id, v
  local count = 2
  repeat
    id = randomid()
    count=count+1
    v = transaction:get(db.dbi, id)
  until v == nil
  transaction:commit()
  
  return id
end

dbmt.getstring = function(db, id)
  local transaction = db.env:txn_begin(nil, 0)
  local v = transaction:get(db.dbi, id)
  transaction:commit()
  return v
end

dbmt.setstring = function(db, id, s)  -- DON'T USE THIS DIRECTLY! Everything that goes into the DB must pass through binser (the serializer), because if not it there's no way to loop through all records, because it's impossible to tell what's a record (i.e. binser table) and what's a raw string
  local transaction = db.env:txn_begin(nil, 0)
  transaction:put(db.dbi, id, s, 0)
  transaction:commit()
end

dbmt.delete = function(db, id)
  local transaction = db.env:txn_begin(nil, 0)
  transaction:del(db.dbi, id, nil)
  transaction:commit()
end


-- gets record cast from raw string into lua table, outfitted with metatable methods
dbmt.getrecord = function(db, id)
  local record_string = db:getstring(id)
  return db:deserializerecord( record_string )
end

-- stores Lua table as record
dbmt.setrecord = function(db, id, t)
  local record_string = db:serializerecord(t)
  db:setstring(id, record_string)
end


-- stores Lua table as new record, assigning it a random ID
dbmt.insertrecord = function(db, t)
  local id = db:newid()
  t.id = id
  setmetatable(t, recordmt)
  local record_string = db:serializerecord(t)
  db:setstring(id, record_string)
  return id
end

dbmt.records = function(db)
  local transaction = db.env:txn_begin(nil, 0)
  local cursor = transaction:cursor_open(db.dbi)
  --cursor:get(nil, nil, MDB.FIRST)
  local k, v
  
  return function()
    k, v = cursor:get(k, MDB.NEXT)
    if k then
      return k, db:deserializerecord(v)
    else
      cursor:close()
      transaction:commit()
    end
  end
end

dbmt.findAll = function(db, fn)
  local t = {}
  for id,record in db:records() do
    if fn(record, id) then table.insert(t, record) end
  end
  return t
end

dbmt.findOne = function(db, fn)
  for id,record in db:records() do
    if fn(record, id) then return record end
  end
  return nil
end

dbmt.close = function(db)
  db.env:dbi_close(db.dbi)
end

dbmt.copyandtransformall = function(db, list, structure)
  local copy = {}
  for _, el in ipairs(list) do
    local record
    if type(el) == 'number' then -- if it's an ID, get the record
      record = db:getrecord(el)
    elseif type(el) == 'table' then -- if it's a record itself
      record = el
    end
    table.insert(copy, record:copyandtransform(structure))
  end
  return copy
end

dbmt.findByStructure = function(db, structure)
  local t = {}
  for _,record in db:records() do
    local does_match = true
    for k,v in pairs(structure) do
      if type(v) == 'table' then  -- TODO: everything within this `if` needs to be fleshed-out; it doesn't do JOINs yet
        if v[1] == 'list' then
          if type(record[k]) ~= 'table' then  -- TODO: find more accurate way to check whether it's a list, which excludes 'object' tables but includes empty lists (i.e. `#record[k] > 0` won't work)
            does_match = false
            break
          end
        end
      else
        if record[k] ~= v then
          does_match = false
          break
        end
      end
    end
    if does_match then table.insert(t, record) end
  end
  return t
end

recordmt.save = function(record)
  record.db:setrecord(record.id, record)
end

recordmt.delete = function(record)
  record.db:delete(record.id)
end

recordmt.getref = function(record, key)
  return record.db:getrecord(record[key..'_id'])
end

recordmt.copyandtransform = function(record, structure)
  local db = record.db
  local copy = {}
  for k,op in pairs(structure) do
    if op=='copy' then
      copy[k] = record[k]
    elseif type(op)=='table' then  -- it can either be a 'list' directive, or a substructure
      if op[1] == 'list' or op[1] == 'array' then
        local subcopy = {}
        local substructure = op[2]
        if substructure == nil then -- if it's a list of scalars, e.g. strings
          for _,el in ipairs(record[k]) do
            table.insert(subcopy, el)
          end
        else  -- if it's a list of tables/records, of which we want just a portion as per the substructure
          if record[k] ~= nil then
            for _,el in ipairs(record[k]) do
              local referenced_record = db:getrecord(el)
              table.insert(subcopy, referenced_record:copyandtransform(substructure))
            end
          end
        end
        copy[k] = subcopy
      elseif op[1] == nil then -- if it's just a single referenced table
        local substructure = op
        local referenced_record = db:getrecord(record[k]) -- e.g. k='some_record_id'
        copy[k] = referenced_record:copyandtransform(substructure)
      end
    end
  end
  return copy
end

return piazzadb