feat: skip lineage analysis for function related statement (#501)

reata · Dec 11, 2023 · fe3e2c5 · fe3e2c5
1 parent 3874881
commit fe3e2c5
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 0 deletions.
diff --git a/sqllineage/core/parser/sqlfluff/extractors/noop.py b/sqllineage/core/parser/sqlfluff/extractors/noop.py
@@ -17,9 +17,14 @@ class NoopExtractor(BaseExtractor):
         "cache_table",
         "uncache_table",
         "show_statement",
+        "describe_statement",
         "use_statement",
         "declare_segment",
         "analyze_statement",
+        "add_jar_statement",
+        "create_function_statement",
+        "drop_function_statement",
+        "set_statement",
     ]
 
     def extract(

diff --git a/tests/sql/table/other_single_statement/test_other_without_lineage_dialect_specific.py b/tests/sql/table/other_single_statement/test_other_without_lineage_dialect_specific.py
@@ -45,6 +45,74 @@ def test_show_create_table(dialect: str):
     assert_table_lineage_equal("show create table tab1", dialect=dialect)
 
 
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_add_jar(dialect: str):
+    assert_table_lineage_equal("ADD JAR /tmp/SimpleUdf.jar", dialect=dialect)
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_create_function(dialect: str):
+    assert_table_lineage_equal(
+        """CREATE FUNCTION simple_udf AS 'SimpleUdf'
+  USING JAR '/tmp/SimpleUdf.jar'""",
+        dialect=dialect,
+    )
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_create_or_replace_function(dialect: str):
+    assert_table_lineage_equal(
+        """CREATE OR REPLACE FUNCTION simple_udf AS 'SimpleUdf'
+      USING JAR '/tmp/SimpleUdf.jar'""",
+        dialect=dialect,
+    )
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_create_temporary_function(dialect: str):
+    assert_table_lineage_equal(
+        """CREATE TEMPORARY FUNCTION simple_temp_udf AS 'SimpleUdf'
+  USING JAR '/tmp/SimpleUdf.jar'""",
+        dialect=dialect,
+    )
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_create_or_replace_temporary_function(dialect: str):
+    assert_table_lineage_equal(
+        """CREATE OR REPLACE TEMPORARY FUNCTION simple_temp_udf AS 'SimpleUdf'
+  USING JAR '/tmp/SimpleUdf.jar'""",
+        dialect=dialect,
+    )
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_show_functions(dialect: str):
+    assert_table_lineage_equal("SHOW FUNCTIONS trim", dialect=dialect)
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_describe_function(dialect: str):
+    assert_table_lineage_equal("DESC FUNCTION abs", dialect=dialect)
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_drop_function(dialect: str):
+    assert_table_lineage_equal("DROP FUNCTION test_avg", dialect=dialect)
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "hive", "sparksql"])
+def test_set_command(dialect: str):
+    assert_table_lineage_equal(
+        "SET spark.sql.variable.substitute=false", dialect=dialect
+    )
+
+
+@pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
+def test_set_command_without_property_value(dialect: str):
+    assert_table_lineage_equal("SET spark.sql.variable.substitute", dialect=dialect)
+
+
 @pytest.mark.parametrize("dialect", ["postgres", "redshift"])
 def test_analyze_table(dialect: str):
     assert_table_lineage_equal("analyze tab", dialect=dialect)