Fixes #21027: [Snowflake] Add snowflake stages lineage support (#25944)

Author: mohittilala

ingestion/src/metadata/ingestion/lineage/parser.py

@@ -17,12 +17,12 @@
 from collections import defaultdict
 from copy import deepcopy
 from logging.config import DictConfigurator
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import sqlparse
 from cached_property import cached_property
 from collate_sqllineage import SQLPARSE_DIALECT
-from collate_sqllineage.core.models import Column, DataFunction, Table
+from collate_sqllineage.core.models import Column, DataFunction, Location, Table
 from collate_sqllineage.core.parser.sqlfluff.analyzer import SqlFluffLineageAnalyzer
 from collate_sqllineage.core.parser.sqlglot.analyzer import SqlGlotLineageAnalyzer
 from collate_sqllineage.exceptions import SQLLineageException
@@ -164,7 +164,7 @@ def intermediate_tables(self) -> List[Table]:
         return []
 
     @cached_property
-    def source_tables(self) -> List[Union[Table, DataFunction]]:
+    def source_tables(self) -> List[Union[Table, DataFunction, Location]]:
         """
         Get a list of source tables
         """
@@ -174,7 +174,7 @@ def source_tables(self) -> List[Union[Table, DataFunction]]:
         return []
 
     @cached_property
-    def target_tables(self) -> List[Table]:
+    def target_tables(self) -> List[Union[Table, Location]]:
         """
         Get a list of target tables
         """
@@ -485,13 +485,15 @@ def table_joins(self) -> Dict[str, List[TableColumnJoin]]:
 
         return join_data
 
-    def retrieve_tables(self, tables: List[Any]) -> List[Table]:
+    def retrieve_tables(
+        self, tables: List[Union[Table, DataFunction, Location]]
+    ) -> List[Union[Table, DataFunction, Location]]:
         if not self._clean_query:
             return []
         return [
             self.clean_table_name(table)
             for table in tables
-            if isinstance(table, (Table, DataFunction))
+            if isinstance(table, (Table, DataFunction, Location))
         ]
 
     @classmethod
@@ -520,9 +522,11 @@ def clean_raw_query(cls, raw_query: str) -> Optional[str]:
 
         # We remove queries of the type 'COPY table FROM path' since they are data manipulation statements
         # that do not provide value for user. However, we keep Snowflake 'COPY INTO table FROM @stage'
-        # statements as they provide lineage from stages to tables.
-        if insensitive_match(clean_query, r"^COPY\s+") and not insensitive_match(
-            clean_query, r"^COPY\s+INTO\s+.*\s+FROM\s+@"
+        # and 'COPY INTO @stage FROM table' as they are used for loading/unloading data and are relevant
+        # for lineage from stages to tables or vice versa.
+        if insensitive_match(clean_query, r"^COPY\s+") and not (
+            insensitive_match(clean_query, r"^COPY\s+INTO\s+.*\s+FROM\s+@")
+            or insensitive_match(clean_query, r"^COPY\s+INTO\s+@.*\s+FROM\s+.*")
         ):
             return None
 
@@ -739,7 +743,9 @@ def get_sqlparse_lineage_runner(query: str) -> LineageRunner:
         return None
 
     @staticmethod
-    def clean_table_name(table: Table) -> Table:
+    def clean_table_name(
+        table: Union[Table, DataFunction, Location],
+    ) -> Union[Table, DataFunction, Location]:
         """
         Clean table name by:
         - Removing brackets from the beginning and end of the table and schema name
@@ -761,4 +767,15 @@ def clean_table_name(table: Table) -> Table:
             clean_table.schema.raw_name = insensitive_replace(
                 clean_table.schema.raw_name, r"\[(.*)\]", r"\1"
             )
+        # Remove leading @ from the location storage objects if present as they are
+        # not used while ingesting location storage objects in OpenMetadata
+        # ex. @STAGE_01 -> STAGE_01 (snowflake stage object)
+        if (
+            isinstance(clean_table, Location)
+            and clean_table.raw_name
+            and insensitive_match(clean_table.raw_name, r"@.*")
+        ):
+            clean_table.raw_name = insensitive_replace(
+                clean_table.raw_name, r"@(.*)", r"\1"
+            )
         return clean_table

ingestion/src/metadata/ingestion/source/database/snowflake/lineage.py

@@ -45,7 +45,7 @@ class SnowflakeLineageSource(
 
     filters = """
         AND (
-            QUERY_TYPE IN ('MERGE', 'UPDATE','CREATE_TABLE_AS_SELECT','COPY')
+            QUERY_TYPE IN ('MERGE', 'UPDATE','CREATE_TABLE_AS_SELECT','COPY','UNLOAD')
             OR (QUERY_TYPE = 'INSERT' and query_text ILIKE '%%insert%%into%%select%%')
             OR (QUERY_TYPE = 'ALTER' and query_text ILIKE '%%alter%%table%%swap%%')
             OR (QUERY_TYPE = 'CREATE_TABLE' and query_text ILIKE '%%clone%%')

ingestion/tests/unit/lineage/queries/test_specific_dialect_queries.py

@@ -105,7 +105,7 @@
 from unittest import TestCase
 
 import pytest
-from collate_sqllineage.core.models import Path
+from collate_sqllineage.core.models import Location, Path
 
 from ingestion.tests.unit.lineage.queries.helpers import (
     TestColumnQualifierTuple,
@@ -1090,3 +1090,83 @@ def test_oracle_create_procedure_insert_select(self):
             test_sqlfluff=False,
             test_sqlparse=False,
         )
+
+    # -------------------------------------------------------------------------
+    # Snowflake Stage Lineage Tests (COPY INTO @stage / COPY INTO table FROM @stage)
+    # -------------------------------------------------------------------------
+
+    def test_snowflake_copy_into_table_from_stage(self):
+        """Test Snowflake COPY INTO table FROM @stage (loading data from stage)"""
+        query = "COPY INTO wine_quality FROM @demo FILE_FORMAT = wine_csv_format;"
+
+        assert_table_lineage_equal(
+            query,
+            {Location("@demo")},
+            {"wine_quality"},
+            dialect=Dialect.SNOWFLAKE.value,
+        )
+
+        # No column lineage expected for COPY INTO stage operations
+        assert_column_lineage_equal(
+            query,
+            [],
+            dialect=Dialect.SNOWFLAKE.value,
+        )
+
+    def test_snowflake_copy_into_stage_from_table(self):
+        """Test Snowflake COPY INTO @stage FROM table (unloading data to stage)"""
+        query = "COPY INTO @my_stage FROM my_table"
+
+        assert_table_lineage_equal(
+            query,
+            {"my_table"},
+            {Location("@my_stage")},
+            dialect=Dialect.SNOWFLAKE.value,
+        )
+
+        assert_column_lineage_equal(
+            query,
+            [],
+            dialect=Dialect.SNOWFLAKE.value,
+        )
+
+    def test_snowflake_copy_into_stage_from_select(self):
+        """Test Snowflake COPY INTO @stage FROM (SELECT ...) - unload with subquery"""
+        query = "COPY INTO @db.schema.my_stage FROM (SELECT col1, col2 FROM my_table)"
+
+        assert_table_lineage_equal(
+            query,
+            {"my_table"},
+            {Location("@db.schema.my_stage")},
+            dialect=Dialect.SNOWFLAKE.value,
+            # SqlParse builds a different internal graph structure for subqueries
+            # (5 nodes vs 3 nodes) even though final lineage results are identical
+            skip_graph_check=True,
+        )
+
+        assert_column_lineage_equal(
+            query,
+            [],
+            dialect=Dialect.SNOWFLAKE.value,
+            # SqlGlot and SqlFluff/SqlParse produce different graph structures for subqueries
+            skip_graph_check=True,
+        )
+
+    def test_snowflake_copy_into_fully_qualified_stage(self):
+        """Test COPY INTO table FROM @db.schema.stage with fully qualified stage name"""
+        query = (
+            "COPY INTO my_table FROM @my_db.my_schema.my_stage FILE_FORMAT=(TYPE=CSV)"
+        )
+
+        assert_table_lineage_equal(
+            query,
+            {Location("@my_db.my_schema.my_stage")},
+            {"my_table"},
+            dialect=Dialect.SNOWFLAKE.value,
+        )
+
+        assert_column_lineage_equal(
+            query,
+            [],
+            dialect=Dialect.SNOWFLAKE.value,
+        )

ingestion/tests/unit/lineage/test_sql_lineage.py

@@ -16,6 +16,8 @@
 from unittest import TestCase
 
 import pytest
+from collate_sqllineage.core.models import Location
+from collate_sqllineage.core.models import Table as LineageTable
 
 from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.lineage.models import Dialect
@@ -335,3 +337,130 @@ def test_copy_to_statements_filtered(self):
                 result,
                 f"Query should be filtered: {query}",
             )
+
+    # -------------------------------------------------------------------------
+    # Snowflake Stage Lineage Tests
+    # -------------------------------------------------------------------------
+
+    def test_copy_into_stage_from_table_not_filtered(self):
+        """
+        Test that Snowflake COPY INTO @stage FROM table (unload) statements
+        are NOT filtered out, as they provide lineage from tables to stages.
+        """
+        snowflake_unload_queries = [
+            "COPY INTO @my_stage FROM my_table",
+            "COPY INTO @db.schema.stage FROM (SELECT * FROM t)",
+            "copy into @stage/path FROM table1",
+            "COPY INTO @~/ FROM my_table FILE_FORMAT = (TYPE = CSV COMPRESSION = GZIP)",
+            "COPY INTO @~/staged FROM sales_data",
+            "COPY INTO @my_stage/daily/2024/ FROM reporting.public.daily_metrics",
+            "COPY INTO @external_stage/path/ FROM (SELECT col1 FROM src_table WHERE id > 100)",
+        ]
+
+        for query in snowflake_unload_queries:
+            result = LineageParser.clean_raw_query(query)
+            self.assertIsNotNone(
+                result,
+                f"COPY INTO @stage FROM table should NOT be filtered: {query}",
+            )
+
+    def test_stage_lineage_source_as_location_type(self):
+        """
+        Verify that COPY INTO table FROM @stage returns Location as source
+        and Table as target with correct types.
+        """
+        query = "COPY INTO wine_quality FROM @demo FILE_FORMAT = wine_csv_format;"
+        parser = LineageParser(query, dialect=Dialect.SNOWFLAKE)
+
+        self.assertEqual(len(parser.source_tables), 1)
+        self.assertEqual(len(parser.target_tables), 1)
+        self.assertIsInstance(parser.source_tables[0], Location)
+        self.assertNotIsInstance(parser.source_tables[0], LineageTable)
+
+    def test_stage_lineage_target_as_location_type(self):
+        """
+        Verify that COPY INTO @stage FROM table returns Table as source
+        and Location as target with correct types.
+        """
+        query = "COPY INTO @my_stage FROM my_table"
+        parser = LineageParser(query, dialect=Dialect.SNOWFLAKE)
+
+        self.assertEqual(len(parser.source_tables), 1)
+        self.assertEqual(len(parser.target_tables), 1)
+        self.assertNotIsInstance(parser.target_tables[0], LineageTable)
+        self.assertIsInstance(parser.target_tables[0], Location)
+
+    def test_stage_lineage_fully_qualified_names(self):
+        """
+        Test stage lineage with fully qualified database.schema.stage names
+        for both source and target directions.
+        """
+        # Stage as source (loading data into table)
+        query_load = "COPY INTO db.schema.target_table FROM @db.schema.my_stage"
+        parser_load = LineageParser(query_load, dialect=Dialect.SNOWFLAKE)
+
+        self.assertEqual(len(parser_load.source_tables), 1)
+        self.assertEqual(len(parser_load.target_tables), 1)
+        self.assertIsInstance(parser_load.source_tables[0], Location)
+        self.assertEqual(str(parser_load.source_tables[0]), "db.schema.my_stage")
+        self.assertEqual(str(parser_load.target_tables[0]), "db.schema.target_table")
+
+        # Stage as target (unloading data from table)
+        query_unload = "COPY INTO @db.schema.my_stage FROM db.schema.source_table"
+        parser_unload = LineageParser(query_unload, dialect=Dialect.SNOWFLAKE)
+
+        self.assertEqual(len(parser_unload.source_tables), 1)
+        self.assertEqual(len(parser_unload.target_tables), 1)
+        self.assertIsInstance(parser_unload.target_tables[0], Location)
+        self.assertEqual(str(parser_unload.source_tables[0]), "db.schema.source_table")
+        self.assertEqual(str(parser_unload.target_tables[0]), "db.schema.my_stage")
+
+    def test_stage_lineage_unload_with_select_subquery(self):
+        """
+        Test COPY INTO @stage FROM (SELECT ...) extracts the underlying
+        source table correctly from the subquery.
+        """
+        query = (
+            "COPY INTO @external_stage/path/ FROM "
+            "(SELECT col1, col2 FROM db.schema.source_table WHERE id > 100)"
+        )
+        parser = LineageParser(query, dialect=Dialect.SNOWFLAKE)
+
+        self.assertEqual(len(parser.source_tables), 1)
+        self.assertEqual(len(parser.target_tables), 1)
+        self.assertEqual(str(parser.source_tables[0]), "db.schema.source_table")
+        self.assertIsInstance(parser.target_tables[0], Location)
+
+    def test_stage_lineage_user_stage(self):
+        """
+        Test COPY INTO with user stage (@~/) is properly handled.
+        """
+        query = "COPY INTO @~/ FROM my_table FILE_FORMAT = (TYPE = CSV)"
+        parser = LineageParser(query, dialect=Dialect.SNOWFLAKE)
+
+        self.assertEqual(len(parser.source_tables), 1)
+        self.assertEqual(len(parser.target_tables), 1)
+        self.assertIsInstance(parser.source_tables[0], LineageTable)
+        self.assertIsInstance(parser.target_tables[0], Location)
+
+    def test_stage_lineage_with_file_format_options(self):
+        """
+        Test that file format options don't interfere with lineage parsing.
+        """
+        queries = [
+            "COPY INTO my_table FROM @stage FILE_FORMAT = (TYPE = CSV SKIP_HEADER = 1)",
+            "COPY INTO my_table FROM @stage FILE_FORMAT = wine_csv_format",
+            "COPY INTO @stage FROM my_table FILE_FORMAT = (TYPE = PARQUET)",
+            "COPY INTO my_table FROM @stage PATTERN='.*[.]csv'",
+        ]
+
+        for query in queries:
+            parser = LineageParser(query, dialect=Dialect.SNOWFLAKE)
+            self.assertTrue(
+                len(parser.source_tables) > 0,
+                f"Expected source tables for query: {query}",
+            )
+            self.assertTrue(
+                len(parser.target_tables) > 0,
+                f"Expected target tables for query: {query}",
+            )