Merge pull request #632 from Mashape/dao/count

Adding a function to count the entities in a table

kong/dao/cassandra/base_dao.lua

@@ -548,6 +548,18 @@ function BaseDao:find_by_keys(where_t, page_size, paging_state)
   return res, err, filtering
 end
 
+-- Retrieve the number of rows from the given columns/value table.
+-- @param `where_t`      (Optional) columns/values table by which to count entities.
+-- @return `res`
+-- @return `err`
+-- @return `filtering`   A boolean indicating if ALLOW FILTERING was needed by the query
+function BaseDao:count_by_keys(where_t)
+  local select_q, where_columns, filtering = query_builder.count(self._table, where_t, self._column_family_details)
+  local res, err = self:execute(select_q, where_columns, where_t, {})
+
+  return (#res >= 1 and table.remove(res, 1).count or 0), err, filtering
+end
+
 -- Retrieve a page of the table attached to the DAO.
 -- @param  `page_size`    Size of the page to retrieve (number of rows).
 -- @param  `paging_state` Start page from given offset. See lua-resty-cassandra's :execute() option.

kong/dao/cassandra/query_builder.lua

@@ -15,6 +15,10 @@ local function select_fragment(column_family, select_columns)
   return string.format("SELECT %s FROM %s", select_columns, column_family)
 end
 
+local function count_fragment(column_family)
+  return string.format("SELECT COUNT(*) FROM %s", column_family)
+end
+
 local function insert_fragment(column_family, insert_values)
   local values_placeholders, columns = {}, {}
   for column, value in pairs(insert_values) do
@@ -129,6 +133,22 @@ function _M.select(column_family, where_t, column_family_details, select_columns
   return trim(string.format("%s %s", select_str, where_str)), columns, needed_filtering
 end
 
+-- Generate a COUNT query with an optional WHERE instruction.
+-- If building a WHERE instruction, we need some additional informations about the column family.
+-- @param `column_family`         Name of the column family
+-- @param `column_family_details` Additional infos about the column family (partition key, clustering key, indexes)
+-- @return `query`                The SELECT query
+-- @return `columns`              An list of columns to bind for the query, in the order of the placeholder markers (?)
+-- @return `needs_filtering`      A boolean indicating if ALLOW FILTERING was added to this query or not
+function _M.count(column_family, where_t, column_family_details)
+  assert(type(column_family) == "string", "column_family must be a string")
+
+  local count_str = count_fragment(column_family)
+  local where_str, columns, needed_filtering = where_fragment(where_t, column_family_details)
+
+  return trim(string.format("%s %s", count_str, where_str)), columns, needed_filtering
+end
+
 -- Generate an INSERT query.
 -- @param `column_family` Name of the column family
 -- @param `insert_values` A columns/values table of values to insert

spec/integration/dao/cassandra/base_dao_spec.lua

@@ -182,6 +182,34 @@ describe("Cassandra", function()
 
     end) -- describe :insert()
 
+    describe(":count_by_keys()", function()
+
+      it("should properly count the items in a table", function()
+        local count, err = dao_factory.apis:count_by_keys()
+        assert.falsy(err)
+        assert.truthy(count)
+        assert.are.same(4, count)
+      end)
+
+      it("should properly count the items in a table with keys", function()
+        local count, err = dao_factory.apis:count_by_keys({name="test.com"})
+        assert.falsy(err)
+        assert.truthy(count)
+        assert.are.same(1, count)
+
+        count, err = dao_factory.apis:count_by_keys({name="test.com.com"})
+        assert.falsy(err)
+        assert.truthy(count)
+        assert.are.same(0, count)
+
+        count, err = dao_factory.apis:count_by_keys({name=""})
+        assert.falsy(err)
+        assert.truthy(count)
+        assert.are.same(0, count)
+      end)
+
+    end)
+
     describe(":update()", function()
 
       it("should error if called with invalid parameters", function()

spec/unit/dao/cassandra/query_builder_spec.lua

@@ -88,6 +88,20 @@ describe("Query Builder", function()
 
   end)
 
+  describe("COUNT", function()
+
+    it("should build a COUNT query", function()
+      local q = builder.count("apis")
+      assert.equal("SELECT COUNT(*) FROM apis", q)
+    end)
+
+    it("should build a COUNT query with WHERE keys", function()
+      local q = builder.count("apis", {id="123", name="mockbin"})
+      assert.equal("SELECT COUNT(*) FROM apis WHERE name = ? AND id = ? ALLOW FILTERING", q)
+    end)
+
+  end)
+
   describe("INSERT", function()
 
     it("should build an INSERT query", function()