Make sure size() can still be called once we reached the end of result set

8 files changed, 28 insertions, 12 deletions

diff --git a/odb/pgsql/no-id-object-result.hxx b/odb/pgsql/no-id-object-result.hxx
index 0bd3bee..2bd1b47 100644
--- a/odb/pgsql/no-id-object-result.hxx
+++ b/odb/pgsql/no-id-object-result.hxx
@@ -59,6 +59,7 @@ namespace odb
     private:
       details::shared_ptr<select_statement> statement_;
       statements_type& statements_;
+      std::size_t count_;
     };
   }
 }
diff --git a/odb/pgsql/no-id-object-result.txx b/odb/pgsql/no-id-object-result.txx
index 98e1b48..6b1e921 100644
--- a/odb/pgsql/no-id-object-result.txx
+++ b/odb/pgsql/no-id-object-result.txx
@@ -25,7 +25,8 @@ namespace odb
                               statements_type& statements)
         : base_type (statements.connection ().database ()),
           statement_ (statement),
-          statements_ (statements)
+          statements_ (statements),
+          count_ (0)
     {
     }
 
@@ -76,7 +77,9 @@ namespace odb
     {
       this->current (pointer_type ());
 
-      if (!statement_->next ())
+      if (statement_->next ())
+        count_++;
+      else
       {
         statement_->free_result ();
         this->end_ = true;
@@ -93,7 +96,7 @@ namespace odb
     std::size_t no_id_object_result_impl<T>::
     size ()
     {
-      return statement_->result_size ();
+      return this->end_ ? count_ : statement_->result_size ();
     }
   }
 }
diff --git a/odb/pgsql/polymorphic-object-result.hxx b/odb/pgsql/polymorphic-object-result.hxx
index 9dc8dcd..795a864 100644
--- a/odb/pgsql/polymorphic-object-result.hxx
+++ b/odb/pgsql/polymorphic-object-result.hxx
@@ -75,6 +75,7 @@ namespace odb
     private:
       details::shared_ptr<select_statement> statement_;
       statements_type& statements_;
+      std::size_t count_;
     };
   }
 }
diff --git a/odb/pgsql/polymorphic-object-result.txx b/odb/pgsql/polymorphic-object-result.txx
index 9d50427..02a71a0 100644
--- a/odb/pgsql/polymorphic-object-result.txx
+++ b/odb/pgsql/polymorphic-object-result.txx
@@ -28,7 +28,8 @@ namespace odb
                                     statements_type& sts)
         : base_type (sts.connection ().database ()),
           statement_ (st),
-          statements_ (sts)
+          statements_ (sts),
+          count_ (0)
     {
     }
 
@@ -167,7 +168,9 @@ namespace odb
     {
       this->current (pointer_type ());
 
-      if (!statement_->next ())
+      if (statement_->next ())
+        count_++;
+      else
       {
         statement_->free_result ();
         this->end_ = true;
@@ -259,7 +262,7 @@ namespace odb
     std::size_t polymorphic_object_result_impl<T>::
     size ()
     {
-      return statement_->result_size ();
+      return this->end_ ? count_ : statement_->result_size ();
     }
   }
 }
diff --git a/odb/pgsql/simple-object-result.hxx b/odb/pgsql/simple-object-result.hxx
index 5edee2c..d572fea 100644
--- a/odb/pgsql/simple-object-result.hxx
+++ b/odb/pgsql/simple-object-result.hxx
@@ -67,6 +67,7 @@ namespace odb
     private:
       details::shared_ptr<select_statement> statement_;
       statements_type& statements_;
+      std::size_t count_;
     };
   }
 }
diff --git a/odb/pgsql/simple-object-result.txx b/odb/pgsql/simple-object-result.txx
index f7b2e1c..81b5f73 100644
--- a/odb/pgsql/simple-object-result.txx
+++ b/odb/pgsql/simple-object-result.txx
@@ -27,7 +27,8 @@ namespace odb
                         statements_type& statements)
         : base_type (statements.connection ().database ()),
           statement_ (statement),
-          statements_ (statements)
+          statements_ (statements),
+          count_ (0)
     {
     }
 
@@ -84,7 +85,9 @@ namespace odb
     {
       this->current (pointer_type ());
 
-      if (!statement_->next ())
+      if (statement_->next ())
+        count_++;
+      else
       {
         statement_->free_result ();
         this->end_ = true;
@@ -136,7 +139,7 @@ namespace odb
     std::size_t object_result_impl<T>::
     size ()
     {
-      return statement_->result_size ();
+      return this->end_ ? count_ : statement_->result_size ();
     }
   }
 }
diff --git a/odb/pgsql/view-result.hxx b/odb/pgsql/view-result.hxx
index 341065d..1bdb2d2 100644
--- a/odb/pgsql/view-result.hxx
+++ b/odb/pgsql/view-result.hxx
@@ -59,6 +59,7 @@ namespace odb
     private:
       details::shared_ptr<select_statement> statement_;
       statements_type& statements_;
+      std::size_t count_;
     };
   }
 }
diff --git a/odb/pgsql/view-result.txx b/odb/pgsql/view-result.txx
index 932d19b..9c22133 100644
--- a/odb/pgsql/view-result.txx
+++ b/odb/pgsql/view-result.txx
@@ -25,7 +25,8 @@ namespace odb
                       statements_type& statements)
         : base_type (statements.connection ().database ()),
           statement_ (statement),
-          statements_ (statements)
+          statements_ (statements),
+          count_ (0)
     {
     }
 
@@ -76,7 +77,9 @@ namespace odb
     {
       this->current (pointer_type ());
 
-      if (!statement_->next ())
+      if (statement_->next ())
+        count_++;
+      else
       {
         statement_->free_result ();
         this->end_ = true;
@@ -93,7 +96,7 @@ namespace odb
     std::size_t view_result_impl<T>::
     size ()
     {
-      return statement_->result_size ();
+      return this->end_ ? count_ : statement_->result_size ();
     }
   }
 }