Repository: astronomer/astro-sdk

Issue:
Unable to use astro.dataframe (without the table arg)
**Describe the bug**

Astro fails when the user attempts to transfer the output of a dataframe to a table. This issue is critical because it breaks Astro Build pipelines.

**Version**
* Astro: 0.6.0, 0.7.0 and main

**To Reproduce**
Steps to reproduce the behavior:

1. Add this to the test to the `TestDataframeFromSQL` class:

```
    def test_postgres_dataframe_without_table_arg(self):
        @df
        def validate_result(df: pandas.DataFrame):
            assert df.iloc[0].to_dict()["colors"] == "red"
   
        @df
        def sample_df():
            return pandas.DataFrame({"numbers": [1, 2, 3], "colors": ["red", "white", "blue"]})
   
        @aql.transform
        def sample_pg(input: Table):
            return "SELECT * FROM {{input}}"
   
        with self.dag:
            plain_df = sample_df()
            pg_df = sample_pg(conn_id="postgres_conn", database="pagila", input=plain_df)
            validate_result(pg_df)
        test_utils.run_dag(self.dag)
```

2. Run the tests:
```
$ PYTHONPATH=`pwd`/src/ pytest tests/operators/test_dataframe.py 
```

3. See the error:
```
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/pandas/io/sql.py", line 853, in _execute_create
    self.table.create()
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/sql/schema.py", line 928, in create
    bind._run_visitor(ddl.SchemaGenerator, self, checkfirst=checkfirst)
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 2095, in _run_visitor
    conn._run_visitor(visitorcallable, element, **kwargs)
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1656, in _run_visitor
    visitorcallable(self.dialect, self, **kwargs).traverse_single(element)
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/sql/visitors.py", line 145, in traverse_single
    return meth(obj, **kw)
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/sql/ddl.py", line 808, in visit_table
    if not create_ok and not self._can_create_table(table):
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/sql/ddl.py", line 732, in _can_create_table
    self.dialect.validate_identifier(table.name)
  File "/home/tati/.virtualenvs/astro-main/lib/python3.9/site-packages/sqlalchemy/engine/default.py", line 501, in validate_identifier
    raise exc.IdentifierError(
sqlalchemy.exc.IdentifierError: Identifier 'test_dag_sample_pg_backfill__2016-01-01T00:00:00+00:00_input_dataframe_input'
```

**Expected behavior**

Users should be able to use the dataframe decorator and the transform operation without raising exceptions.


**Reporters**
* @jlaneve @dimberman 


Discussion:

Resolving patch:
diff --git a/src/astro/sql/operators/sql_decorator.py b/src/astro/sql/operators/sql_decorator.py
--- a/src/astro/sql/operators/sql_decorator.py
+++ b/src/astro/sql/operators/sql_decorator.py
@@ -12,7 +12,7 @@
 
 from astro.constants import Database
 from astro.settings import SCHEMA
-from astro.sql.table import Table, TempTable, create_table_name
+from astro.sql.table import Table, TempTable, create_unique_table_name
 from astro.utils import get_hook, postgres_transform, snowflake_transform
 from astro.utils.database import (
     get_database_from_conn_id,
@@ -104,7 +104,6 @@ def execute(self, context: Dict):
         self._process_params()
 
         output_table_name: str = ""
-
         # Create DEFAULT SCHEMA if not created.
         self._set_schema_if_needed(SCHEMA)
 
@@ -116,7 +115,7 @@ def execute(self, context: Dict):
                 self._set_schema_if_needed()
 
             if not self.output_table or type(self.output_table) == TempTable:
-                output_table_name = create_table_name(context=context)
+                output_table_name = create_unique_table_name()
                 full_output_table_name = self.handle_output_table_schema(
                     # Since there is no output table defined we have to assume default schema
                     output_table_name,
@@ -175,7 +174,6 @@ def read_sql(self):
                 self.sql = file.read().replace("\n", " ")
 
     def handle_params(self, context):
-        # Automatically add any kwargs going into the function
         if self.op_kwargs:
             self.parameters.update(self.op_kwargs)  # type: ignore
         if self.op_args:
@@ -295,7 +293,7 @@ def _add_templates_to_context(self, context):
 
     def default_transform(self, parameters, context):
         for k, v in parameters.items():
-            if type(v) == Table:
+            if isinstance(v, Table):
                 context[k] = v.table_name
             else:
                 context[k] = ":" + k
@@ -311,9 +309,7 @@ def convert_op_arg_dataframes(self):
         for i, arg in enumerate(self.op_args):
             if type(arg) == pd.DataFrame:
                 pandas_dataframe = arg
-                output_table_name = (
-                    f"{self.dag_id}_{self.task_id}_{self.run_id}_input_dataframe_{i}"
-                )
+                output_table_name = create_unique_table_name()
                 output_table = Table(
                     table_name=output_table_name,
                     conn_id=self.conn_id,
@@ -338,15 +334,7 @@ def convert_op_kwarg_dataframes(self):
         for key, value in self.op_kwargs.items():
             if type(value) == pd.DataFrame:
                 pandas_dataframe = value
-                output_table_name = "_".join(
-                    [
-                        self.dag_id,
-                        self.task_id,
-                        self.run_id,
-                        "input_dataframe",
-                        str(key),
-                    ]
-                )
+                output_table_name = create_unique_table_name()
                 output_table = Table(
                     table_name=output_table_name,
                     conn_id=self.conn_id,
@@ -361,7 +349,7 @@ def convert_op_kwarg_dataframes(self):
                     warehouse=self.warehouse,
                 )
                 load_dataframe_into_sql_table(pandas_dataframe, output_table, hook)
-                final_kwargs[key] = output_table_name
+                final_kwargs[key] = output_table
             else:
                 final_kwargs[key] = value
         self.op_kwargs = final_kwargs
diff --git a/src/astro/utils/postgres_transform.py b/src/astro/utils/postgres_transform.py
--- a/src/astro/utils/postgres_transform.py
+++ b/src/astro/utils/postgres_transform.py
@@ -4,9 +4,8 @@
 
 def add_templates_to_context(parameters, context):
     for k, v in parameters.items():
-        if type(v) == Table:
-            final_name = v.schema + "." + v.table_name if v.schema else v.table_name
-            context[k] = final_name
+        if isinstance(v, Table):
+            context[k] = v.qualified_name()
         else:
             context[k] = ":" + k
     return context


Test patch:
diff --git a/tests/operators/test_dataframe.py b/tests/operators/test_dataframe.py
--- a/tests/operators/test_dataframe.py
+++ b/tests/operators/test_dataframe.py
@@ -92,6 +92,29 @@ def create_and_run_task(self, decorator_func, op_args, op_kwargs):
 
         return f
 
+    def test_postgres_dataframe_without_table_arg(self):
+        @df
+        def validate_result(df: pandas.DataFrame):
+            assert df.iloc[0].to_dict()["colors"] == "red"
+
+        @df
+        def sample_df():
+            return pandas.DataFrame(
+                {"numbers": [1, 2, 3], "colors": ["red", "white", "blue"]}
+            )
+
+        @aql.transform
+        def sample_pg(input: Table):
+            return "SELECT * FROM {{input}}"
+
+        with self.dag:
+            plain_df = sample_df()
+            pg_df = sample_pg(
+                conn_id="postgres_conn", database="pagila", input=plain_df
+            )
+            validate_result(pg_df)
+        test_utils.run_dag(self.dag)
+
     def test_dataframe_from_sql_basic(self):
         @df
         def my_df_func(df: pandas.DataFrame):