Repository: sqlfluff/sqlfluff
Pull request: L026: Handle DML statements and multiple levels of nesting
<!--Firstly, thanks for adding this feature! Secondly, please check the key steps against the checklist below to make your contribution easy to merge.-->

<!--Please give the Pull Request a meaningful title (including the dialect this PR is for if it is dialect specific), as this will automatically be added to the release notes, and then the Change Log.-->

### Brief summary of the change made
Fixes #2265

Rewrite of L026:
* Supporting `UPDATE` and `DELETE` statements
* Support for arbitrarily nested subqueries (courtesy of `SelectCrawler`)
* Improve handling of multi-part table references (i.e. schema.table)

<!--If there is an open issue for this, then please include `fixes #XXXX` or `closes #XXXX` replacing `XXXX` with the issue number and it will automatically close the issue when the pull request is merged. Alternatively if not fully closed you can say `makes progress on #XXXX` to create a link on that issue without closing it.-->


### Are there any other side effects of this change that we should be aware of?


### Pull Request checklist
- [x] Please confirm you have completed any of the necessary steps below.

- Included test cases to demonstrate any code changes, which may be one or more of the following:
  - `.yml` rule test cases in `test/fixtures/rules/std_rule_cases`.
  - `.sql`/`.yml` parser test cases in `test/fixtures/dialects` (note YML files can be auto generated with `tox -e generate-fixture-yml`).
  - Full autofix test cases in `test/fixtures/linter/autofix`.
  - Other.
- Added appropriate documentation for the change.
- Created GitHub issues for any relevant followup/future enhancements if appropriate.


Problem statement:
L026: Rule incorrectly flag column does not exist in `FROM` clause in an UPDATE statement.
## Expected Behaviour

L026 should not fail when a subquery in an UPDATE statement references a column from the UPDATE target.

## Observed Behaviour

L026 failed due to reference was not found in the FROM clause with the following error printed (When using `sample.sql` content below)

```
L:   7 | P:  28 | L026 | Reference 'my_table.id' refers to table/view not found
                       | in the FROM clause or found in parent subquery.
```

## Steps to Reproduce

1. Create `sample.sql` with the content below
```
UPDATE my_table
SET row_sum = (
    SELECT COUNT(*) AS row_sum
    FROM
        another_table
    WHERE
        another_table.id = my_table.id
);
```
2. Run SQLFluff by `sqlfluff lint sample.sql`

## Dialect

Default / Ansi (No dialect specified)

## Version
```
(.venv) ~/code/sqlfluff (main) $ sqlfluff --version
sqlfluff, version 0.9.0
```

```
(.venv) ~/code/sqlfluff (main) $ python --version
Python 3.9.9
```

## Configuration
Default. No customization.

Hints from issue discussion:

Commit message: Merge branch 'main' into bhart-issue_2265_l026_handle_update_statements

Patch to review:
diff --git a/src/sqlfluff/core/rules/analysis/select.py b/src/sqlfluff/core/rules/analysis/select.py
index 33eba2e973f..2e0a313ef37 100644
--- a/src/sqlfluff/core/rules/analysis/select.py
+++ b/src/sqlfluff/core/rules/analysis/select.py
@@ -124,7 +124,7 @@ def _has_value_table_function(table_expr, dialect):
         # We need the dialect to get the value table function names. If
         # we don't have it, assume the clause does not have a value table
         # function.
-        return False
+        return False  # pragma: no cover
 
     for function_name in table_expr.recursive_crawl("function_name"):
         # Other rules can increase whitespace in the function name, so use strip to remove
@@ -138,7 +138,7 @@ def _get_pivot_table_columns(segment, dialect):
     if not dialect:
         # We need the dialect to get the pivot table column names. If
         # we don't have it, assume the clause does not have a pivot table
-        return []
+        return []  # pragma: no cover
 
     fc = segment.get_child("from_pivot_expression")
     if not fc:
diff --git a/src/sqlfluff/rules/L025.py b/src/sqlfluff/rules/L025.py
index de116b33524..5b9d0a83265 100644
--- a/src/sqlfluff/rules/L025.py
+++ b/src/sqlfluff/rules/L025.py
@@ -85,16 +85,18 @@ def _eval(self, context: RuleContext) -> EvalResultType:
     def _analyze_table_aliases(cls, query: L025Query, dialect: Dialect):
         # Get table aliases defined in query.
         for selectable in query.selectables:
-            select_info = get_select_statement_info(selectable.selectable, dialect)
+            select_info = selectable.select_info
             if select_info:
                 # Record the aliases.
                 query.aliases += select_info.table_aliases
 
-                # Look at each table reference; if its an alias reference,
+                # Look at each table reference; if it's an alias reference,
                 # resolve the alias: could be an alias defined in "query"
                 # itself or an "ancestor" query.
                 for r in select_info.reference_buffer:
-                    for tr in r.extract_possible_references(level=r.ObjectReferenceLevel.TABLE):  # type: ignore
+                    for tr in r.extract_possible_references(
+                        level=r.ObjectReferenceLevel.TABLE
+                    ):
                         # This function walks up the query's parent stack if necessary.
                         cls._resolve_and_mark_reference(query, tr.part)
 
diff --git a/src/sqlfluff/rules/L026.py b/src/sqlfluff/rules/L026.py
index 1cd2c347614..21d8aea2fe1 100644
--- a/src/sqlfluff/rules/L026.py
+++ b/src/sqlfluff/rules/L026.py
@@ -1,13 +1,33 @@
 """Implementation of Rule L026."""
 
-from sqlfluff.core.rules.analysis.select import get_aliases_from_select
-from sqlfluff.core.rules.base import EvalResultType, LintResult, RuleContext
+from dataclasses import dataclass, field
+from typing import cast, List, Optional
+
+from sqlfluff.core.dialects.base import Dialect
+from sqlfluff.core.rules.analysis.select_crawler import (
+    Query as SelectCrawlerQuery,
+    SelectCrawler,
+)
+from sqlfluff.core.dialects.common import AliasInfo
+from sqlfluff.core.rules.base import (
+    BaseRule,
+    LintResult,
+    RuleContext,
+    EvalResultType,
+)
+from sqlfluff.core.rules.functional import sp
 from sqlfluff.core.rules.doc_decorators import document_configuration
-from sqlfluff.rules.L020 import Rule_L020
+
+
+@dataclass
+class L026Query(SelectCrawlerQuery):
+    """SelectCrawler Query with custom L026 info."""
+
+    aliases: List[AliasInfo] = field(default_factory=list)
 
 
 @document_configuration
-class Rule_L026(Rule_L020):
+class Rule_L026(BaseRule):
     """References cannot reference objects not present in ``FROM`` clause.
 
     NB: This rule is disabled by default for BigQuery due to its use of
@@ -36,61 +56,7 @@ class Rule_L026(Rule_L020):
 
     config_keywords = ["force_enable"]
 
-    @staticmethod
-    def _is_bad_tbl_ref(table_aliases, parent_select, tbl_ref):
-        """Given a table reference, try to find what it's referring to."""
-        # Is it referring to one of the table aliases?
-        if tbl_ref[0] in [a.ref_str for a in table_aliases]:
-            # Yes. Therefore okay.
-            return False
-
-        # Not a table alias. It it referring to a correlated subquery?
-        if parent_select:
-            parent_aliases, _ = get_aliases_from_select(parent_select)
-            if parent_aliases and tbl_ref[0] in [a[0] for a in parent_aliases]:
-                # Yes. Therefore okay.
-                return False
-
-        # It's not referring to an alias or a correlated subquery. Looks like a
-        # bad reference (i.e. referring to something unknown.)
-        return True
-
-    def _lint_references_and_aliases(
-        self,
-        table_aliases,
-        standalone_aliases,
-        references,
-        col_aliases,
-        using_cols,
-        parent_select,
-    ):
-        # A buffer to keep any violations.
-        violation_buff = []
-
-        # Check all the references that we have, do they reference present aliases?
-        for r in references:
-            tbl_refs = r.extract_possible_references(level=r.ObjectReferenceLevel.TABLE)
-            if tbl_refs and all(
-                self._is_bad_tbl_ref(table_aliases, parent_select, tbl_ref)
-                for tbl_ref in tbl_refs
-            ):
-                violation_buff.append(
-                    LintResult(
-                        # Return the first segment rather than the string
-                        anchor=tbl_refs[0].segments[0],
-                        description=f"Reference {r.raw!r} refers to table/view "
-                        "not found in the FROM clause or found in parent "
-                        "subquery.",
-                    )
-                )
-        return violation_buff or None
-
     def _eval(self, context: RuleContext) -> EvalResultType:
-        """Override Rule L020 for dialects that use structs.
-
-        Some dialects use structs (e.g. column.field) which look like
-        table references and so incorrectly trigger this rule.
-        """
         # Config type hints
         self.force_enable: bool
 
@@ -100,4 +66,84 @@ def _eval(self, context: RuleContext) -> EvalResultType:
         ):
             return LintResult()
 
-        return super()._eval(context=context)
+        violations: List[LintResult] = []
+        start_types = ["select_statement", "delete_statement", "update_statement"]
+        if context.segment.is_type(
+            *start_types
+        ) and not context.functional.parent_stack.any(sp.is_type(*start_types)):
+            dml_target_table: Optional[str] = None
+            if not context.segment.is_type("select_statement"):
+                # Extract first table reference.
+                table_reference = next(
+                    context.segment.recursive_crawl("table_reference"), None
+                )
+                if table_reference:
+                    dml_target_table = table_reference.raw
+
+            # Verify table references in any SELECT statements found in or
+            # below context.segment.
+            crawler = SelectCrawler(
+                context.segment, context.dialect, query_class=L026Query
+            )
+            query: L026Query = cast(L026Query, crawler.query_tree)
+            self._analyze_table_references(
+                query, dml_target_table, context.dialect, violations
+            )
+        return violations or None
+
+    def _analyze_table_references(
+        self,
+        query: L026Query,
+        dml_target_table: Optional[str],
+        dialect: Dialect,
+        violations: List[LintResult],
+    ):
+        # For each query...
+        for selectable in query.selectables:
+            select_info = selectable.select_info
+            if select_info:
+                # Record the table references.
+                query.aliases += select_info.table_aliases
+
+                # Try and resolve each table reference.
+                for r in select_info.reference_buffer:
+                    tbl_refs = r.extract_possible_references(
+                        level=r.ObjectReferenceLevel.TABLE
+                    )
+                    # This function walks up the query's parent stack if necessary.
+                    violation = self._resolve_reference(
+                        r, tbl_refs, dml_target_table, query
+                    )
+                    if violation:
+                        violations.append(violation)
+
+        # Visit children.
+        for child in query.children:
+            self._analyze_table_references(
+                cast(L026Query, child), dml_target_table, dialect, violations
+            )
+
+    def _resolve_reference(
+        self, r, tbl_refs, dml_target_table: Optional[str], query: L026Query
+    ):
+        # Does this query define the referenced table?
+        if tbl_refs and all(
+            tbl_ref[0] not in [a.ref_str for a in query.aliases] for tbl_ref in tbl_refs
+        ):
+            # No. Check the parent query, if there is one.
+            if query.parent:
+                return self._resolve_reference(
+                    r, tbl_refs, dml_target_table, cast(L026Query, query.parent)
+                )
+            # No parent query. If there's a DML statement at the root, check its
+            # target table.
+            elif not dml_target_table or all(
+                tbl_ref[0] != dml_target_table for tbl_ref in tbl_refs
+            ):
+                return LintResult(
+                    # Return the first segment rather than the string
+                    anchor=tbl_refs[0].segments[0],
+                    description=f"Reference {r.raw!r} refers to table/view "
+                    "not found in the FROM clause or found in ancestor "
+                    "statement.",
+                )
diff --git a/test/fixtures/rules/std_rule_cases/L026.yml b/test/fixtures/rules/std_rule_cases/L026.yml
index 9ee970569f6..8c9e096eed1 100644
--- a/test/fixtures/rules/std_rule_cases/L026.yml
+++ b/test/fixtures/rules/std_rule_cases/L026.yml
@@ -110,3 +110,39 @@ test_pass_object_referenced_6:
                 table3
                 on table2.y_id = table3.y_id
     ) as cc
+
+test_pass_object_referenced_7:
+  pass_str: |
+    UPDATE my_table
+    SET row_sum = (
+        SELECT COUNT(*) AS row_sum
+        FROM
+            another_table
+        WHERE
+            another_table.id = my_table.id
+    )
+
+test_fail_object_referenced_7:
+  fail_str: |
+    UPDATE my_table
+    SET row_sum = (
+        SELECT COUNT(*) AS row_sum
+        FROM
+            another_table
+        WHERE
+            another_table.id = my_tableeee.id
+    )
+
+test_pass_object_referenced_8:
+  pass_str: |
+    DELETE FROM agent1
+    WHERE EXISTS(
+        SELECT customer.cust_id FROM customer
+        WHERE agent1.agent_code <> customer.agent_code);
+
+test_fail_object_referenced_8:
+  fail_str: |
+    DELETE FROM agent1
+    WHERE EXISTS(
+        SELECT customer.cust_id FROM customer
+        WHERE agent11.agent_code <> customer.agent_code);


Code review comments:
diff_hunk: @@ -100,4 +66,84 @@ def _eval(self, context: RuleContext) -> EvalResultType:
         ):
             return LintResult()
 
-        return super()._eval(context=context)
+        violations: List[LintResult] = []
+        start_types = ["select_statement", "delete_statement", "update_statement"]
+        if context.segment.is_type(
+            *start_types
+        ) and not context.functional.parent_stack.any(sp.is_type(*start_types)):
+            dml_target_table: Optional[str] = None
+            if not context.segment.is_type("select_statement"):
+                # Extract first table reference.
+                table_reference = next(
+                    context.segment.recursive_crawl("table_reference"), None
+                )
+                if table_reference:
+                    dml_target_table = table_reference.raw
path: src/sqlfluff/rules/L026.py
text: @user1:
I think using raw here means the rule would get confused if the table references were more (or less) qualified in different places?

Bit of a toy example but this raises a lint:
```sql
use schema public;

delete from public.agent1
where exists(
    select customer.cust_id from customer
    where agent1.agent_code <> customer.agent_code);
```

It's a bit of a weird one as it seems niche and relies on the context that we're in the public schema already.
We can't just take the last identifier from the reference as we don't want to not flag:
```sql
delete from schema1.agent1
where exists(
    select customer.cust_id from customer
    where schema2.agent1.agent_code <> customer.agent_code);
```

@user1:
More relevant would be that something like this raises a lint:
```sql
delete from public.agent1
where exists(
    select customer.cust_id from customer
    where public.agent1.agent_code <> customer.agent_code);
```

@author:
You're right -- I'll take a look at this. Presumably, some existing rule(s) already handle similar cases.