Add check for subquery ref and add tests

cwida · Jan 19, 2025 · 80dcaa4 · 80dcaa4
1 parent 4697869
commit 80dcaa4
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/src/core/parser/duckpgq_parser.cpp b/src/core/parser/duckpgq_parser.cpp
@@ -15,6 +15,7 @@
 #include <duckpgq/core/functions/table/describe_property_graph.hpp>
 #include <duckpgq/core/functions/table/drop_property_graph.hpp>
 #include "duckdb/parser/query_node/cte_node.hpp"
+#include "duckdb/parser/tableref/subqueryref.hpp"
 
 namespace duckpgq {
 
@@ -34,6 +35,9 @@ ParserExtensionParseResult duckpgq_parse(ParserExtensionInfo *info,
           std::move(parser.statements[0])));
 }
 
+ParserExtensionPlanResult
+duckpgq_handle_statement(SQLStatement *statement, DuckPGQState &duckpgq_state);
+
 void duckpgq_find_match_function(TableRef *table_ref,
                                  DuckPGQState &duckpgq_state) {
   if (auto table_function_ref = dynamic_cast<TableFunctionRef *>(table_ref)) {
@@ -54,6 +58,9 @@ void duckpgq_find_match_function(TableRef *table_ref,
     // Handle JoinRef case
     duckpgq_find_match_function(join_ref->left.get(), duckpgq_state);
     duckpgq_find_match_function(join_ref->right.get(), duckpgq_state);
+  } else if (auto subquery_ref = dynamic_cast<SubqueryRef *>(table_ref)) {
+    // Handle SubqueryRef case
+    duckpgq_handle_statement(subquery_ref->subquery.get(), duckpgq_state);
   }
 }
 

diff --git a/test/sql/nested_subquery.test b/test/sql/nested_subquery.test
@@ -0,0 +1,32 @@
+
+
+require duckpgq
+
+statement ok
+CREATE TABLE Student(id BIGINT, name VARCHAR);INSERT INTO Student VALUES (0, 'Daniel'), (1, 'Tavneet'), (2, 'Gabor'), (3, 'Peter'), (4, 'David');
+
+statement ok
+CREATE TABLE know(src BIGINT, dst BIGINT, createDate BIGINT);INSERT INTO know VALUES (0,1, 10), (0,2, 11), (0,3, 12), (3,0, 13), (1,2, 14), (1,3, 15), (2,3, 16), (4,3, 17);
+
+statement ok
+CREATE TABLE School(name VARCHAR, Id BIGINT, Kind VARCHAR);INSERT INTO School VALUES ('VU', 0, 'University'), ('UVA', 1, 'University');
+
+statement ok
+CREATE TABLE StudyAt(personId BIGINT, schoolId BIGINT);INSERT INTO StudyAt VALUES (0, 0), (1, 0), (2, 1), (3, 1), (4, 1);
+
+statement ok
+-CREATE PROPERTY GRAPH pg
+VERTEX TABLES (
+    Student PROPERTIES ( id, name ) LABEL Person,
+    School LABEL SCHOOL
+    )
+EDGE TABLES (
+    know    SOURCE KEY ( src ) REFERENCES Student ( id )
+            DESTINATION KEY ( dst ) REFERENCES Student ( id ),
+    studyAt SOURCE KEY ( personId ) REFERENCES Student ( id )
+            DESTINATION KEY ( SchoolId ) REFERENCES School ( id )
+    );
+
+
+statement ok
+-select * from (select id, id_1 from graph_table ( pg match (p:person)-[k:know]->(p2:person) columns (p.id, p2.id)));