Automatically fetch any additional result sets from a query.

Move the iterator to an internal seq.

src/kuzu/queries.nim

@@ -7,7 +7,12 @@ proc `=destroy`*( query: KuzuQueryResultObj ) =
         kuzu_query_result_destroy( addr query.handle )
 
 
-func getQueryMetadata( query: KuzuQueryResult ) =
+# Forward declarations.
+func hasNextSet( query: KuzuQueryResult ): bool
+func getNextSet( query: KuzuQueryResult ): KuzuQueryResult
+
+
+func getQueryMetadata( query: KuzuQueryResult, getAllQueryResults=false ) =
     ## Find and retain additional data for the query.
     query.num_columns = kuzu_query_result_get_num_columns( addr query.handle )
     query.num_tuples  = kuzu_query_result_get_num_tuples( addr query.handle )
@@ -19,6 +24,12 @@ func getQueryMetadata( query: KuzuQueryResult ) =
     query.execution_time = kuzu_query_summary_get_execution_time( addr summary )
     kuzu_query_summary_destroy( addr summary )
 
+    # Pull any additional query results.
+    query.sets = @[]
+    if getAllQueryResults:
+        while query.hasNextSet:
+            query.sets.add( query.getNextSet )
+
     # Column information.
     query.column_types = @[]
     query.column_names = @[]
@@ -48,13 +59,29 @@ func getQueryMetadata( query: KuzuQueryResult ) =
         kuzu_destroy_string( name )
 
 
+func hasNextSet( query: KuzuQueryResult ): bool =
+    ## Returns +true+ if there are more result sets to be consumed.
+    result = kuzu_query_result_has_next_query_result( addr query.handle )
+
+
+func getNextSet( query: KuzuQueryResult ): KuzuQueryResult =
+    ## Consume and return the next query set result, or raise a KuzuIterationError
+    ## if at the end of sets.
+    result = new KuzuQueryResult
+    if kuzu_query_result_get_next_query_result( addr query.handle, addr result.handle ) == KuzuSuccess:
+        result.valid = true
+        result.getQueryMetadata()
+    else:
+        raise newException( KuzuIterationError, &"Query iteration past end of set." )
+
+
 func query*( conn: KuzuConnection, query: string ): KuzuQueryResult =
     ## Perform a database +query+ and return the result.
     result = new KuzuQueryResult
 
     if kuzu_connection_query( addr conn.handle, query, addr result.handle ) == KuzuSuccess:
         result.valid = true
-        result.getQueryMetadata()
+        result.getQueryMetadata( getAllQueryResults=true )
     else:
         var err = kuzu_query_result_get_error_message( addr result.handle )
         raise newException( KuzuQueryError, &"Error running query: {err}" )
@@ -162,32 +189,6 @@ func rewind*( query: KuzuQueryResult ) =
     kuzu_query_result_reset_iterator( addr query.handle )
 
 
-func hasNextSet*( query: KuzuQueryResult ): bool =
-    ## Returns +true+ if there are more result sets to be consumed.
-    result = kuzu_query_result_has_next_query_result( addr query.handle )
-
-
-# Keeping this private, because it's only safe to call from within
-# an iterator -- overwriting a query variable with the next result
-# goes boom.
-func getNextSet( query: KuzuQueryResult ): KuzuQueryResult =
-    ## Consume and return the next query set result, or raise a KuzuIterationError
-    ## if at the end of sets.
-    result = new KuzuQueryResult
-    if kuzu_query_result_get_next_query_result( addr query.handle, addr result.handle ) == KuzuSuccess:
-        result.valid = true
-        result.getQueryMetadata()
-    else:
-        raise newException( KuzuIterationError, &"Query iteration past end of set." )
-
-
-iterator sets*( query: KuzuQueryResult ): KuzuQueryResult =
-    ## Iterate available query result sets, yielding to the block.
-    while query.hasNextSet:
-        yield query.getNextSet
-    # NOTE: There is no 'rewind' mechanism for result sets.
-
-
 iterator items*( query: KuzuQueryResult ): KuzuFlatTuple =
     ## Iterate available tuples, yielding to the block.
     while query.hasNext:

src/kuzu/types.nim

@@ -21,6 +21,7 @@ type
         execution_time*: cdouble = 0
         column_types*:   seq[ kuzu_data_type_id ]
         column_names*:   seq[ string ]
+        sets*:           seq[ KuzuQueryResult ]
         valid = false
     KuzuQueryResult* = ref KuzuQueryResultObj
 

tests/queries/t_can_contain_multiple_result_sets.nim

@@ -1,7 +1,7 @@
 # vim: set et sta sw=4 ts=4 :
 
 discard """
-output: "Jenny|Lenny\nLenny\nJenny\n"
+output: "a\nb\nc\nd\ne\nf\n"
 """
 
 import kuzu
@@ -9,34 +9,26 @@ import kuzu
 let db = newKuzuDatabase()
 let conn = db.connect
 
-var q = conn.query """
+var q = conn.query( "RETURN 'hi'" )
-CREATE NODE TABLE User(
-    id SERIAL PRIMARY KEY,
-    name STRING
-);
-
-CREATE REL TABLE FOLLOWS(
-   From User To User
-);
-
-MERGE (a:User {name: "Lenny"})-[f:Follows]->(b:User {name: "Jenny"});
-"""
-
-q = conn.query( "MATCH (u:User) RETURN *" )
 
 assert typeOf( q ) is KuzuQueryResult
-assert q.hasNextSet == false
+assert q.sets.len == 0
 
 q = conn.query """
-    MATCH (a:User)<-[f:Follows]-(b:User) RETURN a.name, b.name;
+    RETURN "a";
-    MATCH (u:User) RETURN u.name;
+    RETURN "b";
+    RETURN "c";
+    RETURN "d";
+    RETURN "e";
+    RETURN "f";
 """
 
 assert typeOf( q ) is KuzuQueryResult
-assert q.hasNextSet == true
+assert q.sets.len == 5
 
 echo q.getNext
-for query_result in q.sets:
+for set in q.sets:
-    for row in query_result.items:
+    for row in set.items:
         echo row
 
+