Repository: astronomer/astro-sdk

Issue:
Change `export_file` to return `File` object
**Context**

In order to allow users to perform subsequent actions on an exported file (while maintaining a functional structure), we should allow the `export_file` function to return a file object.

* Astro-SDK version: 0.9..1
* Request by: @jlaneve
* Analysed by @dimberman 

**Problem**

At the moment a user who wants to use the `output_file` object would need to explicitly set dependencies like this:

```
    output_file = File(path="/tmp/saved_df.csv")
    with sample_dag:
        table = aql.load_file(input_file=File(path=data_path), output_table=test_table)
        export = aql.export_file(
            input_data=table,
            output_file=output_file,
            if_exists="replace",
        )
        res_df = aql.load_file(input_file=output_file)
        export >> res_df
```

**Desired behaviour**

```
    with sample_dag:
        table = aql.load_file(input_file=File(path=data_path), output_table=test_table)
        exported_file = aql.export_file(
            input_data=table,
            output_file=File(path="/tmp/saved_df.csv"),
            if_exists="replace",
        )
        res_df = aql.load_file(input_file=exported_file)
```

**Acceptance criteria**
* Change `export_file` so it returns the `File` instance, as opposed to `None`

Since there is no documentation about this task, we don't need to update the documentation for it. To create documentation for this feature should be part of another issue.


Discussion:

Resolving patch:
diff --git a/src/astro/sql/operators/export_file.py b/src/astro/sql/operators/export_file.py
--- a/src/astro/sql/operators/export_file.py
+++ b/src/astro/sql/operators/export_file.py
@@ -34,7 +34,7 @@ def __init__(
         self.if_exists = if_exists
         self.kwargs = kwargs
 
-    def execute(self, context: dict) -> None:
+    def execute(self, context: dict) -> File:
         """Write SQL table to csv/parquet on local/S3/GCS.
 
         Infers SQL database type based on connection.
@@ -53,6 +53,7 @@ def execute(self, context: dict) -> None:
         # Write file if overwrite == True or if file doesn't exist.
         if self.if_exists == "replace" or not self.output_file.exists():
             self.output_file.create_from_dataframe(df)
+            return self.output_file
         else:
             raise FileExistsError(f"{self.output_file.path} file already exists.")
 
@@ -66,7 +67,20 @@ def export_file(
 ) -> XComArg:
     """Convert SaveFile into a function. Returns XComArg.
 
-    Returns an XComArg object.
+    Returns an XComArg object of type File which matches the output_file parameter.
+
+    This will allow users to perform further actions with the exported file.
+
+    e.g.
+
+        with sample_dag:
+        table = aql.load_file(input_file=File(path=data_path), output_table=test_table)
+        exported_file = aql.export_file(
+            input_data=table,
+            output_file=File(path="/tmp/saved_df.csv"),
+            if_exists="replace",
+        )
+        res_df = aql.load_file(input_file=exported_file)
 
     :param output_file: Path and conn_id
     :param input_data: Input table / dataframe


Test patch:
diff --git a/tests/sql/operators/test_export_file.py b/tests/sql/operators/test_export_file.py
--- a/tests/sql/operators/test_export_file.py
+++ b/tests/sql/operators/test_export_file.py
@@ -88,6 +88,33 @@ def test_save_temp_table_to_local(sample_dag, database_table_fixture):
     assert input_df.equals(output_df)
 
 
+@pytest.mark.parametrize(
+    "database_table_fixture", [{"database": Database.SQLITE}], indirect=True
+)
+def test_save_returns_output_file(sample_dag, database_table_fixture):
+    _, test_table = database_table_fixture
+
+    @aql.dataframe
+    def validate(df: pd.DataFrame):
+        assert not df.empty
+
+    data_path = str(CWD) + "/../../data/homes.csv"
+    with sample_dag:
+        table = aql.load_file(input_file=File(path=data_path), output_table=test_table)
+        file = aql.export_file(
+            input_data=table,
+            output_file=File(path="/tmp/saved_df.csv"),
+            if_exists="replace",
+        )
+        res_df = aql.load_file(input_file=file)
+        validate(res_df)
+    test_utils.run_dag(sample_dag)
+
+    output_df = pd.read_csv("/tmp/saved_df.csv")
+    input_df = pd.read_csv(data_path)
+    assert input_df.equals(output_df)
+
+
 @pytest.mark.parametrize(
     "database_table_fixture",
     [