How to return a DatabaseStatementBase object
Bandaid fix, instead of return $query; test to see if the class of $query is of DatabaseStatementBase or else create a new empty instance of that class and return that ... Eg:

if (is_a($query, 'DatabaseStatementBase')) {
   return $query;
} else {
  return $empty = new DatabaseStatementBase(Database::getConnection());
}