Allow a more streaming interface to scans

This permits the NoSQL implementation to be informed when we abort
a scan early, because its scan routine can implement and receive the
close() method.  It also gives us a way to stream the JDBC results,
so we aren't creating a full result set in memory.

Change-Id: Iaf53a1803fe985e49236e10bde5d838846e0c8f0
Signed-off-by: Shawn O. Pearce <sop@google.com>

8 files changed, 258 insertions, 44 deletions

diff --git a/src/main/java/com/google/gwtorm/client/OrmRuntimeException.java b/src/main/java/com/google/gwtorm/client/OrmRuntimeException.java
new file mode 100644
index 0000000..d021072
--- /dev/null
+++ b/src/main/java/com/google/gwtorm/client/OrmRuntimeException.java
@@ -0,0 +1,32 @@
+// Copyright 2008 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.gwtorm.client;
+
+/**
+ * Any data store read or write error.
+ */
+public class OrmRuntimeException extends RuntimeException {
+  public OrmRuntimeException(final String message) {
+    super(message);
+  }
+
+  public OrmRuntimeException(final String message, final Throwable why) {
+    super(message, why);
+  }
+
+  public OrmRuntimeException(final Throwable why) {
+    super(why);
+  }
+}
diff --git a/src/main/java/com/google/gwtorm/client/impl/AbstractResultSet.java b/src/main/java/com/google/gwtorm/client/impl/AbstractResultSet.java
new file mode 100644
index 0000000..ab5f860
--- /dev/null
+++ b/src/main/java/com/google/gwtorm/client/impl/AbstractResultSet.java
@@ -0,0 +1,62 @@
+// Copyright 2010 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.gwtorm.client.impl;
+
+import com.google.gwtorm.client.ResultSet;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * Simple implementation of a {@link ResultSet}.
+ *
+ * @param <T> type of the object to be returned from the result set.
+ */
+public abstract class AbstractResultSet<T> implements ResultSet<T> {
+  @Override
+  public final Iterator<T> iterator() {
+    return new Iterator<T>() {
+      @Override
+      public boolean hasNext() {
+        return AbstractResultSet.this.hasNext();
+      }
+
+      @Override
+      public T next() {
+        return AbstractResultSet.this.next();
+      }
+
+      @Override
+      public void remove() {
+        throw new UnsupportedOperationException();
+      }
+    };
+  }
+
+  public List<T> toList() {
+    List<T> r = new ArrayList<T>();
+    for (T obj : this) {
+      r.add(obj);
+    }
+    return r;
+  }
+
+  /** @return true if another result remains, false otherwise. */
+  protected abstract boolean hasNext();
+
+  /** @return next result. */
+  protected abstract T next();
+}
diff --git a/src/main/java/com/google/gwtorm/jdbc/AccessGen.java b/src/main/java/com/google/gwtorm/jdbc/AccessGen.java
index 4be2247..ede59a0 100644
--- a/src/main/java/com/google/gwtorm/jdbc/AccessGen.java
+++ b/src/main/java/com/google/gwtorm/jdbc/AccessGen.java
@@ -750,7 +750,7 @@ class AccessGen implements Opcodes {
     mv.visitVarInsn(ALOAD, 0);
     mv.visitVarInsn(ALOAD, psvar);
     mv.visitMethodInsn(INVOKEVIRTUAL, superTypeName, "queryList", Type
-        .getMethodDescriptor(Type.getType(ListResultSet.class),
+        .getMethodDescriptor(Type.getType(com.google.gwtorm.client.ResultSet.class),
             new Type[] {Type.getType(PreparedStatement.class)}));
     mv.visitInsn(ARETURN);
     mv.visitMaxs(-1, -1);
diff --git a/src/main/java/com/google/gwtorm/jdbc/JdbcAccess.java b/src/main/java/com/google/gwtorm/jdbc/JdbcAccess.java
index afcab77..b75011a 100644
--- a/src/main/java/com/google/gwtorm/jdbc/JdbcAccess.java
+++ b/src/main/java/com/google/gwtorm/jdbc/JdbcAccess.java
@@ -126,28 +126,25 @@ public abstract class JdbcAccess<T, K extends Key<?>> extends
     }
   }
 
-  protected ListResultSet<T> queryList(final PreparedStatement ps)
-      throws OrmException {
+  protected com.google.gwtorm.client.ResultSet<T> queryList(
+      final PreparedStatement ps) throws OrmException {
+    final ResultSet rs;
     try {
+      rs = ps.executeQuery();
+      if (!rs.next()) {
+        rs.close();
+        ps.close();
+        return new ListResultSet<T>(Collections.<T> emptyList());
+      }
+    } catch (SQLException err) {
       try {
-        final ResultSet rs = ps.executeQuery();
-        try {
-          final ArrayList<T> r = new ArrayList<T>();
-          while (rs.next()) {
-            final T o = newEntityInstance();
-            bindOneFetch(rs, o);
-            r.add(o);
-          }
-          return new ListResultSet<T>(r);
-        } finally {
-          rs.close();
-        }
-      } finally {
         ps.close();
+      } catch (SQLException e) {
+        // Ignored.
       }
-    } catch (SQLException e) {
-      throw convertError("fetch", e);
+      throw convertError("fetch", err);
     }
+    return new JdbcResultSet<T, K>(this, rs, ps);
   }
 
   @Override
@@ -294,7 +291,7 @@ public abstract class JdbcAccess<T, K extends Key<?>> extends
     }
   }
 
-  private OrmException convertError(final String op, final SQLException err) {
+  protected OrmException convertError(final String op, final SQLException err) {
     if (err.getCause() == null && err.getNextException() != null) {
       err.initCause(err.getNextException());
     }
diff --git a/src/main/java/com/google/gwtorm/jdbc/JdbcResultSet.java b/src/main/java/com/google/gwtorm/jdbc/JdbcResultSet.java
new file mode 100644
index 0000000..f041964
--- /dev/null
+++ b/src/main/java/com/google/gwtorm/jdbc/JdbcResultSet.java
@@ -0,0 +1,100 @@
+// Copyright 2008 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.gwtorm.jdbc;
+
+import com.google.gwtorm.client.Key;
+import com.google.gwtorm.client.OrmRuntimeException;
+import com.google.gwtorm.client.impl.AbstractResultSet;
+
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.NoSuchElementException;
+
+class JdbcResultSet<T, K extends Key<?>> extends AbstractResultSet<T> {
+  private final JdbcAccess<T, K> access;
+  private final ResultSet rs;
+  private final PreparedStatement ps;
+  private Boolean haveRow;
+  private boolean closed;
+
+  JdbcResultSet(JdbcAccess<T, K> jdbcAccess, ResultSet rs, PreparedStatement ps) {
+    this.access = jdbcAccess;
+    this.rs = rs;
+    this.ps = ps;
+    this.haveRow = Boolean.TRUE;
+  }
+
+  @Override
+  protected boolean hasNext() {
+    if (closed) {
+      return false;
+    }
+
+    if (haveRow == null) {
+      try {
+        if (rs.next()) {
+          haveRow = Boolean.TRUE;
+        } else {
+          haveRow = Boolean.FALSE;
+          close();
+        }
+      } catch (SQLException err) {
+        close();
+        throw new OrmRuntimeException(access.convertError("fetch", err));
+      }
+    }
+
+    return haveRow;
+  }
+
+  @Override
+  protected T next() {
+    if (!hasNext()) {
+      throw new NoSuchElementException();
+    }
+
+    final T o = access.newEntityInstance();
+    try {
+      access.bindOneFetch(rs, o);
+    } catch (SQLException err) {
+      close();
+      throw new OrmRuntimeException(access.convertError("fetch", err));
+    }
+
+    haveRow = null;
+    hasNext();
+    return o;
+  }
+
+  @Override
+  public void close() {
+    if (!closed) {
+      closed = true;
+
+      try {
+        rs.close();
+      } catch (SQLException e) {
+        // Ignore
+      }
+
+      try {
+        ps.close();
+      } catch (SQLException e) {
+        // Ignore
+      }
+    }
+  }
+}
diff --git a/src/main/java/com/google/gwtorm/nosql/generic/GenericAccess.java b/src/main/java/com/google/gwtorm/nosql/generic/GenericAccess.java
index 08e03bc..3884a31 100644
--- a/src/main/java/com/google/gwtorm/nosql/generic/GenericAccess.java
+++ b/src/main/java/com/google/gwtorm/nosql/generic/GenericAccess.java
@@ -21,6 +21,7 @@ import com.google.gwtorm.client.OrmConcurrencyException;
 import com.google.gwtorm.client.OrmDuplicateKeyException;
 import com.google.gwtorm.client.OrmException;
 import com.google.gwtorm.client.ResultSet;
+import com.google.gwtorm.client.impl.AbstractResultSet;
 import com.google.gwtorm.client.impl.ListResultSet;
 import com.google.gwtorm.nosql.IndexFunction;
 import com.google.gwtorm.nosql.IndexKeyBuilder;
@@ -31,6 +32,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 
@@ -124,15 +126,28 @@ public abstract class GenericAccess<T, K extends Key<?>> extends
     b.addRaw(toKey);
     toKey = b.toByteArray();
 
-    final ArrayList<T> res = new ArrayList<T>();
-    Iterator<Map.Entry<byte[], byte[]>> i = db.scan(fromKey, toKey, limit);
-    while (i.hasNext()) {
-      byte[] bin = i.next().getValue();
-      T obj = getObjectCodec().decode(bin);
-      res.add(obj);
-      cache().put(primaryKey(obj), bin);
-    }
-    return new ListResultSet<T>(res);
+    final ResultSet<Map.Entry<byte[], byte[]>> rs = db.scan(fromKey, toKey, limit);
+    final Iterator<Map.Entry<byte[], byte[]>> i = rs.iterator();
+
+    return new AbstractResultSet<T>() {
+      @Override
+      protected boolean hasNext() {
+        return i.hasNext();
+      }
+
+      @Override
+      protected T next() {
+        byte[] bin = i.next().getValue();
+        T obj = getObjectCodec().decode(bin);
+        cache().put(primaryKey(obj), bin);
+        return obj;
+      }
+
+      @Override
+      public void close() {
+        rs.close();
+      }
+    };
   }
 
   /**
@@ -173,9 +188,8 @@ public abstract class GenericAccess<T, K extends Key<?>> extends
 
     SCAN: for (;;) {
       int scanned = 0;
-      Iterator<Map.Entry<byte[], byte[]>> i = db.scan(lastKey, toKey, limit);
-      while (i.hasNext()) {
-        final Map.Entry<byte[], byte[]> ent = i.next();
+      ResultSet<Entry<byte[], byte[]>> rs = db.scan(lastKey, toKey, limit);
+      for (Map.Entry<byte[], byte[]> ent : rs) {
         final byte[] idxkey = ent.getKey();
         lastKey = idxkey;
         scanned++;
@@ -211,6 +225,7 @@ public abstract class GenericAccess<T, K extends Key<?>> extends
           cache().put(primaryKey(obj), objData);
           res.add(obj);
           if (limit > 0 && res.size() == limit) {
+            rs.close();
             break SCAN;
           }
         } else {
@@ -223,6 +238,7 @@ public abstract class GenericAccess<T, K extends Key<?>> extends
       // a match, and no further rows would exist.
       //
       if (limit == 0 || scanned < limit) {
+        rs.close();
         break SCAN;
       }
 
diff --git a/src/main/java/com/google/gwtorm/nosql/generic/GenericSchema.java b/src/main/java/com/google/gwtorm/nosql/generic/GenericSchema.java
index a70fdae..37fdde3 100644
--- a/src/main/java/com/google/gwtorm/nosql/generic/GenericSchema.java
+++ b/src/main/java/com/google/gwtorm/nosql/generic/GenericSchema.java
@@ -17,6 +17,7 @@ package com.google.gwtorm.nosql.generic;
 import com.google.gwtorm.client.AtomicUpdate;
 import com.google.gwtorm.client.OrmDuplicateKeyException;
 import com.google.gwtorm.client.OrmException;
+import com.google.gwtorm.client.ResultSet;
 import com.google.gwtorm.client.Schema;
 import com.google.gwtorm.nosql.CounterShard;
 import com.google.gwtorm.nosql.IndexKeyBuilder;
@@ -129,16 +130,21 @@ public abstract class GenericSchema extends NoSqlSchema {
     final byte[] toKey = new byte[key.length + 1];
     System.arraycopy(key, 0, toKey, 0, key.length);
 
-    Iterator<Entry<byte[], byte[]>> i = scan(fromKey, toKey, 2);
-    if (!i.hasNext()) {
-      return null;
-    }
+    ResultSet<Entry<byte[], byte[]>> r = scan(fromKey, toKey, 2);
+    try {
+      Iterator<Entry<byte[], byte[]>> i = r.iterator();
+      if (!i.hasNext()) {
+        return null;
+      }
 
-    byte[] data = i.next().getValue();
-    if (i.hasNext()) {
-      throw new OrmDuplicateKeyException("Unexpected duplicate keys");
+      byte[] data = i.next().getValue();
+      if (i.hasNext()) {
+        throw new OrmDuplicateKeyException("Unexpected duplicate keys");
+      }
+      return data;
+    } finally {
+      r.close();
     }
-    return data;
   }
 
   /**
@@ -160,7 +166,7 @@ public abstract class GenericSchema extends NoSqlSchema {
    *         lazily filled, or filled completely.
    * @throws OrmException an error occurred preventing the scan from completing.
    */
-  public abstract Iterator<Map.Entry<byte[], byte[]>> scan(byte[] fromKey,
+  public abstract ResultSet<Map.Entry<byte[], byte[]>> scan(byte[] fromKey,
       byte[] toKey, int limit) throws OrmException;
 
   /**
diff --git a/src/main/java/com/google/gwtorm/nosql/heap/TreeMapSchema.java b/src/main/java/com/google/gwtorm/nosql/heap/TreeMapSchema.java
index 310b328..4cc1a2a 100644
--- a/src/main/java/com/google/gwtorm/nosql/heap/TreeMapSchema.java
+++ b/src/main/java/com/google/gwtorm/nosql/heap/TreeMapSchema.java
@@ -16,11 +16,12 @@ package com.google.gwtorm.nosql.heap;
 
 import com.google.gwtorm.client.AtomicUpdate;
 import com.google.gwtorm.client.OrmException;
+import com.google.gwtorm.client.ResultSet;
 import com.google.gwtorm.client.Schema;
+import com.google.gwtorm.client.impl.ListResultSet;
 import com.google.gwtorm.nosql.generic.GenericSchema;
 
 import java.util.ArrayList;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -46,8 +47,8 @@ public abstract class TreeMapSchema extends GenericSchema {
   }
 
   @Override
-  public Iterator<Map.Entry<byte[], byte[]>> scan(byte[] fromKey, byte[] toKey,
-      int limit) {
+  public ResultSet<Map.Entry<byte[], byte[]>> scan(byte[] fromKey,
+      byte[] toKey, int limit) {
     db.lock.lock();
     try {
       final List<Map.Entry<byte[], byte[]>> res =
@@ -78,7 +79,7 @@ public abstract class TreeMapSchema extends GenericSchema {
           break;
         }
       }
-      return res.iterator();
+      return new ListResultSet<Entry<byte[],byte[]>>(res);
     } finally {
       db.lock.unlock();
     }