Repository: tobymao/sqlglot
Pull request: fix(optimizer): Enable USING expansion with multiple joins
fixes tobymao/sqlglot#4112
also added test for this case

Problem statement:
`optimizer.qualify.qualify` raises "Cannot automatically join" for multiple joins with `USING` and `ON` clauses
When there are two or more `JOIN` operations, and a `JOIN` with an `ON` condition precedes a `JOIN` with a `USING` clause, the `optimizer.qualify.qualify` method throws the error: **"Cannot automatically join"**.


## Steps to Reproduce

```python
from sqlglot import parse_one
from sqlglot.optimizer.qualify import qualify

schema = {
    'A': {'b_id': 'int'}, 
    'B': {'b_id': 'int', 'c_id': 'int'}, 
    'C': {'c_id': 'int'}
}

sql = "SELECT A.b_id FROM A JOIN B ON A.b_id=B.b_id JOIN C USING(c_id)"
print(qualify(parse_one(sql, read='postgres'), schema=schema, quote_identifiers=False).sql(dialect='postgres'))
```

### Current Behavior
The code throws the following error:

```python
sqlglot.errors.OptimizeError: Cannot automatically join: c_id
```

### Expected Behavior
The optimizer should properly qualify the columns and return the following SQL:

```sql
SELECT a.b_id AS b_id FROM a AS a JOIN b AS b ON a.b_id = b.b_id JOIN c AS c ON b.c_id = c.c_id
```


## Additional Context
The error seems to originate from the `_expand_using` method's handling of `JOIN` clauses. It misinterprets the mix of `ON` and `USING` join strategies, leading to failure.

Hints from issue discussion:

Commit message: Fix for #4112 optimizer.qualify._expand_using now correctly process multiple joins with USING and ON clauses

Patch to review:
diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py
index c7f3cda0a4..e1bdd6b5bb 100644
--- a/sqlglot/optimizer/qualify_columns.py
+++ b/sqlglot/optimizer/qualify_columns.py
@@ -155,9 +155,6 @@ def _expand_using(scope: Scope, resolver: Resolver) -> t.Dict[str, t.Any]:
     for i, join in enumerate(joins):
         using = join.args.get("using")
 
-        if not using:
-            continue
-
         join_table = join.alias_or_name
 
         columns = {}
@@ -170,6 +167,10 @@ def _expand_using(scope: Scope, resolver: Resolver) -> t.Dict[str, t.Any]:
 
         source_table = ordered[-1]
         ordered.append(join_table)
+
+        if not using:
+            continue
+
         join_columns = resolver.get_source_columns(join_table)
         conditions = []
         using_identifier_count = len(using)
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index fe5a4d7c54..a71644a39c 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -397,6 +397,22 @@ def test_qualify_columns(self, logger):
             "SELECT u.user_id AS user_id, l.log_date AS log_date FROM users AS u CROSS JOIN LATERAL (SELECT l1.log_date AS log_date FROM (SELECT l.log_date AS log_date FROM logs AS l WHERE l.user_id = u.user_id AND l.log_date <= 100 ORDER BY l.log_date LIMIT 1) AS l1) AS l",
         )
 
+        self.assertEqual(
+            optimizer.qualify.qualify(
+                parse_one(
+                    "SELECT A.b_id FROM A JOIN B ON A.b_id=B.b_id JOIN C USING(c_id)",
+                    dialect="postgres",
+                ),
+                schema={
+                    "A": {"b_id": "int"},
+                    "B": {"b_id": "int", "c_id": "int"},
+                    "C": {"c_id": "int"},
+                },
+                quote_identifiers=False,
+            ).sql("postgres"),
+            "SELECT a.b_id AS b_id FROM a AS a JOIN b AS b ON a.b_id = b.b_id JOIN c AS c ON b.c_id = c.c_id",
+        )
+
         self.check_file(
             "qualify_columns",
             qualify_columns,


Code review comments:
diff_hunk: @@ -170,6 +167,10 @@ def _expand_using(scope: Scope, resolver: Resolver) -> t.Dict[str, t.Any]:
 
         source_table = ordered[-1]
         ordered.append(join_table)
+
+        if not using:
+            continue
+
path: sqlglot/optimizer/qualify_columns.py
text: @user1:
Moving this branch down is enough to solve the bug (idea being `ordered` should contain all join tables regardless of a `USING` clause if I'm not mistaken) however we don't have compute `columns` if that's the case; The order of operations could be:

```
for i, join in enumerate(join):
 join_table = join.alias_or_name
 ordered.append(join_table)
 
  using = join.args.get("using")
  if not using:
    continue

 ...
```

As a sidenote, that `columns` nested for-loop looks like a code smell, I'm working on a refactor that can be included in this PR / right after it once it's merged.  

@author:
You are welcome.
I've just wanted to provide the simplest way to fix it. 

Unfortunately, your suggestion will not work as we add `join_table` to `ordered` and it looks like it influences to `columns` dict and some tests will fail, for instance:

```
======================================================================
FAIL: test_qualify_columns (tests.test_optimizer.TestOptimizer.test_qualify_columns) [136, SELECT b FROM t JOIN x USING(a)]
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/dg/projects/test/sqlglot/tests/test_optimizer.py", line 167, in check_file
    self.assertEqual(
AssertionError: 'SELECT x.b AS b FROM t AS t JOIN x AS x ON t.a = x.a' != 'SELECT x.b AS b FROM t AS t JOIN x AS x ON x.a = x.a'
- SELECT x.b AS b FROM t AS t JOIN x AS x ON t.a = x.a
?                                            ^
+ SELECT x.b AS b FROM t AS t JOIN x AS x ON x.a = x.a
?     
```

What will work is code like this:
```python
    for i, join in enumerate(joins):
        columns = {}

        for source_name in scope.selected_sources:
            if source_name in ordered:
                for column_name in resolver.get_source_columns(source_name):
                    if column_name not in columns:
                        columns[column_name] = source_name

        source_table = ordered[-1]
        using = join.args.get("using")
        join_table = join.alias_or_name
        ordered.append(join_table)

        if not using:
            continue
```

so it would be a great idea to refactor code with `columns`

Which solution you will prefer, already committed or proposed?

@user1:
Right, we should also hoist up the line `source_table = ordered[-1]` and swap `ordered` for `ordered[:-1]` in the nested for-loop, that should restore the functionality while skipping the (potentially) unnecessary computations.

My refactor regarding `columns` has to do with the observation that we append one item at a time to `ordered` but loop through it again on each iteration, so we can instead construct it incrementally by moving `columns = {}` outside of the loop and checking `scope.selected_sources` with the latest entry in ordered afaict; Locally all tests are passing for a quick implementation of this approach, so we can revisit this once this PR is approved.