zeus/website/backend/app/libs/resty/moongoo/cursor.lua

local cbson = require("cbson")
local bit = require("bit")


local function check_bit(num, bitnum)
  return bit.band(num,math.pow(2,bitnum)) ~= 0 -- and true or false
end

local _M = {}

local mt = { __index = _M }

function _M.new(collection, query, fields, explain, id)
  return setmetatable(
    {
      _collection = collection,
      _query = query,
      _fields = fields,
      _id = id or cbson.uint(0),
      _skip = 0,
      _limit = 0,
      _docs = {},
      _started = false,
      _cnt = 0,
      _comment = nil,
      _hint = nil,
      _max_scan = nil ,
      _max_time_ms = nil,
      _read_preference = nil,
      _snapshot = nil,
      _sort = nil,
      _await = false,
      _tailable = false,
      _explain = explain or false
    },
  mt)
end

function _M.tailable(self, tailable)
  self._tailable = tailable
  return self
end

function _M.await(self, await)
  self._await = await
  return self
end


function _M.comment(self, comment)
  self._comment = comment
  return self
end

function _M.hint(self, hint)
  self._hint = hint
  return self
end

function _M.max_scan(self, max_scan)
  self._max_scan = max_scan
  return self
end

function _M.max_time_ms(self, max_time_ms)
  self._max_time_ms = max_time_ms
  return self
end

function _M.read_preference(self, read_preference)
  self._read_preference = read_preference
  return self
end

function _M.snapshot(self, snapshot)
  self._snapshot = snapshot
  return self
end

function _M.sort(self, sort)
  self._sort = sort
  return self
end


function _M.clone(self, explain)
  local clone = self.new(self._collection, self._query, self._fields, explain)
  clone:limit(self._limit)
  clone:skip(self._skip)

  clone:comment(self._comment)
  clone:hint(self._hint)
  clone:max_scan(self._max_scan)
  clone:max_time_ms(self._max_time_ms)
  clone:read_preference(self._read_preference)
  clone:snapshot(self._snapshot)
  clone:sort(self._sort)

  return clone
end

function _M.skip(self, skip)
  if self._started then
    print("Can's set skip after starting cursor")
  else
    self._skip = skip
  end
  return self
end

function _M.limit(self, limit)
  if self._started then
    print("Can's set limit after starting cursor")
  else
    self._limit = limit
  end
  return self
end

function _M._build_query(self)
  local ext = {}
  if self._comment then ext['$comment'] = self._comment end
  if self._explain then ext['$explain'] = true end


  if self._hint then ext['$hint'] = self._hint end
  if self._max_scan then ext['$maxScan'] = self._max_scan end
  if self._max_time_ms then ext['$maxTimeMS'] = self._max_time_ms end
  if self._read_preference then ext['$readPreference'] = self._read_preference end
  if self._snapshot then ext['$snapshot'] = true end
  if self._sort then ext['$orderby'] = self._sort end

  ext['$query'] = self._query

  return cbson.encode(ext)
end

function _M.next(self)
  local moongoo, err = self._collection._db._moongoo:connect()
  if not moongoo then
    return nil, err
  end

  if self:_finished() then
    if self._id ~= cbson.uint(0) then
      self._collection._db._moongoo.connection:_kill_cursors(self._id)
      self._id = cbson.uint(0)
    end
    return nil, "no more data"
  end

  if (not self._started) and (self._id == cbson.uint(0)) then

    -- query and add id and batch
    local flags, id, from, number, docs = self._collection._db._moongoo.connection:_query(self._collection:full_name(), self:_build_query(), self._skip, self._limit, self._fields, {tailable = self._tailable, await = self._await})

    flags = tonumber(tostring(flags)) -- bitop can't work with cbson.int, so...

    if check_bit(flags, 1) then -- QueryFailure
      return nil, docs[1]['$err'] -- why is this $err and not errmsg, like others??
    end
    self._id = id
    self:add_batch(docs)
  elseif #self._docs == 0 and self._id ~= cbson.uint(0) then
    -- we have something to fetch - get_more and add_batch
    local flags, id, from, number, docs = self._collection._db._moongoo.connection:_get_more(self._collection:full_name(), self._limit, self._id)

    flags = tonumber(tostring(flags)) -- bitop can't work with cbson.int, so...

    if check_bit(flags, 0) then -- QueryFailure
      return nil, "wrong cursor id"
    end
    self:add_batch(docs)
    self._id = id

  elseif #self._docs == 0 then--or self._id == cbson.uint(0) then
    return nil, "no more data"
  end
  self._cnt = self._cnt+1
  return table.remove(self._docs, 1) or nil, 'No more data'
end

function _M.all(self)
  local docs = {}
  while true do
    local doc = self:next()
    if doc == nil then break end
    table.insert(docs, doc)
  end
  return docs
end

function _M.rewind(self)
  self._started = false
  self._docs = {}
  self._collection._db._moongoo.connection:_kill_cursors(self._id)
  self._id = cbson.uint(0)
  return self
end

function _M.count(self)
  local doc, err = self._collection._db:cmd(
    { count = self._collection.name },
    {
      query = self._query,
      skip = self._skip,
      limit = self._limit
    }
  )
  if not doc then
    return nil, err
  end

  return doc and doc.n or 0
end

function _M.distinct(self, key)
  local doc, err = self._collection._db:cmd(
    { distinct = self._collection.name },
    {
      query = self._query,
      key = key
    }
  )
  if not doc then
    return nil, err
  end

  return doc and doc.values or {}
end

function _M.explain(self)
  return self:clone(true):sort(nil):next()
end

function _M.add_batch(self, docs)
  self._started = true
  for k,v in ipairs(docs) do
    table.insert(self._docs, v)
  end
  return self
end

function _M._finished(self)
  if self._limit == 0 then
    return false
  else
    if self._cnt >= math.abs(self._limit) then
      return true
    else
      return false
    end
  end
end

return _M