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/__init__.py b/src/astro/__init__.py
--- a/src/astro/__init__.py
+++ b/src/astro/__init__.py
@@ -1,6 +1,6 @@
 """A decorator that allows users to run SQL queries natively in Airflow."""
 
-__version__ = "0.9.1"
+__version__ = "0.9.2"
 
 # The following line is an import work-around to avoid raising a circular dependency issue related to `create_database`
 # Without this, if we run the following imports, in this specific order:
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
@@ -85,6 +85,43 @@ def test_save_temp_table_to_local(sample_dag, sql_server, test_table):
     assert input_df.equals(output_df)
 
 
+@pytest.mark.parametrize("sql_server", [Database.SQLITE.value], indirect=True)
+@pytest.mark.parametrize(
+    "test_table",
+    [
+        {
+            "path": str(CWD) + "/../../data/homes.csv",
+            "load_table": True,
+            "param": {
+                "name": test_utils.get_table_name("test_stats_check_1"),
+            },
+        }
+    ],
+    indirect=True,
+    ids=["temp_table"],
+)
+def test_save_returns_output_file(sample_dag, test_table, sql_server):
+    @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("sql_server", SUPPORTED_DATABASES, indirect=True)
 @pytest.mark.parametrize(
     "test_table",