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