How to extract pairs of tables such that they are joined?

[galenseilis]

I am trying to write a function which takes a string of a SQL query and returns a Networkx graph where the nodes are the table names and the edges connect two table names where they are joined on one or more columns. I want to make sure that it works for simple queries, queries with one or more CTEs, and queries with one or more nested queries.

I attempted this something like 11 months ago so my memory is a little foggy. Below is a script I extracted from the Jupyter notebook I wrote at that time to try to exract the pairs. I remember getting stuck on trying to get all the tables that a given table is joined to. The script accesses one of a collection of example sql scripts. The script currently accesses example 13.

#!/usr/bin/env python
# coding: utf-8

import glob
from typing import NoReturn

import networkx as nx
import sqlglot
from sqlglot import exp
from sqlglot.optimizer.scope import build_scope
from sqlglot.optimizer.qualify import qualify
from sqlglot.expressions import Join, Table, Column

def read_file(file):
    with open(file, "r") as f:
        return f.read()

files = sorted(glob.glob("*.sql"))
str_query = read_file(files[12])
ast = sqlglot.parse_one(str_query)
root = build_scope(qualify(ast))

tables = [
    source.name
    for scope in root.traverse()
    for alias, (node, source) in (scope.selected_sources.items())
    if isinstance(source, exp.Table)
]

for join_exp in root.find_all(exp.Join):
    for j in join_exp.find_all(exp.Table):
        print(j.parent, j)

joins = [
    source.name
    for scope in root.traverse()
    for alias, (node, source) in (scope.selected_sources.items())
    if isinstance(source, exp.Join)
]

print(joins)

I 'feel' like this is something that should be easy to extract if I knew the underlying implementation better. Of course, eventually, I will resort to carefully doing static and dynamic analysis of SQLGlot if it comes to it, but I am hoping not to get into that.

It isn't evident to me how to proceed from AST Introspection which appears to cover only a minimal example. I think more topical guides in the docs would be useful; especially if they followed Diátaxis.

Suggestions and feedback would be appreciated. 🙏

---

I should add, I did ask on Slack way back at the time that I was first looking at this. I recall only getting some partial information, perhaps how to get table names. However, I cannot see those messages anymore.

[georgesittas]

I'd suggest taking a look at these two docs, you may find useful patterns:

• https://github.com/tobymao/sqlglot/blob/main/posts/ast_primer.md
• https://github.com/tobymao/sqlglot/blob/main/posts/onboarding.md

There's quite a bit of info in this post and it's not very clear which part you're stuck on exactly. Can you reduce it to a simple example, demonstrating the initial & expected state of things? That way it'll be easier to point you to places to look.

[galenseilis]

I'd suggest taking a look at these two docs, you may find useful patterns:

• https://github.com/tobymao/sqlglot/blob/main/posts/ast_primer.md
• https://github.com/tobymao/sqlglot/blob/main/posts/onboarding.md

Thanks! I will review those materials.

There's quite a bit of info in this post and it's not very clear which part you're stuck on exactly. Can you reduce it to a simple example, demonstrating the initial & expected state of things? That way it'll be easier to point you to places to look.

Thanks for this feedback. I will come up with a simpler example that demonstrates the initial and expected states.

[galenseilis]

@georgesittas Thank you again for the documents. After that reading I still don't know why my example above does not extract the collection of pairs of tables that are joined.

Let's consider this example:

# bar.py
import sqlglot
from sqlglot import exp
from sqlglot.optimizer.scope import build_scope
from sqlglot.expressions import Join


str_query = """
WITH MonthlySales AS (
    SELECT
        customer_id,
        MONTH(order_date) AS order_month,
        SUM(order_amount) AS total_sales
    FROM orders
    GROUP BY customer_id, order_month
)

SELECT
    e.employee_id,
    e.first_name,
    e.last_name,
    d.department_name,
    ms.order_month,
    ms.total_sales
FROM employees AS e
JOIN departments AS d ON e.department_id = d.department_id
JOIN (
    SELECT
        c.customer_id,
        MONTH(o.order_date) AS order_month,
        SUM(o.order_amount) AS total_sales
    FROM customers AS c
    JOIN orders AS o ON c.customer_id = o.customer_id
    WHERE c.country = 'USA'
    GROUP BY c.customer_id, order_month
) AS ms ON e.employee_id = ms.customer_id
WHERE ms.total_sales > 1000
ORDER BY e.employee_id, ms.order_month;
"""

ast = sqlglot.parse_one(str_query)
root = build_scope(ast)

joins = [
    source.name
    for scope in root.traverse()
    for alias, (node, source) in scope.selected_sources.items()
    if isinstance(source, exp.Join)
]

print(joins)

We can run this example:

> python bar.py
[]

There are joins in the SQL query, but they are not showing up. I don't expect this example to have the pairs of tables/queries participating in the joins, but I did expect that the joins would show up in some way.

[galenseilis]

It is unclear to me why this discussion has been closed. The original question has not been answered and the thread has not been abandoned for any substantial length of time.