Repository: reata/sqllineage

Issue:
Support Redshift `Unload` Command
**Is your feature request related to a problem? Please describe.**
Similar to how Redshift imports data using `COPY` (supported via [Issue#164](https://github.com/reata/sqllineage/issues/164)), `UNLOAD` is when Redshift exports data.

Using this SQL example from [Redshift Docs](https://docs.aws.amazon.com/redshift/latest/dg/r_UNLOAD_command_examples.html):
```
unload ('select venueid, venuename, venueseats from venue')
to 's3://amzn-s3-demo-bucket/venue_tab_'
iam_role 'arn:aws:iam::0123456789012:role/MyRedshiftRole'
delimiter as '\t';
```

(There are a bunch of options that can occur after the `UNLOAD ({CTE}) TO s3://...`, but I do not believe those would be relevant for Lineage)

**Describe the solution you'd like**

```
Statement #1: unload ('select venueid, venuename, venueseats from ve...
    table read: [Table: <default>.venue]
    table write: [Path: s3://amzn-s3-demo-bucket/]
    table cte: []
    table drop: []
    table rename: []
==========
Summary:
Statements(#): 1
Source Tables:
    <default>.venue
Target Tables:
    s3://amzn-s3-demo-bucket/
```

**Describe alternatives you've considered**
N/A

**Additional context**
Unload's output files written is non-deterministic and is dependent on Threads, filesize, etc. The path is a prefix-path. I think cutting to the path parent folder would be the most straight forward approach.

From the docs:
https://docs.aws.amazon.com/redshift/latest/dg/r_UNLOAD_command_examples.html#unload-examples-venue

By default, UNLOAD writes one or more files per slice. 
```
unload ('select * from venue')
to 's3://amzn-s3-demo-bucket/unload/venue_pipe_'
iam_role 'arn:aws:iam::0123456789012:role/MyRedshiftRole';
```
The result is these four files in the unload folder, again assuming four slices.
```
venue_pipe_0000_part_00
venue_pipe_0001_part_00
venue_pipe_0002_part_00
venue_pipe_0003_part_00
```

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/723?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=reata) Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 99.56%. Comparing base ([`c4c9b4f`](https://app.codecov.io/gh/reata/sqllineage/commit/c4c9b4fc3121438c54430cef6abf9d790227eec9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=reata)) to head ([`1fe2635`](https://app.codecov.io/gh/reata/sqllineage/commit/1fe2635d1507ab2da9f195d05c2dda81c887b917?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=reata)).

<details><summary>Additional details and impacted files</summary>



```diff
@@           Coverage Diff           @@
##           master     #723   +/-   ##
=======================================
  Coverage   99.55%   99.56%           
=======================================
  Files          41       42    +1     
  Lines        2268     2289   +21     
=======================================
+ Hits         2258     2279   +21     
  Misses         10       10           
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/723?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=reata).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=reata).
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>

Fix (patch):
diff --git a/sqllineage/core/parser/sqlfluff/extractors/unload.py b/sqllineage/core/parser/sqlfluff/extractors/unload.py
new file mode 100644
index 0000000..0e638a5
--- /dev/null
+++ b/sqllineage/core/parser/sqlfluff/extractors/unload.py
@@ -0,0 +1,35 @@
+from sqlfluff.core.parser import BaseSegment
+
+from sqllineage.core.holders import StatementLineageHolder
+from sqllineage.core.models import Path
+from sqllineage.core.parser.sqlfluff.extractors.base import BaseExtractor
+from sqllineage.utils.entities import AnalyzerContext
+
+
+class UnloadExtractor(BaseExtractor):
+    """
+    Unload statement lineage extractor
+    """
+
+    SUPPORTED_STMT_TYPES = [
+        "unload_statement",
+    ]
+
+    def extract(
+        self, statement: BaseSegment, context: AnalyzerContext
+    ) -> StatementLineageHolder:
+        holder = StatementLineageHolder()
+        if bracketed := statement.get_child("bracketed"):
+            if quoted_literal := bracketed.get_child("quoted_literal"):
+                raw_sql = quoted_literal.raw.strip("'").strip('"')
+                holder |= self._analyze_inner_sql(raw_sql)
+        if quoted_literal := statement.get_child("quoted_literal"):
+            raw_path = quoted_literal.raw.strip("'").strip('"')
+            holder.add_write(Path(raw_path))
+        return holder
+
+    def _analyze_inner_sql(self, sql: str) -> StatementLineageHolder:
+        from sqllineage.core.parser.sqlfluff.analyzer import SqlFluffLineageAnalyzer
+
+        analyzer = SqlFluffLineageAnalyzer("", self.dialect)
+        return analyzer.analyze(sql, self.metadata_provider)


Test changes:
diff --git a/tests/sql/table/test_path_dialect_specific.py b/tests/sql/table/test_path_dialect_specific.py
index 5d7e160..2f4173b 100644
--- a/tests/sql/table/test_path_dialect_specific.py
+++ b/tests/sql/table/test_path_dialect_specific.py
@@ -67,3 +67,19 @@ SELECT * FROM tab1""",
         {Path("hdfs://path/to/folder")},
         dialect=dialect,
     )
+
+
+@pytest.mark.parametrize("dialect", ["redshift"])
+def test_redshift_unloading_data_to_s3(dialect: str):
+    """
+    https://docs.aws.amazon.com/redshift/latest/dg/t_Unloading_tables.html
+    """
+    assert_table_lineage_equal(
+        """unload ('select * from venue')   
+to 's3://amzn-s3-demo-bucket/tickit/unload/venue_' 
+iam_role 'arn:aws:iam::0123456789012:role/MyRedshiftRole';""",
+        {"venue"},
+        {Path("s3://amzn-s3-demo-bucket/tickit/unload/venue_")},
+        dialect=dialect,
+        test_sqlparse=False,
+    )