From a7044b672667f2a0b48bd0b326b5a55b0815ef79 Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Sat, 4 May 2024 18:13:01 +0200 Subject: Merging upstream version 23.13.1. Signed-off-by: Daniel Baumann --- tests/test_lineage.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 60 insertions(+), 4 deletions(-) (limited to 'tests/test_lineage.py') diff --git a/tests/test_lineage.py b/tests/test_lineage.py index c782d9a..3e17f95 100644 --- a/tests/test_lineage.py +++ b/tests/test_lineage.py @@ -224,16 +224,50 @@ class TestLineage(unittest.TestCase): downstream.source.sql(dialect="snowflake"), "LATERAL FLATTEN(INPUT => TEST_TABLE.RESULT, OUTER => TRUE) AS FLATTENED(SEQ, KEY, PATH, INDEX, VALUE, THIS)", ) - self.assertEqual( - downstream.expression.sql(dialect="snowflake"), - "VALUE", - ) + self.assertEqual(downstream.expression.sql(dialect="snowflake"), "VALUE") self.assertEqual(len(downstream.downstream), 1) downstream = downstream.downstream[0] self.assertEqual(downstream.name, "TEST_TABLE.RESULT") self.assertEqual(downstream.source.sql(dialect="snowflake"), "TEST_TABLE AS TEST_TABLE") + node = lineage( + "FIELD", + "SELECT FLATTENED.VALUE:field::text AS FIELD FROM SNOWFLAKE.SCHEMA.MODEL AS MODEL_ALIAS, LATERAL FLATTEN(INPUT => MODEL_ALIAS.A) AS FLATTENED", + schema={"SNOWFLAKE": {"SCHEMA": {"TABLE": {"A": "integer"}}}}, + sources={"SNOWFLAKE.SCHEMA.MODEL": "SELECT A FROM SNOWFLAKE.SCHEMA.TABLE"}, + dialect="snowflake", + ) + self.assertEqual(node.name, "FIELD") + + downstream = node.downstream[0] + self.assertEqual(downstream.name, "FLATTENED.VALUE") + self.assertEqual( + downstream.source.sql(dialect="snowflake"), + "LATERAL FLATTEN(INPUT => MODEL_ALIAS.A) AS FLATTENED(SEQ, KEY, PATH, INDEX, VALUE, THIS)", + ) + self.assertEqual(downstream.expression.sql(dialect="snowflake"), "VALUE") + self.assertEqual(len(downstream.downstream), 1) + + downstream = downstream.downstream[0] + self.assertEqual(downstream.name, "MODEL_ALIAS.A") + self.assertEqual(downstream.source_name, "SNOWFLAKE.SCHEMA.MODEL") + self.assertEqual( + downstream.source.sql(dialect="snowflake"), + "SELECT TABLE.A AS A FROM SNOWFLAKE.SCHEMA.TABLE AS TABLE", + ) + self.assertEqual(downstream.expression.sql(dialect="snowflake"), "TABLE.A AS A") + self.assertEqual(len(downstream.downstream), 1) + + downstream = downstream.downstream[0] + self.assertEqual(downstream.name, "TABLE.A") + self.assertEqual( + downstream.source.sql(dialect="snowflake"), "SNOWFLAKE.SCHEMA.TABLE AS TABLE" + ) + self.assertEqual( + downstream.expression.sql(dialect="snowflake"), "SNOWFLAKE.SCHEMA.TABLE AS TABLE" + ) + def test_subquery(self) -> None: node = lineage( "output", @@ -266,6 +300,7 @@ class TestLineage(unittest.TestCase): self.assertEqual(node.name, "a") node = node.downstream[0] self.assertEqual(node.name, "cte.a") + self.assertEqual(node.reference_node_name, "cte") node = node.downstream[0] self.assertEqual(node.name, "z.a") @@ -304,6 +339,27 @@ class TestLineage(unittest.TestCase): node = a.downstream[0] self.assertEqual(node.name, "foo.a") + # Select from derived table + node = lineage( + "a", + "SELECT a FROM (SELECT a FROM x) subquery", + ) + self.assertEqual(node.name, "a") + self.assertEqual(len(node.downstream), 1) + node = node.downstream[0] + self.assertEqual(node.name, "subquery.a") + self.assertEqual(node.reference_node_name, "subquery") + + node = lineage( + "a", + "SELECT a FROM (SELECT a FROM x)", + ) + self.assertEqual(node.name, "a") + self.assertEqual(len(node.downstream), 1) + node = node.downstream[0] + self.assertEqual(node.name, "_q_0.a") + self.assertEqual(node.reference_node_name, "_q_0") + def test_lineage_cte_union(self) -> None: query = """ WITH dataset AS ( -- cgit v1.2.3