Repository: astronomer/astro-sdk

Issue:
load_file & save_file from GCS/S3 to a Pandas data-frame
Acceptance criteria
- A user should be able to load a file from GCS/S3 into a dataframe without needing any SQL database connection.
- does not need to test with large scale data
- user should also be able to save the dataframe directly


Discussion:

Resolving patch:
diff --git a/example_dags/example_amazon_s3_postgres_load_and_save.py b/example_dags/example_amazon_s3_postgres_load_and_save.py
--- a/example_dags/example_amazon_s3_postgres_load_and_save.py
+++ b/example_dags/example_amazon_s3_postgres_load_and_save.py
@@ -45,10 +45,12 @@ def example_amazon_s3_postgres_load_and_save():
     )
 
     aql.save_file(
+        input=t1,
         output_file_path="s3://tmp9/homes.csv",
-        input_table=t1,
         overwrite=True,
     )
 
 
-dag = example_amazon_s3_postgres_load_and_save()
+example_amazon_s3_postgres_load_and_save_dag = (
+    example_amazon_s3_postgres_load_and_save()
+)
diff --git a/src/astro/sql/operators/agnostic_save_file.py b/src/astro/sql/operators/agnostic_save_file.py
--- a/src/astro/sql/operators/agnostic_save_file.py
+++ b/src/astro/sql/operators/agnostic_save_file.py
@@ -14,8 +14,7 @@
 limitations under the License.
 """
 
-import os
-from typing import Optional
+from typing import Optional, Union
 from urllib.parse import urlparse
 
 import pandas as pd
@@ -47,7 +46,7 @@ class SaveFile(BaseOperator):
     """
 
     template_fields = (
-        "input_table",
+        "input",
         "output_file_path",
         "output_conn_id",
         "output_file_format",
@@ -55,7 +54,7 @@ class SaveFile(BaseOperator):
 
     def __init__(
         self,
-        input_table: Table = None,
+        input: Optional[Union[Table, pd.DataFrame]] = None,
         output_file_path="",
         output_conn_id=None,
         output_file_format="csv",
@@ -64,7 +63,7 @@ def __init__(
     ) -> None:
         super().__init__(**kwargs)
         self.output_file_path = output_file_path
-        self.input_table = input_table
+        self.input = input
         self.output_conn_id = output_conn_id
         self.overwrite = overwrite
         self.output_file_format = output_file_format
@@ -77,7 +76,40 @@ def execute(self, context):
         """
 
         # Infer db type from `input_conn_id`.
-        input_table = self.input_table
+        if type(self.input) == Table:
+            df = self.convert_sql_table_to_dataframe()
+        elif type(self.input) == pd.DataFrame:
+            df = self.input
+        else:
+            raise ValueError(
+                "Expected input_table to be Table or dataframe. Got %s",
+                type(self.input),
+            )
+
+        # Write file if overwrite == True or if file doesn't exist.
+        if self.overwrite == True or not self.file_exists(
+            self.output_file_path, self.output_conn_id
+        ):
+            self.agnostic_write_file(df, self.output_file_path, self.output_conn_id)
+        else:
+            raise FileExistsError
+
+    def file_exists(self, output_file_path, output_conn_id=None):
+        transport_params = {
+            "s3": s3fs_creds,
+            "gs": gcs_client,
+            "": lambda: None,
+        }[urlparse(output_file_path).scheme]()
+        try:
+            with smart_open.open(
+                output_file_path, mode="r", transport_params=transport_params
+            ):
+                return True
+        except IOError:
+            return False
+
+    def convert_sql_table_to_dataframe(self):
+        input_table = self.input
         conn_type = BaseHook.get_connection(input_table.conn_id).conn_type
 
         # Select database Hook based on `conn` type
@@ -117,33 +149,11 @@ def execute(self, context):
         else:
             table_name = f"{get_table_name(input_table)}"
         # Load table from SQL db.
-        df = pd.read_sql(
+        return pd.read_sql(
             f"SELECT * FROM {table_name}",
             con=input_hook.get_sqlalchemy_engine(),
         )
 
-        # Write file if overwrite == True or if file doesn't exist.
-        if self.overwrite == True or not self.file_exists(
-            self.output_file_path, self.output_conn_id
-        ):
-            self.agnostic_write_file(df, self.output_file_path, self.output_conn_id)
-        else:
-            raise FileExistsError
-
-    def file_exists(self, output_file_path, output_conn_id=None):
-        transport_params = {
-            "s3": s3fs_creds,
-            "gs": gcs_client,
-            "": lambda: None,
-        }[urlparse(output_file_path).scheme]()
-        try:
-            with smart_open.open(
-                output_file_path, mode="r", transport_params=transport_params
-            ):
-                return True
-        except IOError:
-            return False
-
     def agnostic_write_file(self, df, output_file_path, output_conn_id=None):
         """Write dataframe to csv/parquet files formats
 
@@ -182,7 +192,7 @@ def create_table_name(context):
 
 def save_file(
     output_file_path,
-    input_table=None,
+    input=None,
     output_conn_id=None,
     overwrite=False,
     output_file_format="csv",
@@ -214,7 +224,7 @@ def save_file(
     return SaveFile(
         task_id=task_id,
         output_file_path=output_file_path,
-        input_table=input_table,
+        input=input,
         output_conn_id=output_conn_id,
         overwrite=overwrite,
         output_file_format=output_file_format,


Test patch:
diff --git a/tests/data/output.csv b/tests/data/output.csv
new file mode 100644
diff --git a/tests/operators/test_agnostic_save_file.py b/tests/operators/test_agnostic_save_file.py
--- a/tests/operators/test_agnostic_save_file.py
+++ b/tests/operators/test_agnostic_save_file.py
@@ -43,6 +43,7 @@
 from airflow.utils.state import State
 from airflow.utils.types import DagRunType
 
+import astro.dataframe as adf
 import astro.sql as aql
 
 # Import Operator
@@ -117,6 +118,22 @@ def create_and_run_tasks(self, decorator_funcs):
         test_utils.run_dag(self.dag)
         return tasks
 
+    def test_save_dataframe_to_local(self):
+        @adf
+        def make_df():
+            d = {"col1": [1, 2], "col2": [3, 4]}
+            return pd.DataFrame(data=d)
+
+        with self.dag:
+            df = make_df()
+            aql.save_file(
+                input=df, output_file_path="/tmp/saved_df.csv", overwrite=True
+            )
+
+        test_utils.run_dag(self.dag)
+        df = pd.read_csv("/tmp/saved_df.csv")
+        assert df.equals(pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}))
+
     def test_save_postgres_table_to_local_file_exists_overwrite_false(self):
 
         OUTPUT_FILE_PATH = str(self.cwd) + "/../data/save_file_out.csv"
@@ -134,7 +151,7 @@ def run_task():
                 save_file,
                 (),
                 {
-                    "input_table": Table(
+                    "input": Table(
                         INPUT_TABLE_NAME,
                         conn_id="postgres_conn",
                         database="pagila",
@@ -171,7 +188,7 @@ def test_save_postgres_table_to_s3(self):
             save_file,
             (),
             {
-                "input_table": Table(
+                "input": Table(
                     INPUT_TABLE_NAME,
                     conn_id="postgres_conn",
                     database="pagila",
@@ -213,7 +230,7 @@ def run_task():
                 save_file,
                 (),
                 {
-                    "input_table": Table(
+                    "input": Table(
                         INPUT_TABLE_NAME,
                         conn_id="postgres_conn",
                         database="pagila",
@@ -242,7 +259,7 @@ def test_unique_task_id_for_same_path(self):
                     "func": save_file,
                     "op_args": (),
                     "op_kwargs": {
-                        "input_table": Table(
+                        "input": Table(
                             INPUT_TABLE_NAME,
                             conn_id="postgres_conn",
                             database="pagila",
@@ -285,7 +302,7 @@ def test_save_bigquery_table_to_local_file_exists_overwrite_false(self):
             save_file,
             (),
             {
-                "input_table": Table(
+                "input": Table(
                     INPUT_TABLE_NAME,
                     conn_id="bigquery",
                     schema=test_utils.DEFAULT_SCHEMA,
@@ -323,7 +340,7 @@ def test_save_sqlite_table_to_local_file_exists_overwrite_false(self):
             save_file,
             (),
             {
-                "input_table": Table(INPUT_TABLE_NAME, conn_id="sqlite_conn"),
+                "input": Table(INPUT_TABLE_NAME, conn_id="sqlite_conn"),
                 "output_file_path": OUTPUT_FILE_PATH,
                 "output_conn_id": None,
                 "overwrite": True,
@@ -410,7 +427,7 @@ def test_save_file(sample_dag, sql_server, file_type):
         filepath = Path(tmp_dir, f"sample.{file_type}")
 
         task_params = {
-            "input_table": Table(table_name=INPUT_TABLE_NAME, **sql_server_params),
+            "input": Table(table_name=INPUT_TABLE_NAME, **sql_server_params),
             "output_file_path": str(filepath),
             "output_file_format": file_type,
             "output_conn_id": None,