diff --git a/src/kuzu/queries.nim b/src/kuzu/queries.nim index 03fe293..71058bb 100644 --- a/src/kuzu/queries.nim +++ b/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: diff --git a/src/kuzu/types.nim b/src/kuzu/types.nim index f5026eb..dc2dbb9 100644 --- a/src/kuzu/types.nim +++ b/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 diff --git a/tests/queries/t_can_contain_multiple_result_sets.nim b/tests/queries/t_can_contain_multiple_result_sets.nim index d0be2d5..534f871 100644 --- a/tests/queries/t_can_contain_multiple_result_sets.nim +++ b/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 """ -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 *" ) +var q = conn.query( "RETURN 'hi'" ) 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; - MATCH (u:User) RETURN u.name; + RETURN "a"; + 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 row in query_result.items: +for set in q.sets: + for row in set.items: echo row +