Repository: astronomer/astro-sdk

Issue:
Astro Build's Integration Test breaking on 0.8.1
broken on - 0.8.1 , but works with 0.7.0.

In this test dag, task_5 joins tables from task_3 (snowflake) and task_4 (postgres). The task’s print statement shows the joined table, suggesting successful ingestion and transformation. However, the error below suggests the the postgres output conn is expected to contain the database element of the snowflake connection.
In other words, the task fails becausepostgres_conn doesn’t have the database attribute associated with snowflake_conn.

```
import time
from datetime import datetime

import pandas as pd
from airflow.decorators import dag, task
from airflow.models import Variable
from airflow.utils import timezone
from airflow.utils.dates import days_ago
from astro import dataframe as df
from astro import sql as aql
from astro.sql.table import Table, TempTable

@df()
def task_1_func():
    return pd.DataFrame({'a':[1,2,3]})

@aql.transform(conn_id='postgres_conn')
def task_2_func(execution_date: Table):
    return """SELECT * FROM actor WHERE startdate < '{{ execution_date }}'"""

@aql.transform(conn_id='snowflake_conn')
def task_3_func():
    return """SELECT * FROM "ASTROBUILD"."BUILDSCHEMA"."MYTABLE" LIMIT 10;"""

@aql.transform(conn_id='postgres_conn')
def task_4_func():
    return """SELECT * FROM actor LIMIT 10;"""

@df(conn_id='postgres_conn')
def task_5_func(task_3: pd.DataFrame, task_4: pd.DataFrame):
    df=task_3.join(task_4)
    print(df)
    return df

@dag(schedule_interval='0 0 * * *', start_date=datetime(2022, 4, 15, 11, 28, 8), catchup=False, tags=['tag_1', 'tag_1'])
def dag_1():
    task_1 = task_1_func()
    task_2 = task_2_func(output_table=Table(conn_id='postgres_conn', schema='tmp_astro', table_name='tmp_astro_dag_1_task_2'), execution_date=Table(conn_id='postgres_conn', table_name='execution_date'))
    task_3 = task_3_func(output_table=Table(conn_id='snowflake_conn', schema='tmp_astro', table_name='tmp_astro_dag_1_task_3'))
    task_4 = task_4_func(output_table=Table(conn_id='postgres_conn', schema='tmp_astro', table_name='tmp_astro_dag_1_task_4'))
    task_5 = task_5_func(task_3, task_4, output_table=Table(conn_id='postgres_conn', schema='tmp_astro', table_name='tmp_astro_dag_1_task_5'))

dag_obj = dag_1()
```

Error:

```
INFO - Using connection to: id: postgres_conn. Host: 127.0.0.1, Port: 8999, Schema: postgres, Login: postgres, Password: ***, extra: {}
*** psycopg2.OperationalError: connection to server at "127.0.0.1", port 8999 failed: FATAL:  database "ASTROBUILD" does not exist
```


Discussion:

Resolving patch:
diff --git a/src/astro/utils/table_handler.py b/src/astro/utils/table_handler.py
--- a/src/astro/utils/table_handler.py
+++ b/src/astro/utils/table_handler.py
@@ -16,28 +16,51 @@ def _set_variables_from_first_table(self):
         """
         first_table: Optional[Table] = None
         if self.op_args:
-            table_index = [x for x, t in enumerate(self.op_args) if type(t) == Table]
-            if table_index:
+            table_index = [
+                x for x, t in enumerate(self.op_args) if isinstance(t, Table)
+            ]
+            conn_id_set = {x.conn_id for x in self.op_args if isinstance(x, Table)}
+            # Check to see if all tables belong to same conn_id. Otherwise, we this can go wrong for cases
+            # 1. When we have tables from different DBs.
+            # 2. When we have tables from different conn_id, since they can be configured with different
+            # database/schema etc.
+            if table_index and len(conn_id_set) == 1:
                 first_table = self.op_args[table_index[0]]
-        elif not first_table:
+
+        if not first_table and self.op_kwargs and self.python_callable:
             table_kwargs = [
                 x
                 for x in inspect.signature(self.python_callable).parameters.values()
                 if (
                     x.annotation == Table
-                    and type(self.op_kwargs[x.name]) == Table
+                    and isinstance(self.op_kwargs[x.name], Table)
                     or x.annotation == pandas.DataFrame
-                    and type(self.op_kwargs[x.name]) == Table
+                    and isinstance(self.op_kwargs[x.name], Table)
                 )
             ]
-            if table_kwargs:
+            conn_id_set = {
+                self.op_kwargs[x.name].conn_id
+                for x in inspect.signature(self.python_callable).parameters.values()
+                if (
+                    x.annotation == Table
+                    and isinstance(self.op_kwargs[x.name], Table)
+                    or x.annotation == pandas.DataFrame
+                    and isinstance(self.op_kwargs[x.name], Table)
+                )
+            }
+            if table_kwargs and len(conn_id_set) == 1:
                 first_table = self.op_kwargs[table_kwargs[0].name]
 
         # If there is no first table via op_ags or kwargs, we check the parameters
-        elif not first_table:
+        if not first_table and self.parameters:
             if self.parameters:
-                param_tables = [t for t in self.parameters.values() if type(t) == Table]
-                if param_tables:
+                param_tables = [
+                    t for t in self.parameters.values() if isinstance(t, Table)
+                ]
+                conn_id_set = {
+                    t.conn_id for t in self.parameters.values() if isinstance(t, Table)
+                }
+                if param_tables and len(conn_id_set) == 1:
                     first_table = param_tables[0]
 
         if first_table:


Test patch:
diff --git a/tests/miscellaneous/test_table_handler.py b/tests/miscellaneous/test_table_handler.py
new file mode 100644
--- /dev/null
+++ b/tests/miscellaneous/test_table_handler.py
@@ -0,0 +1,216 @@
+from astro.sql.table import Table
+from astro.utils.table_handler import TableHandler
+
+
+def test__set_variables_from_first_table_with_same_db_tables_in_op_args():
+    """Test _set_variables_from_first_table() when the tables passed are with same tables"""
+
+    def dummy_function(param_1: str, param_2: str):  # skipcq: PTC-W0049, PY-D0003
+        pass
+
+    handler = TableHandler()
+    handler.python_callable = dummy_function
+    handler.op_kwargs = {"param_1": "dummy_value", "param_2": "dummy_value"}
+    handler.parameters = {}
+    handler.op_args = (
+        Table(
+            conn_id="conn_1",
+            database="database_1",
+            schema="scheme_1",
+            warehouse="warehouse_1",
+            role="role_1",
+        ),
+        Table(
+            conn_id="conn_1",
+            database="database_1",
+            schema="scheme_1",
+            warehouse="warehouse_1",
+            role="role_1",
+        ),
+    )
+    handler._set_variables_from_first_table()
+    assert handler.conn_id == "conn_1"
+    assert handler.database == "database_1"
+    assert handler.schema == "scheme_1"
+    assert handler.warehouse == "warehouse_1"
+    assert handler.role == "role_1"
+
+
+def test__set_variables_from_first_table_with_different_db_tables_in_op_args():
+    """Test _set_variables_from_first_table() when the tables passed are with same tables"""
+
+    def dummy_function(param_1: str, param_2: str):  # skipcq: PTC-W0049, PY-D0003
+        pass
+
+    handler = TableHandler()
+    handler.python_callable = dummy_function
+    handler.op_kwargs = {"param_1": "dummy_value", "param_2": "dummy_value"}
+    handler.parameters = {}
+    handler.op_args = (
+        Table(
+            conn_id="conn_1",
+            database="database_1",
+            schema="scheme_1",
+            warehouse="warehouse_1",
+            role="role_1",
+        ),
+        Table(
+            conn_id="conn_2",
+            database="database_2",
+            schema="scheme_2",
+            warehouse="warehouse_2",
+            role="role_2",
+        ),
+    )
+
+    handler._set_variables_from_first_table()
+    assert not hasattr(handler, "conn_id")
+    assert not hasattr(handler, "database")
+    assert not hasattr(handler, "schema")
+    assert not hasattr(handler, "warehouse")
+    assert not hasattr(handler, "role")
+
+
+def test__set_variables_from_first_table_with_same_db_tables_in_python_callable():
+    """Test _set_variables_from_first_table() when the tables passed are with same tables in python_callable"""
+    table_1 = Table(
+        conn_id="conn_1",
+        database="database_1",
+        schema="scheme_1",
+        warehouse="warehouse_1",
+        role="role_1",
+    )
+    table_2 = Table(
+        conn_id="conn_1",
+        database="database_1",
+        schema="scheme_1",
+        warehouse="warehouse_1",
+        role="role_1",
+    )
+
+    def dummy_function(param_1: Table, param_2: Table):  # skipcq: PTC-W0049, PY-D0003
+        pass
+
+    handler = TableHandler()
+    handler.python_callable = dummy_function
+    handler.op_args = ()
+    handler.parameters = {}
+    handler.op_kwargs = {
+        "param_1": table_1,
+        "param_2": table_2,
+    }
+
+    handler._set_variables_from_first_table()
+    assert handler.conn_id == "conn_1"
+    assert handler.database == "database_1"
+    assert handler.schema == "scheme_1"
+    assert handler.warehouse == "warehouse_1"
+    assert handler.role == "role_1"
+
+
+def test__set_variables_from_first_table_with_different_db_tables_in_python_callable():
+    """Test _set_variables_from_first_table() when the tables passed are with same tables in python_callable"""
+    table_1 = Table(
+        conn_id="conn_1",
+        database="database_1",
+        schema="scheme_1",
+        warehouse="warehouse_1",
+        role="role_1",
+    )
+    table_2 = Table(
+        conn_id="conn_2",
+        database="database_2",
+        schema="scheme_2",
+        warehouse="warehouse_2",
+        role="role_2",
+    )
+
+    def dummy_function(param_1: Table, param_2: Table):  # skipcq: PTC-W0049, PY-D0003
+        pass
+
+    handler = TableHandler()
+    handler.python_callable = dummy_function
+    handler.op_args = ()
+    handler.parameters = {}
+    handler.op_kwargs = {
+        "param_1": table_1,
+        "param_2": table_2,
+    }
+
+    handler._set_variables_from_first_table()
+    assert not hasattr(handler, "conn_id")
+    assert not hasattr(handler, "database")
+    assert not hasattr(handler, "schema")
+    assert not hasattr(handler, "warehouse")
+    assert not hasattr(handler, "role")
+
+
+def test__set_variables_from_first_table_with_same_db_tables_in_parameters():
+    """Test _set_variables_from_first_table() when the tables passed are with same tables in parameters"""
+
+    def dummy_function(param_1: str, param_2: str):  # skipcq: PTC-W0049, PY-D0003
+        pass
+
+    handler = TableHandler()
+    handler.op_args = ()
+    handler.python_callable = dummy_function
+    handler.op_kwargs = {"param_1": "dummy_value", "param_2": "dummy_value"}
+    handler.parameters = {
+        "param_1": Table(
+            conn_id="conn_1",
+            database="database_1",
+            schema="scheme_1",
+            warehouse="warehouse_1",
+            role="role_1",
+        ),
+        "param_3": Table(
+            conn_id="conn_1",
+            database="database_1",
+            schema="scheme_1",
+            warehouse="warehouse_1",
+            role="role_1",
+        ),
+    }
+
+    handler._set_variables_from_first_table()
+
+    assert handler.conn_id == "conn_1"
+    assert handler.database == "database_1"
+    assert handler.schema == "scheme_1"
+    assert handler.warehouse == "warehouse_1"
+    assert handler.role == "role_1"
+
+
+def test__set_variables_from_first_table_with_different_db_tables_in_parameters():
+    """Test _set_variables_from_first_table() when the tables passed are with same tables in parameters"""
+
+    def dummy_function(param_1: str, param_2: str):  # skipcq: PTC-W0049, PY-D0003
+        pass
+
+    handler = TableHandler()
+    handler.op_args = ()
+    handler.python_callable = dummy_function
+    handler.op_kwargs = {"param_1": "dummy_value", "param_2": "dummy_value"}
+    handler.parameters = {
+        "param_1": Table(
+            conn_id="conn_1",
+            database="database_1",
+            schema="scheme_1",
+            warehouse="warehouse_1",
+            role="role_1",
+        ),
+        "param_3": Table(
+            conn_id="conn_2",
+            database="database_2",
+            schema="scheme_2",
+            warehouse="warehouse_2",
+            role="role_2",
+        ),
+    }
+
+    handler._set_variables_from_first_table()
+    assert not hasattr(handler, "conn_id")
+    assert not hasattr(handler, "database")
+    assert not hasattr(handler, "schema")
+    assert not hasattr(handler, "warehouse")
+    assert not hasattr(handler, "role")