Repository: sqlfluff/sqlfluff

Issue:
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.



Discussion:

Patch:
diff --git a/src/sqlfluff/core/rules/analysis/select.py b/src/sqlfluff/core/rules/analysis/select.py
--- 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/core/rules/reference.py b/src/sqlfluff/core/rules/reference.py
new file mode 100644
--- /dev/null
+++ b/src/sqlfluff/core/rules/reference.py
@@ -0,0 +1,26 @@
+"""Components for working with object and table references."""
+from typing import Sequence, Tuple
+
+
+def object_ref_matches_table(
+    possible_references: Sequence[Tuple[str, ...]], targets: Sequence[Tuple[str, ...]]
+) -> bool:
+    """Return True if any of the possible references matches a target."""
+    # Simple case: If there are no references, assume okay
+    # (i.e. no mismatch = good).
+    if not possible_references:
+        return True
+    # Simple case: Reference exactly matches a target.
+    if any(pr in targets for pr in possible_references):
+        return True
+    # Tricky case: If one is shorter than the other, check for a suffix match.
+    # (Note this is an "optimistic" check, i.e. it assumes the ignored parts of
+    # the target don't matter. In a SQL context, this is basically assuming
+    # there was an earlier "USE <<database>>" or similar directive.
+    for pr in possible_references:
+        for t in targets:
+            if (len(pr) < len(t) and pr == t[-len(pr) :]) or (
+                len(t) < len(pr) and t == pr[-len(t) :]
+            ):
+                return True
+    return False
diff --git a/src/sqlfluff/dialects/dialect_ansi.py b/src/sqlfluff/dialects/dialect_ansi.py
--- a/src/sqlfluff/dialects/dialect_ansi.py
+++ b/src/sqlfluff/dialects/dialect_ansi.py
@@ -734,6 +734,18 @@ def extract_possible_references(
             return [refs[-level]]
         return []
 
+    def extract_possible_multipart_references(
+        self, levels: List[Union[ObjectReferenceLevel, int]]
+    ) -> List[Tuple[ObjectReferencePart, ...]]:
+        """Extract possible multipart references, e.g. schema.table."""
+        levels_tmp = [self._level_to_int(level) for level in levels]
+        min_level = min(levels_tmp)
+        max_level = max(levels_tmp)
+        refs = list(self.iter_raw_references())
+        if len(refs) >= max_level:
+            return [tuple(refs[-max_level : 1 - min_level])]
+        return []
+
     @staticmethod
     def _level_to_int(level: Union[ObjectReferenceLevel, int]) -> int:
         # If it's an ObjectReferenceLevel, get the value. Otherwise, assume it's
@@ -1156,7 +1168,6 @@ def get_eventual_alias(self) -> Optional[AliasInfo]:
             return AliasInfo(segment.raw, segment, True, self, alias_expression, ref)
 
         # If not return the object name (or None if there isn't one)
-        # ref = self.get_child("object_reference")
         if ref:
             # Return the last element of the reference.
             penultimate_ref: ObjectReferenceSegment.ObjectReferencePart = list(
diff --git a/src/sqlfluff/dialects/dialect_bigquery.py b/src/sqlfluff/dialects/dialect_bigquery.py
--- a/src/sqlfluff/dialects/dialect_bigquery.py
+++ b/src/sqlfluff/dialects/dialect_bigquery.py
@@ -740,6 +740,18 @@ def extract_possible_references(self, level):
             return [refs[1], refs[2]]
         return super().extract_possible_references(level)  # pragma: no cover
 
+    def extract_possible_multipart_references(self, levels):
+        """Extract possible multipart references, e.g. schema.table."""
+        levels_tmp = [self._level_to_int(level) for level in levels]
+        min_level = min(levels_tmp)
+        max_level = max(levels_tmp)
+        refs = list(self.iter_raw_references())
+        if max_level == self.ObjectReferenceLevel.SCHEMA.value and len(refs) >= 3:
+            return [tuple(refs[0 : max_level - min_level + 1])]
+        # Note we aren't handling other possible cases. We'll add these as
+        # needed.
+        return super().extract_possible_multipart_references(levels)
+
 
 @bigquery_dialect.segment()
 class HyphenatedObjectReferenceSegment(ObjectReferenceSegment):  # type: ignore
diff --git a/src/sqlfluff/rules/L025.py b/src/sqlfluff/rules/L025.py
--- 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
--- 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, Tuple
+
+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
+from sqlfluff.core.rules.reference import object_ref_matches_table
+
+
+@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,128 @@ 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[Tuple[str, ...]] = None
+            if not context.segment.is_type("select_statement"):
+                # Extract first table reference. This will be the target
+                # table in a DELETE or UPDATE statement.
+                table_reference = next(
+                    context.segment.recursive_crawl("table_reference"), None
+                )
+                if table_reference:
+                    dml_target_table = self._table_ref_as_tuple(table_reference)
+
+            # Verify table references in any SELECT statements found in or
+            # below context.segment in the parser tree.
+            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
+
+    @classmethod
+    def _alias_info_as_tuples(cls, alias_info: AliasInfo) -> List[Tuple[str, ...]]:
+        result: List[Tuple[str, ...]] = []
+        if alias_info.aliased:
+            result.append((alias_info.ref_str,))
+        if alias_info.object_reference:
+            result.append(cls._table_ref_as_tuple(alias_info.object_reference))
+        return result
+
+    @staticmethod
+    def _table_ref_as_tuple(table_reference) -> Tuple[str, ...]:
+        return tuple(ref.part for ref in table_reference.iter_raw_references())
+
+    def _analyze_table_references(
+        self,
+        query: L026Query,
+        dml_target_table: Optional[Tuple[str, ...]],
+        dialect: Dialect,
+        violations: List[LintResult],
+    ):
+        # For each query...
+        for selectable in query.selectables:
+            select_info = selectable.select_info
+            if select_info:
+                # Record the available tables.
+                query.aliases += select_info.table_aliases
+
+                # Try and resolve each reference to a value in query.aliases (or
+                # in an ancestor query).
+                for r in select_info.reference_buffer:
+                    # This function walks up the query's parent stack if necessary.
+                    violation = self._resolve_reference(
+                        r, self._get_table_refs(r, dialect), 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
+            )
+
+    @staticmethod
+    def _get_table_refs(ref, dialect):
+        """Given ObjectReferenceSegment, determine possible table references."""
+        tbl_refs = []
+        # First, handle any schema.table references.
+        for sr, tr in ref.extract_possible_multipart_references(
+            levels=[
+                ref.ObjectReferenceLevel.SCHEMA,
+                ref.ObjectReferenceLevel.TABLE,
+            ]
+        ):
+            tbl_refs.append((tr, (sr.part, tr.part)))
+        # Maybe check for simple table references. Two cases:
+        # - For most dialects, skip this if it's a schema+table reference -- the
+        #   reference was specific, so we shouldn't ignore that by looking
+        #   elsewhere.)
+        # - Always do this in BigQuery. BigQuery table references are frequently
+        #   ambiguous because BigQuery SQL supports structures, making some
+        #   multi-level "." references impossible to interpret with certainty.
+        #   We may need to genericize this code someday to support other
+        #   dialects. If so, this check should probably align somehow with
+        #   whether the dialect overrides
+        #   ObjectReferenceSegment.extract_possible_references().
+        if not tbl_refs or dialect.name in ["bigquery"]:
+            for tr in ref.extract_possible_references(
+                level=ref.ObjectReferenceLevel.TABLE
+            ):
+                tbl_refs.append((tr, (tr.part,)))
+        return tbl_refs
+
+    def _resolve_reference(
+        self, r, tbl_refs, dml_target_table: Optional[Tuple[str, ...]], query: L026Query
+    ):
+        # Does this query define the referenced table?
+        possible_references = [tbl_ref[1] for tbl_ref in tbl_refs]
+        targets = []
+        for alias in query.aliases:
+            targets += self._alias_info_as_tuples(alias)
+        if not object_ref_matches_table(possible_references, targets):
+            # 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 not object_ref_matches_table(
+                possible_references, [dml_target_table]
+            ):
+                return LintResult(
+                    # Return the first segment rather than the string
+                    anchor=tbl_refs[0][0].segments[0],
+                    description=f"Reference {r.raw!r} refers to table/view "
+                    "not found in the FROM clause or found in ancestor "
+                    "statement.",
+                )


Test patch:
diff --git a/test/core/rules/reference_test.py b/test/core/rules/reference_test.py
new file mode 100644
--- /dev/null
+++ b/test/core/rules/reference_test.py
@@ -0,0 +1,72 @@
+"""Test components for working with object and table references."""
+import pytest
+
+from sqlfluff.core.rules import reference
+
+
+@pytest.mark.parametrize(
+    "possible_references, targets, result",
+    [
+        # Empty list of references is always True.
+        [[], [("abc",)], True],
+        # Simple cases: one reference, one target.
+        [[("agent1",)], [("agent1",)], True],
+        [[("agent1",)], [("customer",)], False],
+        # Multiple references. If any match, good.
+        [[("bar",), ("user_id",)], [("bar",)], True],
+        [[("foo",), ("user_id",)], [("bar",)], False],
+        # Multiple targets. If any reference matches, good.
+        [[("table1",)], [("table1",), ("table2",), ("table3",)], True],
+        [[("tbl2",)], [("db", "sc", "tbl1")], False],
+        [[("tbl2",)], [("db", "sc", "tbl2")], True],
+        # Multi-part references and targets. If one tuple is shorter than
+        # the other, checks for a suffix match.
+        [
+            [
+                (
+                    "rc",
+                    "tbl1",
+                )
+            ],
+            [("db", "sc", "tbl1")],
+            False,
+        ],
+        [
+            [
+                (
+                    "sc",
+                    "tbl1",
+                )
+            ],
+            [("db", "sc", "tbl1")],
+            True,
+        ],
+        [
+            [
+                (
+                    "cb",
+                    "sc",
+                    "tbl1",
+                )
+            ],
+            [("db", "sc", "tbl1")],
+            False,
+        ],
+        [
+            [
+                (
+                    "db",
+                    "sc",
+                    "tbl1",
+                )
+            ],
+            [("db", "sc", "tbl1")],
+            True,
+        ],
+        [[("public", "agent1")], [("agent1",)], True],
+        [[("public", "agent1")], [("public",)], False],
+    ],
+)
+def test_object_ref_matches_table(possible_references, targets, result):
+    """Test object_ref_matches_table()."""
+    assert reference.object_ref_matches_table(possible_references, targets) == result
diff --git a/test/fixtures/rules/std_rule_cases/L026.yml b/test/fixtures/rules/std_rule_cases/L026.yml
--- a/test/fixtures/rules/std_rule_cases/L026.yml
+++ b/test/fixtures/rules/std_rule_cases/L026.yml
@@ -110,3 +110,64 @@ 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_pass_two_part_reference_8:
+  pass_str: |
+    delete from public.agent1
+    where exists(
+        select customer.cust_id from customer
+        where agent1.agent_code <> customer.agent_code)
+
+test_pass_two_part_reference_9:
+  pass_str: |
+    delete from public.agent1
+    where exists(
+        select customer.cust_id from customer
+        where public.agent1.agent_code <> customer.agent_code)
+
+test_fail_two_part_reference_10:
+  fail_str: |
+    select *
+    from schema1.agent1
+    where schema2.agent1.agent_code <> 'abc'
+
+test_fail_two_part_reference_11:
+  fail_str: |
+    delete from schema1.agent1
+    where exists(
+        select customer.cust_id from customer
+        where schema2.agent1.agent_code <> customer.agent_code)
+
+test_pass_two_part_reference_11:
+  pass_str: |
+    select * from agent1
+    where public.agent1.agent_code <> '3'