Add support for ORDER/GROUP BY, HAVING, and booleans in WHERE.

lib/pg_query/parse.rb

@@ -117,11 +117,16 @@ def load_tables_and_aliases! # rubocop:disable Metrics/CyclomaticComplexity
           end
         end
 
-        subselect_items += statement.values[0]['targetList'] if !statement.empty? && statement.values[0]['targetList']
-        subselect_items << statement.values[0]['whereClause'] if !statement.empty? && statement.values[0]['whereClause']
+        statement_value = statement.values[0]
+        unless statement.empty?
+          subselect_items.concat(statement_value['targetList']) if statement_value['targetList']
+          subselect_items << statement_value['whereClause'] if statement_value['whereClause']
+          subselect_items.concat(statement_value['sortClause'].collect { |h| h[SORT_BY]['node'] }) if statement_value['sortClause']
+          subselect_items.concat(statement_value['groupClause']) if statement_value['groupClause']
+          subselect_items << statement_value['havingClause'] if statement_value['havingClause']
+        end
       end
 
-      # Find subselects in WHERE clause
       next_item = subselect_items.shift
       if next_item
         case next_item.keys[0]
@@ -135,6 +140,8 @@ def load_tables_and_aliases! # rubocop:disable Metrics/CyclomaticComplexity
               subselect_items << elem
             end
           end
+        when BOOL_EXPR
+          subselect_items.concat(next_item.values[0]['args'])
         when RES_TARGET
           subselect_items << next_item[RES_TARGET]['val']
         when SUB_LINK

spec/lib/parse_spec.rb

@@ -614,7 +614,7 @@
   end
 
   # https://github.com/lfittl/pg_query/issues/38
-  it 'correctly finds nested tables' do
+  it 'correctly finds nested tables in select clause' do
     query = described_class.parse("select u.email, (select count(*) from enrollments e where e.user_id = u.id) as num_enrollments from users u")
     expect(query.warnings).to eq []
     expect(query.tables).to eq ['users', 'enrollments']
@@ -627,6 +627,125 @@
     expect(query.tables).to eq ['table_name']
   end
 
+  it 'correctly finds nested tables in from clause' do
+    query = described_class.parse("select u.* from (select * from users) u")
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users']
+  end
+
+  it 'correctly finds nested tables in where clause' do
+    query = described_class.parse("select users.id from users where 1 = (select count(*) from user_roles)")
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles']
+  end
+
+  it 'traverse boolean expressions in where clause' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      where users.id IN (
+        select user_roles.user_id
+        from user_roles
+      ) and (users.created_at between '2016-06-01' and '2016-06-30')
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles']
+  end
+
+  it 'correctly finds nested tables in the order by clause' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      order by (
+        select max(user_roles.role_id)
+        from user_roles
+        where user_roles.user_id = users.id
+      )
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles']
+  end
+
+  it 'correctly finds nested tables in the order by clause with multiple entries' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      order by (
+        select max(user_roles.role_id)
+        from user_roles
+        where user_roles.user_id = users.id
+      ) asc, (
+        select max(user_logins.role_id)
+        from user_logins
+        where user_logins.user_id = users.id
+      ) desc
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles', 'user_logins']
+  end
+
+  it 'correctly finds nested tables in the group by clause' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      group by (
+        select max(user_roles.role_id)
+        from user_roles
+        where user_roles.user_id = users.id
+      )
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles']
+  end
+
+  it 'correctly finds nested tables in the group by clause with multiple entries' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      group by (
+        select max(user_roles.role_id)
+        from user_roles
+        where user_roles.user_id = users.id
+      ), (
+        select max(user_logins.role_id)
+        from user_logins
+        where user_logins.user_id = users.id
+      )
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles', 'user_logins']
+  end
+
+  it 'correctly finds nested tables in the having clause' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      group by users.id
+      having 1 > (
+        select count(user_roles.role_id)
+        from user_roles
+        where user_roles.user_id = users.id
+      )
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles']
+  end
+
+  it 'correctly finds nested tables in the having clause with a boolean expression' do
+    query = described_class.parse(<<-SQL)
+      select users.*
+      from users
+      group by users.id
+      having true and 1 > (
+        select count(user_roles.role_id)
+        from user_roles
+        where user_roles.user_id = users.id
+      )
+    SQL
+    expect(query.warnings).to eq []
+    expect(query.tables).to eq ['users', 'user_roles']
+  end
+
   it 'handles DROP TYPE' do
     query = described_class.parse("DROP TYPE IF EXISTS repack.pk_something")
     expect(query.warnings).to eq []