Repository: marimo-team/marimo
Pull request: improvement: parse sql CREATE table names
Fixes https://github.com/marimo-team/marimo/issues/2019

(thanks @akshayka for finishing this PR).

This adds "refs" and "defs" between SQL <-> SQL statements and Python -> SQL. (python leaks into sql's namepsace, but not vice-versa).

We define "definitions" as:
* `CREATE` for tables and views
* `ATTACH` for DBs

We draw references when:
* A `SELECT` contains a sql table, schema, or python df

Reference detection for schemas is incomplete but hopefully good enough for many use cases.

Problem statement:
Track database objects in  dataflow graph
### Description

I really like the SQL integration.
However, when database objects like tables and views are created in SQL cells, other cells that depend on these might not be executed in the correct order.

```python
import marimo

__generated_with = "0.7.20"
app = marimo.App(width="medium")


@app.cell
def __(mo, my_table):
    _df = mo.sql(
        f"""
        SELECT * FROM my_table
        """
    )
    return


@app.cell
def __(mo, my_view):
    _df = mo.sql(
        f"""
        SELECT * FROM my_view
        """
    )
    return


@app.cell
def __(mo):
    _df = mo.sql(
        f"""
        CREATE OR REPlACE TABLE my_table (a text);
        INSERT INTO my_table (VALUES ('foo'), ('bar'));
        """
    )
    return


@app.cell
def __(mo):
    _df = mo.sql(
        f"""
        CREATE VIEW my_view AS (SELECT * FROM my_table WHERE a LIKE 'f%o')
        """
    )
    return


@app.cell(hide_code=True)
def __():
    import marimo as mo
    import pandas as pd
    return mo, pd


if __name__ == "__main__":
    app.run()
```

### Suggested solution

Include the created tables, views, etc and their references in the dataflow graph.

### Alternative

As a workaround users can track the dependencies through python objects:
```python
import marimo

__generated_with = "0.7.20"
app = marimo.App(width="medium")


@app.cell
def __(db_my_table, mo, my_table):
    db_my_table

    _df = mo.sql(
        f"""
        SELECT * FROM my_table
        """
    )
    return


@app.cell
def __(db_my_view, mo, my_view):
    db_my_view

    _df = mo.sql(
        f"""
        SELECT * FROM my_view
        """
    )
    return


@app.cell
def __(mo):
    db_my_table = None
    _df = mo.sql(
        f"""
        CREATE OR REPlACE TABLE my_table (a text);
        INSERT INTO my_table (VALUES ('foo'), ('bar'));
        """
    )
    return db_my_table,


@app.cell
def __(db_my_table, mo):
    db_my_table
    db_my_view = None
    _df = mo.sql(
        f"""
        CREATE VIEW my_view AS (SELECT * FROM my_table WHERE a LIKE 'f%o')
        """
    )
    return db_my_view,


@app.cell(hide_code=True)
def __():
    import marimo as mo
    import pandas as pd
    return mo, pd


if __name__ == "__main__":
    app.run()
```
For this pure SQL cells can't be used.
Even worse, this leaves the user in charge of making sure that all database objects are tracked correctly. If one is missed the reproducibility of the notebook relies on the placement of the cells (just like in jupyter notebooks).

### Additional context

_No response_

Hints from issue discussion:
Extracting references from the SQL is already done, so this should work for handling workflow without the verbosity

```python
    my_view = mo.sql(
        f"""
        CREATE VIEW my_view AS (SELECT * FROM my_table WHERE a LIKE 'f%o')
        """
    )
```

```python
    my_table = mo.sql(
        f"""
        CREATE OR REPlACE TABLE my_table (a text);
        INSERT INTO my_table (VALUES ('foo'), ('bar'));
        """
    )
```

But `my_view` and `my_table` will be None. It would be nice if the creation statements created dataframes as well.
So if `CREATE OR REPlACE TABLE my_table (a text)` is a standalone statement, the variable `my_table` is auto set if the reference is not private- to enforce that the sql table and variable matched.


Thanks for reporting. I have it on my TODO list to look into parsing SQL definitions and references more thoroughly, perhaps using `sqlglot`, and integrating them into the dataflow graph. 
So just a note, duckdb has a native `duckdb.tokenize()`, and the grammar is pretty fixed: https://github.com/duckdb/duckdb/blob/main/third_party/libpg_query/grammar/statements/create.y so could directly leverage that
> So just a note, duckdb has a native `duckdb.tokenize()`, and the grammar is pretty fixed: https://github.com/duckdb/duckdb/blob/main/third_party/libpg_query/grammar/statements/create.y so could directly leverage that

Oh nice, even better. I'll look into that.

Commit message: docs

Patch to review:
diff --git a/marimo/_ast/cell.py b/marimo/_ast/cell.py
index 95c99346c6f..f475860d9be 100644
--- a/marimo/_ast/cell.py
+++ b/marimo/_ast/cell.py
@@ -6,8 +6,8 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Literal, Mapping, Optional
 
+from marimo._ast.sql_visitor import SQLVisitor
 from marimo._ast.visitor import ImportData, Name, VariableData
-from marimo._data.sql_visitor import SQLVisitor
 from marimo._utils.deep_merge import deep_merge
 
 CellId_t = str
diff --git a/marimo/_ast/sql_visitor.py b/marimo/_ast/sql_visitor.py
new file mode 100644
index 00000000000..b1001163faa
--- /dev/null
+++ b/marimo/_ast/sql_visitor.py
@@ -0,0 +1,160 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import re
+from typing import Any, Optional
+
+from marimo._dependencies.dependencies import DependencyManager
+
+
+class SQLVisitor(ast.NodeVisitor):
+    """
+    Find any SQL queries in the AST.
+    This should be inside a function called `.execute` or `.sql`.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._sqls: list[str] = []
+
+    def visit_Call(self, node: ast.Call) -> None:
+        # Check if the call is a method call and the method is named
+        # either 'execute' or 'sql'
+        if isinstance(node.func, ast.Attribute) and node.func.attr in (
+            "execute",
+            "sql",
+        ):
+            # Check if there are arguments and the first argument is a
+            # string or f-string
+            if node.args:
+                first_arg = node.args[0]
+                sql: Optional[str] = None
+                if isinstance(first_arg, ast.Constant):
+                    sql = first_arg.s
+                elif isinstance(first_arg, ast.JoinedStr):
+                    sql = normalize_sql_f_string(first_arg)
+
+                if sql is not None:
+                    # Append the SQL query to the list
+                    self._sqls.append(sql)
+        # Continue walking through the AST
+        self.generic_visit(node)
+
+    def get_sqls(self) -> list[str]:
+        return self._sqls
+
+
+def normalize_sql_f_string(node: ast.JoinedStr) -> str:
+    """
+    Normalize a f-string to a string by joining the parts.
+
+    We add placeholder for {...} expressions in the f-string.
+    This is so we can create a valid SQL query to be passed to
+    other utilities.
+    """
+
+    def print_part(part: ast.expr) -> str:
+        if isinstance(part, ast.FormattedValue):
+            return print_part(part.value)
+        elif isinstance(part, ast.JoinedStr):
+            return normalize_sql_f_string(part)
+        elif isinstance(part, ast.Constant):
+            return str(part.s)
+        else:
+            # Just add '_' as a placeholder for {...} expressions
+            return "'_'"
+
+    result = "".join(print_part(part) for part in node.values)
+    # remove any double '' created by the f-string
+    return result.replace("''", "'")
+
+
+def find_created_tables(sql_statement: str) -> list[str]:
+    """
+    Find the tables created in a SQL statement.
+
+    This function uses the DuckDB tokenizer to find the tables created
+    in a SQL statement. It returns a list of the table names created
+    in the statement.
+
+    Args:
+        sql_statement: The SQL statement to parse.
+
+    Returns:
+        A list of the table names created in the statement.
+    """
+    if not DependencyManager.duckdb.has():
+        return []
+
+    import duckdb
+
+    def remove_comments(sql: str) -> str:
+        # Function to replace comments with spaces, preserving newlines
+        def replace_with_spaces(match: re.Match[Any]) -> str:
+            return " " * len(match.group())
+
+        # Split the SQL into strings and non-strings
+        parts = re.split(r'(\'(?:\'\'|[^\'])*\'|"(?:""|[^"])*")', sql)
+
+        for i in range(0, len(parts), 2):
+            # Remove single-line comments
+            parts[i] = re.sub(
+                r"--.*$", replace_with_spaces, parts[i], flags=re.MULTILINE
+            )
+
+            # Remove multi-line comments
+            parts[i] = re.sub(r"/\*[\s\S]*?\*/", replace_with_spaces, parts[i])
+
+        # Join the parts back together
+        return "".join(parts)
+
+    sql_statement = remove_comments(sql_statement)
+
+    tokens = duckdb.tokenize(sql_statement)
+    created_tables: list[str] = []
+    i = 0
+
+    def token_str(i: int) -> str:
+        token = tokens[i]
+        start = token[0]
+        end = len(sql_statement) - 1
+        if i + 1 < len(tokens):
+            end = tokens[i + 1][0]
+        return sql_statement[start:end].strip()
+
+    def keyword_token_str(i: int) -> str:
+        return token_str(i).lower()
+
+    def token_type(i: int) -> str:
+        return tokens[i][1]
+
+    # See https://duckdb.org/docs/sql/statements/create_table#syntax
+    # for more information on the CREATE TABLE syntax.
+    while i < len(tokens):
+        if (
+            keyword_token_str(i) == "create"
+            and token_type(i) == duckdb.token_type.keyword
+        ):
+            i += 1
+            if i < len(tokens) and keyword_token_str(i) == "or":
+                i += 2  # Skip 'OR REPLACE'
+            if i < len(tokens) and keyword_token_str(i) in (
+                "temporary",
+                "temp",
+            ):
+                i += 1  # Skip 'TEMPORARY' or 'TEMP'
+            if i < len(tokens) and keyword_token_str(i) == "table":
+                i += 1
+                if i < len(tokens) and keyword_token_str(i) == "if":
+                    i += 3  # Skip 'IF NOT EXISTS'
+                if i < len(tokens):
+                    table_name = token_str(i)
+                    # Remove quotes if present
+                    if table_name.startswith('"') and table_name.endswith('"'):
+                        table_name = table_name[1:-1]
+
+                    created_tables.append(table_name)
+        i += 1
+
+    return created_tables
diff --git a/marimo/_ast/visitor.py b/marimo/_ast/visitor.py
index fb68326a84e..9da29c6ab22 100644
--- a/marimo/_ast/visitor.py
+++ b/marimo/_ast/visitor.py
@@ -9,7 +9,7 @@
 from uuid import uuid4
 
 from marimo import _loggers
-from marimo._data.sql_visitor import normalize_sql_f_string
+from marimo._ast.sql_visitor import find_created_tables, normalize_sql_f_string
 from marimo._dependencies.dependencies import DependencyManager
 from marimo._utils.variables import is_local
 
@@ -401,11 +401,11 @@ def visit_Call(self, node: ast.Call) -> None:
                         tables = duckdb.get_table_names(statement.query)
                     except duckdb.ProgrammingError:
                         self.generic_visit(node)
-                        return
+                        continue
                     except BaseException as e:
                         LOGGER.warning("Unexpected duckdb error %s", e)
                         self.generic_visit(node)
-                        return
+                        continue
 
                     for table in tables:
                         # Table may be a URL or something else that
@@ -413,6 +413,22 @@ def visit_Call(self, node: ast.Call) -> None:
                         if table.isidentifier():
                             self._add_ref(table, deleted=False)
 
+                    # Add all tables created in the query to the defs
+                    try:
+                        created_tables = find_created_tables(sql)
+                    except duckdb.ProgrammingError:
+                        self.generic_visit(node)
+                        continue
+                    except BaseException as e:
+                        LOGGER.warning("Unexpected duckdb error %s", e)
+                        self.generic_visit(node)
+                        continue
+
+                    for _table in created_tables:
+                        # TODO:
+                        # self._define_sql()
+                        pass
+
         # Visit arguments, keyword args, etc.
         self.generic_visit(node)
 
diff --git a/marimo/_data/sql_visitor.py b/marimo/_data/sql_visitor.py
deleted file mode 100644
index b85794a272c..00000000000
--- a/marimo/_data/sql_visitor.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# Copyright 2024 Marimo. All rights reserved.
-from __future__ import annotations
-
-import ast
-from typing import Optional
-
-
-class SQLVisitor(ast.NodeVisitor):
-    """
-    Find any SQL queries in the AST.
-    This should be inside a function called `.execute` or `.sql`.
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-        self._sqls: list[str] = []
-
-    def visit_Call(self, node: ast.Call) -> None:
-        # Check if the call is a method call and the method is named
-        # either 'execute' or 'sql'
-        if isinstance(node.func, ast.Attribute) and node.func.attr in (
-            "execute",
-            "sql",
-        ):
-            # Check if there are arguments and the first argument is a
-            # string or f-string
-            if node.args:
-                first_arg = node.args[0]
-                sql: Optional[str] = None
-                if isinstance(first_arg, ast.Constant):
-                    sql = first_arg.s
-                elif isinstance(first_arg, ast.JoinedStr):
-                    sql = normalize_sql_f_string(first_arg)
-
-                if sql is not None:
-                    # Append the SQL query to the list
-                    self._sqls.append(sql)
-        # Continue walking through the AST
-        self.generic_visit(node)
-
-    def get_sqls(self) -> list[str]:
-        return self._sqls
-
-
-def normalize_sql_f_string(node: ast.JoinedStr) -> str:
-    """
-    Normalize a f-string to a string by joining the parts.
-
-    We add placeholder for {...} expressions in the f-string.
-    This is so we can create a valid SQL query to be passed to
-    other utilities.
-    """
-
-    def print_part(part: ast.expr) -> str:
-        if isinstance(part, ast.FormattedValue):
-            return print_part(part.value)
-        elif isinstance(part, ast.JoinedStr):
-            return normalize_sql_f_string(part)
-        elif isinstance(part, ast.Constant):
-            return str(part.s)
-        else:
-            # Just add '_' as a placeholder for {...} expressions
-            return "'_'"
-
-    result = "".join(print_part(part) for part in node.values)
-    # remove any double '' created by the f-string
-    return result.replace("''", "'")
diff --git a/marimo/_smoke_tests/sql/duckdb_tokenize.py b/marimo/_smoke_tests/sql/duckdb_tokenize.py
new file mode 100644
index 00000000000..9b793e16e9d
--- /dev/null
+++ b/marimo/_smoke_tests/sql/duckdb_tokenize.py
@@ -0,0 +1,129 @@
+import marimo
+
+__generated_with = "0.7.20"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def __():
+    import marimo as mo
+    import duckdb
+    import re
+    return duckdb, mo, re
+
+
+@app.cell
+def __(mo):
+    default_code = """
+    CREATE TABLE
+            -- This is a comment
+            IF NOT EXISTS
+            -- This is another comment
+            table1 (id INT)
+    """
+    code = mo.ui.code_editor(language="sql", value=default_code.strip())
+    code
+    return code, default_code
+
+
+@app.cell
+def __(code, duckdb):
+    statements = duckdb.extract_statements(code.value)
+    print("Number of statements: " + str(len(statements)))
+
+    print(duckdb.tokenize(code.value))
+    return statements,
+
+
+@app.cell
+def __(code, find_created_tables):
+    find_created_tables(code.value)
+    return
+
+
+@app.cell
+def __(re):
+    def find_created_tables(sql_statement: str) -> list[str]:
+        """
+        Find the tables created in a SQL statement.
+
+        This function uses the DuckDB tokenizer to find the tables created
+        in a SQL statement. It returns a list of the table names created
+        in the statement.
+
+        Args:
+            sql_statement: The SQL statement to parse.
+
+        Returns:
+            A list of the table names created in the statement.
+        """
+
+        import duckdb
+
+        def remove_comments(sql):
+            # Function to replace comments with spaces, preserving newlines
+            def replace_with_spaces(match):
+                return " " * len(match.group())
+
+            # Split the SQL into strings and non-strings
+            parts = re.split(r'(\'(?:\'\'|[^\'])*\'|"(?:""|[^"])*")', sql)
+
+            for i in range(0, len(parts), 2):
+                # Remove single-line comments
+                parts[i] = re.sub(
+                    r"--.*$", replace_with_spaces, parts[i], flags=re.MULTILINE
+                )
+
+                # Remove multi-line comments
+                parts[i] = re.sub(r"/\*[\s\S]*?\*/", replace_with_spaces, parts[i])
+
+            # Join the parts back together
+            return "".join(parts)
+
+        sql_statement = remove_comments(sql_statement)
+
+        tokens = duckdb.tokenize(sql_statement)
+        created_tables: list[str] = []
+        i = 0
+
+        def token_str(i: int) -> str:
+            token = tokens[i]
+            start = token[0]
+            end = len(sql_statement) - 1
+            if i + 1 < len(tokens):
+                end = tokens[i + 1][0]
+            return sql_statement[start:end].strip()
+
+        def keyword_token_str(i: int) -> str:
+            return token_str(i).lower()
+
+        def token_type(i: int) -> str:
+            return tokens[i][1]
+
+        while i < len(tokens):
+            if (
+                keyword_token_str(i) == "create"
+                and token_type(i) == duckdb.token_type.keyword
+            ):
+                i += 1
+                if i < len(tokens) and keyword_token_str(i) == "or":
+                    i += 2  # Skip 'OR REPLACE'
+                if i < len(tokens) and keyword_token_str(i) in (
+                    "temporary",
+                    "temp",
+                ):
+                    i += 1  # Skip 'TEMPORARY' or 'TEMP'
+                if i < len(tokens) and keyword_token_str(i) == "table":
+                    i += 1
+                    if i < len(tokens) and keyword_token_str(i) == "if":
+                        i += 3  # Skip 'IF NOT EXISTS'
+                    if i < len(tokens):
+                        created_tables.append(token_str(i))
+            i += 1
+
+        return created_tables
+    return find_created_tables,
+
+
+if __name__ == "__main__":
+    app.run()
diff --git a/marimo/_smoke_tests/table_urls.py b/marimo/_smoke_tests/table_urls.py
index f29fe514ea3..9fb8102d17f 100644
--- a/marimo/_smoke_tests/table_urls.py
+++ b/marimo/_smoke_tests/table_urls.py
@@ -13,11 +13,20 @@ def __():
 @app.cell
 def __(mo):
     mo.ui.table(
-      [
-        { "title": 'New York', "url": 'https://en.wikipedia.org/wiki/New_York_City', },
-        { "title": 'London', "url": 'https://en.wikipedia.org/wiki/London', },
-        { "title": 'Paris', "url": 'https://en.wikipedia.org/wiki/Paris', },
-      ],
+        [
+            {
+                "title": "New York",
+                "url": "https://en.wikipedia.org/wiki/New_York_City",
+            },
+            {
+                "title": "London",
+                "url": "https://en.wikipedia.org/wiki/London",
+            },
+            {
+                "title": "Paris",
+                "url": "https://en.wikipedia.org/wiki/Paris",
+            },
+        ],
     )
     return
 
diff --git a/tests/_ast/test_sql_visitor.py b/tests/_ast/test_sql_visitor.py
new file mode 100644
index 00000000000..fd833fc49d6
--- /dev/null
+++ b/tests/_ast/test_sql_visitor.py
@@ -0,0 +1,232 @@
+from __future__ import annotations
+
+import ast
+from textwrap import dedent
+
+import pytest
+
+from marimo._ast.sql_visitor import SQLVisitor, find_created_tables
+from marimo._dependencies.dependencies import DependencyManager
+
+HAS_DUCKDB = DependencyManager.duckdb.has()
+
+
+def test_execute_with_string_literal() -> None:
+    source_code = "db.execute('SELECT * FROM users')"
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == ["SELECT * FROM users"]
+
+
+def test_sql_with_string_literal() -> None:
+    source_code = "db.sql('UPDATE users SET name = \\'Alice\\' WHERE id = 1')"
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == [
+        "UPDATE users SET name = 'Alice' WHERE id = 1"
+    ]
+
+
+def test_execute_with_f_string() -> None:
+    source_code = 'db.execute(f"SELECT * FROM users WHERE name = {name}")'
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == ["SELECT * FROM users WHERE name = '_'"]
+
+
+def test_no_sql_calls() -> None:
+    source_code = "print('Hello, world!')"
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == []
+
+
+def test_sql_with_multiple_arguments() -> None:
+    source_code = (
+        "db.sql('SELECT * FROM users', 'This should not be captured')"
+    )
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == ["SELECT * FROM users"]
+
+
+def test_multiple_sql_calls() -> None:
+    source_code = dedent("""
+        a = db.sql('SELECT * FROM users', 'This should not be captured')
+        b = db.sql('ALTER TABLE users ADD COLUMN name TEXT')
+        c = db.sql('UPDATE users SET name = \\'Alice\\' WHERE id = 1')
+        """)
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == [
+        "SELECT * FROM users",
+        "ALTER TABLE users ADD COLUMN name TEXT",
+        "UPDATE users SET name = 'Alice' WHERE id = 1",
+    ]
+
+
+def test_sql_with_variable() -> None:
+    source_code = dedent("""
+      var = f"SELECT * FROM users WHERE name = {name}"
+      db.sql(var)
+    """)
+    tree = ast.parse(source_code)
+    visitor = SQLVisitor()
+    visitor.visit(tree)
+    assert visitor.get_sqls() == []
+
+
+@pytest.mark.skipif(not HAS_DUCKDB, reason="Missing DuckDB")
+class TestFindCreatedTables:
+    @staticmethod
+    def test_find_created_tables_simple() -> None:
+        sql = "CREATE TABLE test_table (id INT, name VARCHAR(255));"
+        assert find_created_tables(sql) == ["test_table"]
+
+    @staticmethod
+    def test_find_created_tables_multiple() -> None:
+        sql = """
+        CREATE TABLE table1 (id INT);
+        CREATE TABLE table2 (name VARCHAR(255));
+        """
+        assert find_created_tables(sql) == ["table1", "table2"]
+
+    @staticmethod
+    def test_find_created_tables_with_comments() -> None:
+        sql = """
+        CREATE TABLE
+        -- This is a comment
+        IF NOT EXISTS
+        -- This is another comment
+        table1 (id INT);
+        -- This is a comment
+        CREATE TABLE table2 (name VARCHAR(255));
+        """
+        assert find_created_tables(sql) == ["table1", "table2"]
+
+    @staticmethod
+    def test_find_created_tables_with_or_replace() -> None:
+        sql = "CREATE OR REPLACE TABLE test_table (id INT);"
+        assert find_created_tables(sql) == ["test_table"]
+
+    @staticmethod
+    def test_find_created_tables_temporary() -> None:
+        sql = "CREATE TEMPORARY TABLE temp_table (id INT);"
+        assert find_created_tables(sql) == ["temp_table"]
+
+    @staticmethod
+    def test_find_created_tables_if_not_exists() -> None:
+        sql = "CREATE TABLE IF NOT EXISTS new_table (id INT);"
+        assert find_created_tables(sql) == ["new_table"]
+
+    @staticmethod
+    def test_find_created_tables_complex() -> None:
+        sql = """
+        CREATE TABLE table1 (id INT);
+        CREATE OR REPLACE TEMPORARY TABLE IF NOT EXISTS table2 (name VARCHAR(255));
+        CREATE TABLE table3 (date DATE);
+        """  # noqa: E501
+        assert find_created_tables(sql) == ["table1", "table2", "table3"]
+
+    @staticmethod
+    def test_find_created_tables_no_create() -> None:
+        sql = "SELECT * FROM existing_table;"
+        assert find_created_tables(sql) == []
+
+    @staticmethod
+    def test_find_created_tables_case_insensitive() -> None:
+        sql = "create TABLE Test_Table (id INT);"
+        assert find_created_tables(sql) == ["Test_Table"]
+
+    @staticmethod
+    @pytest.mark.parametrize(
+        "query",
+        [
+            "",
+            "   ",
+            ";",
+        ],
+    )
+    def test_find_created_tables_empty_input(query: str) -> None:
+        assert find_created_tables(query) == []
+
+    @staticmethod
+    @pytest.mark.parametrize(
+        "query",
+        [
+            """
+            -- This is a comment
+            CREATE TABLE my_table (
+                my_column INT, -- Inline comment
+                my_other_column INT
+            )
+            """,
+            """
+            /* Multi-line
+            comment */
+            CREATE OR REPLACE TABLE my_table (
+                my_column INT,
+                my_other_column INT
+            )
+            """,
+            """
+            CREATE TABLE IF NOT EXISTS my_table
+            -- Comment before AS
+            AS
+            /* Comment
+            before SELECT */
+            SELECT * FROM read_csv()
+            """,
+            """
+            CREATE TEMPORARY TABLE my_table AS
+            -- Comment in the middle
+            SELECT * FROM existing_table
+            """,
+            """
+            -- Comment at the start
+            CREATE OR REPLACE TEMP TABLE IF NOT EXISTS my_table (
+                id INT, -- Comment after column
+                name VARCHAR
+            ) -- Comment at the end
+            """,
+        ],
+    )
+    def test_find_created_tables_many_comments(query: str) -> None:
+        assert find_created_tables(query) == ["my_table"]
+
+    @staticmethod
+    def test_find_created_tables_weird_names() -> None:
+        sql = """
+        CREATE TABLE "my--table" (
+            "column/*with*/comment" INT,
+            "another--column" VARCHAR
+        );
+
+        CREATE TABLE my_table_with_select AS
+        SELECT *
+        FROM (
+            VALUES
+            ('a', 1),
+            ('b', 2)
+        ) AS t("col--1", "col--2");
+
+        CREATE TABLE "my/*weird*/table" (id INT);
+        """
+        assert find_created_tables(sql) == [
+            "my--table",
+            "my_table_with_select",
+            "my/*weird*/table",
+        ]
+
+
+@pytest.mark.skipif(
+    HAS_DUCKDB, reason="Test requires DuckDB to be unavailable"
+)
+def test_find_created_tables_duckdb_not_available() -> None:
+    assert find_created_tables("CREATE TABLE test (id INT);") is None
diff --git a/tests/_data/test_sql_visitor.py b/tests/_data/test_sql_visitor.py
deleted file mode 100644
index 3c578efd5a4..00000000000
--- a/tests/_data/test_sql_visitor.py
+++ /dev/null
@@ -1,77 +0,0 @@
-from __future__ import annotations
-
-import ast
-from textwrap import dedent
-
-from marimo._data.sql_visitor import SQLVisitor
-
-
-def test_execute_with_string_literal():
-    source_code = "db.execute('SELECT * FROM users')"
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == ["SELECT * FROM users"]
-
-
-def test_sql_with_string_literal():
-    source_code = "db.sql('UPDATE users SET name = \\'Alice\\' WHERE id = 1')"
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == [
-        "UPDATE users SET name = 'Alice' WHERE id = 1"
-    ]
-
-
-def test_execute_with_f_string():
-    source_code = 'db.execute(f"SELECT * FROM users WHERE name = {name}")'
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == ["SELECT * FROM users WHERE name = '_'"]
-
-
-def test_no_sql_calls():
-    source_code = "print('Hello, world!')"
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == []
-
-
-def test_sql_with_multiple_arguments():
-    source_code = (
-        "db.sql('SELECT * FROM users', 'This should not be captured')"
-    )
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == ["SELECT * FROM users"]
-
-
-def test_multiple_sql_calls():
-    source_code = dedent("""
-        a = db.sql('SELECT * FROM users', 'This should not be captured')
-        b = db.sql('ALTER TABLE users ADD COLUMN name TEXT')
-        c = db.sql('UPDATE users SET name = \\'Alice\\' WHERE id = 1')
-        """)
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == [
-        "SELECT * FROM users",
-        "ALTER TABLE users ADD COLUMN name TEXT",
-        "UPDATE users SET name = 'Alice' WHERE id = 1",
-    ]
-
-
-def test_sql_with_variable():
-    source_code = dedent("""
-      var = f"SELECT * FROM users WHERE name = {name}"
-      db.sql(var)
-    """)
-    tree = ast.parse(source_code)
-    visitor = SQLVisitor()
-    visitor.visit(tree)
-    assert visitor.get_sqls() == []


Code review comments:
diff_hunk: @@ -0,0 +1,160 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import re
+from typing import Any, Optional
+
+from marimo._dependencies.dependencies import DependencyManager
+
+
+class SQLVisitor(ast.NodeVisitor):
+    """
+    Find any SQL queries in the AST.
+    This should be inside a function called `.execute` or `.sql`.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._sqls: list[str] = []
+
+    def visit_Call(self, node: ast.Call) -> None:
+        # Check if the call is a method call and the method is named
+        # either 'execute' or 'sql'
+        if isinstance(node.func, ast.Attribute) and node.func.attr in (
+            "execute",
+            "sql",
+        ):
+            # Check if there are arguments and the first argument is a
+            # string or f-string
+            if node.args:
+                first_arg = node.args[0]
+                sql: Optional[str] = None
+                if isinstance(first_arg, ast.Constant):
+                    sql = first_arg.s
+                elif isinstance(first_arg, ast.JoinedStr):
+                    sql = normalize_sql_f_string(first_arg)
+
+                if sql is not None:
+                    # Append the SQL query to the list
+                    self._sqls.append(sql)
+        # Continue walking through the AST
+        self.generic_visit(node)
+
+    def get_sqls(self) -> list[str]:
+        return self._sqls
+
+
+def normalize_sql_f_string(node: ast.JoinedStr) -> str:
+    """
+    Normalize a f-string to a string by joining the parts.
+
+    We add placeholder for {...} expressions in the f-string.
+    This is so we can create a valid SQL query to be passed to
+    other utilities.
+    """
+
+    def print_part(part: ast.expr) -> str:
+        if isinstance(part, ast.FormattedValue):
+            return print_part(part.value)
+        elif isinstance(part, ast.JoinedStr):
+            return normalize_sql_f_string(part)
+        elif isinstance(part, ast.Constant):
+            return str(part.s)
+        else:
+            # Just add '_' as a placeholder for {...} expressions
+            return "'_'"
+
+    result = "".join(print_part(part) for part in node.values)
+    # remove any double '' created by the f-string
+    return result.replace("''", "'")
+
+
+def find_created_tables(sql_statement: str) -> list[str]:
+    """
+    Find the tables created in a SQL statement.
+
+    This function uses the DuckDB tokenizer to find the tables created
+    in a SQL statement. It returns a list of the table names created
+    in the statement.
+
+    Args:
+        sql_statement: The SQL statement to parse.
+
+    Returns:
+        A list of the table names created in the statement.
+    """
+    if not DependencyManager.duckdb.has():
+        return []
+
+    import duckdb
+
+    def remove_comments(sql: str) -> str:
+        # Function to replace comments with spaces, preserving newlines
+        def replace_with_spaces(match: re.Match[Any]) -> str:
+            return " " * len(match.group())
+
+        # Split the SQL into strings and non-strings
+        parts = re.split(r'(\'(?:\'\'|[^\'])*\'|"(?:""|[^"])*")', sql)
+
+        for i in range(0, len(parts), 2):
+            # Remove single-line comments
+            parts[i] = re.sub(
+                r"--.*$", replace_with_spaces, parts[i], flags=re.MULTILINE
+            )
+
+            # Remove multi-line comments
+            parts[i] = re.sub(r"/\*[\s\S]*?\*/", replace_with_spaces, parts[i])
+
+        # Join the parts back together
+        return "".join(parts)
+
+    sql_statement = remove_comments(sql_statement)
+
+    tokens = duckdb.tokenize(sql_statement)
+    created_tables: list[str] = []
+    i = 0
+
+    def token_str(i: int) -> str:
+        token = tokens[i]
+        start = token[0]
+        end = len(sql_statement) - 1
+        if i + 1 < len(tokens):
+            end = tokens[i + 1][0]
+        return sql_statement[start:end].strip()
+
+    def keyword_token_str(i: int) -> str:
+        return token_str(i).lower()
+
+    def token_type(i: int) -> str:
+        return tokens[i][1]
+
+    # See https://duckdb.org/docs/sql/statements/create_table#syntax
+    # for more information on the CREATE TABLE syntax.
+    while i < len(tokens):
+        if (
+            keyword_token_str(i) == "create"
+            and token_type(i) == duckdb.token_type.keyword
path: marimo/_ast/sql_visitor.py
text: @user1:
Could fold this into keyword_token_str (maybe as `is_keyword_token`)? or not do it at all, since this is the only place you check for the type

@author:
yea good call

diff_hunk: @@ -0,0 +1,160 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import re
+from typing import Any, Optional
+
+from marimo._dependencies.dependencies import DependencyManager
+
+
+class SQLVisitor(ast.NodeVisitor):
+    """
+    Find any SQL queries in the AST.
+    This should be inside a function called `.execute` or `.sql`.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._sqls: list[str] = []
+
+    def visit_Call(self, node: ast.Call) -> None:
+        # Check if the call is a method call and the method is named
+        # either 'execute' or 'sql'
+        if isinstance(node.func, ast.Attribute) and node.func.attr in (
+            "execute",
+            "sql",
+        ):
+            # Check if there are arguments and the first argument is a
+            # string or f-string
+            if node.args:
+                first_arg = node.args[0]
+                sql: Optional[str] = None
+                if isinstance(first_arg, ast.Constant):
+                    sql = first_arg.s
+                elif isinstance(first_arg, ast.JoinedStr):
+                    sql = normalize_sql_f_string(first_arg)
+
+                if sql is not None:
+                    # Append the SQL query to the list
+                    self._sqls.append(sql)
+        # Continue walking through the AST
+        self.generic_visit(node)
+
+    def get_sqls(self) -> list[str]:
+        return self._sqls
+
+
+def normalize_sql_f_string(node: ast.JoinedStr) -> str:
+    """
+    Normalize a f-string to a string by joining the parts.
+
+    We add placeholder for {...} expressions in the f-string.
+    This is so we can create a valid SQL query to be passed to
+    other utilities.
+    """
+
+    def print_part(part: ast.expr) -> str:
+        if isinstance(part, ast.FormattedValue):
+            return print_part(part.value)
+        elif isinstance(part, ast.JoinedStr):
+            return normalize_sql_f_string(part)
+        elif isinstance(part, ast.Constant):
+            return str(part.s)
+        else:
+            # Just add '_' as a placeholder for {...} expressions
+            return "'_'"
+
+    result = "".join(print_part(part) for part in node.values)
+    # remove any double '' created by the f-string
+    return result.replace("''", "'")
+
+
+def find_created_tables(sql_statement: str) -> list[str]:
+    """
+    Find the tables created in a SQL statement.
+
+    This function uses the DuckDB tokenizer to find the tables created
+    in a SQL statement. It returns a list of the table names created
+    in the statement.
+
+    Args:
+        sql_statement: The SQL statement to parse.
+
+    Returns:
+        A list of the table names created in the statement.
+    """
+    if not DependencyManager.duckdb.has():
+        return []
+
+    import duckdb
+
+    def remove_comments(sql: str) -> str:
+        # Function to replace comments with spaces, preserving newlines
+        def replace_with_spaces(match: re.Match[Any]) -> str:
+            return " " * len(match.group())
+
+        # Split the SQL into strings and non-strings
+        parts = re.split(r'(\'(?:\'\'|[^\'])*\'|"(?:""|[^"])*")', sql)
+
+        for i in range(0, len(parts), 2):
+            # Remove single-line comments
+            parts[i] = re.sub(
+                r"--.*$", replace_with_spaces, parts[i], flags=re.MULTILINE
+            )
+
+            # Remove multi-line comments
+            parts[i] = re.sub(r"/\*[\s\S]*?\*/", replace_with_spaces, parts[i])
+
+        # Join the parts back together
+        return "".join(parts)
+
+    sql_statement = remove_comments(sql_statement)
+
+    tokens = duckdb.tokenize(sql_statement)
+    created_tables: list[str] = []
+    i = 0
+
+    def token_str(i: int) -> str:
+        token = tokens[i]
+        start = token[0]
+        end = len(sql_statement) - 1
path: marimo/_ast/sql_visitor.py
text: @user1:
You can remove the comment logic if you extract the tokens on boundaries:
```python
def token_str(i: int) -> str:
    token = tokens[i]
    start = token[0]
    maybe_end = re.search(r'[\s\-/]', sql_statement[start:])
    end = start + maybe_end.start() if maybe_end else len(sql_statement)
    if i + 1 < len(tokens):
        # For tokens squashed together e.g. '(select' or 'x);;'
        # in (select * from x);;
        end = min(end, tokens[i + 1][0]) 
    return sql_statement[start:end]
```

tokenizer also handles strings

@user1:
Oh, just saw the tests below with the weird table names. You might also be missing the `SELECT '''' AS s` edge case: https://duckdb.org/docs/sql/data_types/literal_types.html#escape-string-literals

and string consts `e'\n\b'`

```python
def token_str(i: int) -> str:
    token = tokens[i]
    start = token[0]
    boundary = r'[\s\-/]'
    quote = None
    if token_type(i) == duckdb.token_type.string_const:
        quote = sql_statement[start].lower()
        # Handle e'string' notation
        if quote == "e":
            start += 1
            # Should be '
            quote = sql_statement[start]
        # Quotes are released by doubling
        # From the duckdb docs:
        # SELECT '''' AS s
        # is valid (table name ')
        boundary = fr"{quote}(?=[^{quote}])"
        start += 1
    maybe_end = re.search(boundary, sql_statement[start:])
    end = start + maybe_end.start() if maybe_end else len(sql_statement)
    if i + 1 < len(tokens):
        # For tokens squashed together e.g. '(select' or 'x);;;;'
        # in (select * from x);;;;
        end = min(end, tokens[i + 1][0]) 
    if quote:
        return re.sub(quote * 2, quote, sql_statement[start:end])
    return sql_statement[start:end]
 ```
 
which also strips quotes for free

@author:
this doesnt handle the case:
```sql
       CREATE TABLE "my--table" (
            "column/*with*/comment" INT,
            "another--column" VARCHAR
        );

        CREATE TABLE my_table_with_select AS
        SELECT *
        FROM (
            VALUES
            ('a', 1),
            ('b', 2)
        ) AS t("col--1", "col--2");

        CREATE TABLE "my/*weird*/table" (id INT);

        CREATE TABLE "with a space" (id INT);
```

@author:
i added:
```python
if sql_statement[start] == '"':
            # If it starts with a quote, find the matching end quote
            end = sql_statement.find('"', start + 1) + 1
```

@user1:
Oh woops, I broke it when checking for the e-strings. This still doesn't get single quotes or escaped quotes

fixed my bug, string table names are still "identifiers" not string_const unless they have the `e`

```python
def token_str(i: int) -> str:
    token = tokens[i]
    start = token[0]
    boundary = r'([\s\-/])'
    quote = None
    if token_type(i) == duckdb.token_type.string_const:
        # Handle e'string' notation
        if sql_statement[start].lower() == "e":
            start += 1
    if sql_statement[start] in ("\"", "'"):
        quote = sql_statement[start]
        # Quotes are released by doubling
        # From the duckdb docs:
        # SELECT '''' AS s
        # is valid (table name ')
        boundary = fr"(?<!{quote})(?:{quote * 2})*({quote})(?!{quote})"
        start += 1
    maybe_end = re.search(boundary, sql_statement[start:])
    end = start + maybe_end.end() - 1 if maybe_end else len(sql_statement)
    if i + 1 < len(tokens):
        # For tokens squashed together e.g. '(select' or 'x);;;;'
        # in (select * from x);;;;
        end = min(end, tokens[i + 1][0])
    if quote:
        return re.sub(quote * 2, quote, sql_statement[start:end])
    return sql_statement[start:end]
```

Which seems to work: https://marimo.app/l/tc7cqv

@user1:
cc @user2 I think the current logic will still break for the escaped string literal notation https://duckdb.org/docs/sql/data_types/literal_types.html#escape-string-literals

@author:
im going to merge this as-is and can follow-up with this change

diff_hunk: @@ -0,0 +1,160 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+import ast
+import re
+from typing import Any, Optional
+
+from marimo._dependencies.dependencies import DependencyManager
+
+
+class SQLVisitor(ast.NodeVisitor):
+    """
+    Find any SQL queries in the AST.
+    This should be inside a function called `.execute` or `.sql`.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._sqls: list[str] = []
+
+    def visit_Call(self, node: ast.Call) -> None:
+        # Check if the call is a method call and the method is named
+        # either 'execute' or 'sql'
+        if isinstance(node.func, ast.Attribute) and node.func.attr in (
+            "execute",
+            "sql",
+        ):
+            # Check if there are arguments and the first argument is a
+            # string or f-string
+            if node.args:
+                first_arg = node.args[0]
+                sql: Optional[str] = None
+                if isinstance(first_arg, ast.Constant):
+                    sql = first_arg.s
+                elif isinstance(first_arg, ast.JoinedStr):
+                    sql = normalize_sql_f_string(first_arg)
+
+                if sql is not None:
+                    # Append the SQL query to the list
+                    self._sqls.append(sql)
+        # Continue walking through the AST
+        self.generic_visit(node)
+
+    def get_sqls(self) -> list[str]:
+        return self._sqls
+
+
+def normalize_sql_f_string(node: ast.JoinedStr) -> str:
+    """
+    Normalize a f-string to a string by joining the parts.
+
+    We add placeholder for {...} expressions in the f-string.
+    This is so we can create a valid SQL query to be passed to
+    other utilities.
+    """
+
+    def print_part(part: ast.expr) -> str:
+        if isinstance(part, ast.FormattedValue):
+            return print_part(part.value)
+        elif isinstance(part, ast.JoinedStr):
+            return normalize_sql_f_string(part)
+        elif isinstance(part, ast.Constant):
+            return str(part.s)
+        else:
+            # Just add '_' as a placeholder for {...} expressions
+            return "'_'"
+
+    result = "".join(print_part(part) for part in node.values)
+    # remove any double '' created by the f-string
+    return result.replace("''", "'")
+
+
+def find_created_tables(sql_statement: str) -> list[str]:
+    """
+    Find the tables created in a SQL statement.
+
+    This function uses the DuckDB tokenizer to find the tables created
+    in a SQL statement. It returns a list of the table names created
+    in the statement.
+
+    Args:
+        sql_statement: The SQL statement to parse.
+
+    Returns:
+        A list of the table names created in the statement.
+    """
+    if not DependencyManager.duckdb.has():
+        return []
+
+    import duckdb
+
+    def remove_comments(sql: str) -> str:
+        # Function to replace comments with spaces, preserving newlines
+        def replace_with_spaces(match: re.Match[Any]) -> str:
+            return " " * len(match.group())
+
+        # Split the SQL into strings and non-strings
+        parts = re.split(r'(\'(?:\'\'|[^\'])*\'|"(?:""|[^"])*")', sql)
+
+        for i in range(0, len(parts), 2):
+            # Remove single-line comments
+            parts[i] = re.sub(
+                r"--.*$", replace_with_spaces, parts[i], flags=re.MULTILINE
+            )
+
+            # Remove multi-line comments
+            parts[i] = re.sub(r"/\*[\s\S]*?\*/", replace_with_spaces, parts[i])
+
+        # Join the parts back together
+        return "".join(parts)
+
+    sql_statement = remove_comments(sql_statement)
+
+    tokens = duckdb.tokenize(sql_statement)
+    created_tables: list[str] = []
+    i = 0
+
+    def token_str(i: int) -> str:
+        token = tokens[i]
+        start = token[0]
+        end = len(sql_statement) - 1
+        if i + 1 < len(tokens):
+            end = tokens[i + 1][0]
+        return sql_statement[start:end].strip()
+
+    def keyword_token_str(i: int) -> str:
+        return token_str(i).lower()
+
+    def token_type(i: int) -> str:
+        return tokens[i][1]
+
+    # See https://duckdb.org/docs/sql/statements/create_table#syntax
+    # for more information on the CREATE TABLE syntax.
+    while i < len(tokens):
+        if (
+            keyword_token_str(i) == "create"
+            and token_type(i) == duckdb.token_type.keyword
+        ):
+            i += 1
+            if i < len(tokens) and keyword_token_str(i) == "or":
+                i += 2  # Skip 'OR REPLACE'
+            if i < len(tokens) and keyword_token_str(i) in (
+                "temporary",
+                "temp",
+            ):
+                i += 1  # Skip 'TEMPORARY' or 'TEMP'
+            if i < len(tokens) and keyword_token_str(i) == "table":
+                i += 1
+                if i < len(tokens) and keyword_token_str(i) == "if":
+                    i += 3  # Skip 'IF NOT EXISTS'
+                if i < len(tokens):
+                    table_name = token_str(i)
+                    # Remove quotes if present
+                    if table_name.startswith('"') and table_name.endswith('"'):
+                        table_name = table_name[1:-1]
+
+                    created_tables.append(table_name)
path: marimo/_ast/sql_visitor.py
text: @user1:
strip quotes and make lowercase
`create table "my_table"` is valid, and duckdb is case insensitive

@author:
i think its worth preserving the case in this and then check for case insensitive later, since you may want clearer error messages with the correct case

diff_hunk: @@ -401,18 +401,34 @@ def visit_Call(self, node: ast.Call) -> None:
                         tables = duckdb.get_table_names(statement.query)
                     except duckdb.ProgrammingError:
                         self.generic_visit(node)
-                        return
+                        continue
                     except BaseException as e:
                         LOGGER.warning("Unexpected duckdb error %s", e)
                         self.generic_visit(node)
-                        return
+                        continue
 
                     for table in tables:
                         # Table may be a URL or something else that
                         # isn't a Python variable
                         if table.isidentifier():
                             self._add_ref(table, deleted=False)
 
+                    # Add all tables created in the query to the defs
+                    try:
+                        created_tables = find_created_tables(sql)
+                    except duckdb.ProgrammingError:
+                        self.generic_visit(node)
+                        continue
+                    except BaseException as e:
+                        LOGGER.warning("Unexpected duckdb error %s", e)
+                        self.generic_visit(node)
+                        continue
+
+                    for _table in created_tables:
path: marimo/_ast/visitor.py
text: @user1:
Nice, I think the errors around duplicate defs need to maybe tell the difference between a table def and normal def. idk, what would the errors look like now I wonder