Repository: astronomer/astro-sdk

Issue:
Unable to run @transform without the `input_table` parameter (using a Snowflake connection)
**Describe the bug**
Unable to run @transform without the `input_table` parameter (using a Snowflake connection)

**Version**
* Astro: 0.8.0

**To Reproduce**
Steps to reproduce the behavior:
1. Write the DAG:
```
"""
testing
DAG file auto-generated by AstroBuild.
"""
import time
from datetime import datetime

from airflow.decorators import dag, taskle
from astro.sql.table import Table

@aql.transform(conn_id='snowflake_conn')
def task_1_func():
    return """
            select catalog_name as database,
                schema_name,
                schema_owner,
                created,
                last_altered
            from information_schema.schemata
            order by schema_name;
        """

@dag(schedule_interval='0 0 * * *', start_date=datetime(2022, 4, 14, 0, 17, 20), catchup=False, tags=[])
def julian():
    task_1 = task_1_func(output_table=Table(conn_id='snowflake_conn', schema='TATIANAALCHUEYR', table_name='tmp_astro_testing_task_1'))

dag_obj = julian()
```
3. Make sure the connection `snowflake_conn` exists
4. Run the DAG
5. See error:
```
      File /home/tati/Code/clean-astro/src/astro/sql/operators/sql_decorator.py, line 91, in execute
        self.database = self.database or self.database_from_conn_id
    AttributeError: 'SqlDecoratedOperator' object has no attribute 'database'
```

**Expected behavior**
The DAG works without raising any exception.


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
@@ -52,8 +52,8 @@ def __init__(
         if self.op_kwargs.get("handler"):
             self.handler = self.op_kwargs.pop("handler")
 
-        self.database = self.op_kwargs.pop("database", database)
         self.conn_id = self.op_kwargs.pop("conn_id", conn_id)
+        self.database = self.op_kwargs.pop("database", database)
         self.schema = self.op_kwargs.pop("schema", schema)
         self.warehouse = self.op_kwargs.pop("warehouse", warehouse)
         self.role = self.op_kwargs.pop("role", role)
@@ -66,6 +66,10 @@ def __init__(
     def conn_type(self):
         return BaseHook.get_connection(self.conn_id).conn_type
 
+    @property
+    def database_from_conn_id(self):
+        return BaseHook.get_connection(self.conn_id).extra_dejson.get("database")
+
     @property
     def hook(self):
         return get_hook(
@@ -84,6 +88,7 @@ def execute(self, context: Dict):
             return cursor
 
         self._set_variables_from_first_table()
+        self.database = self.database or self.database_from_conn_id
 
         conn = BaseHook.get_connection(self.conn_id)
 


Test patch:
diff --git a/tests/operators/transform/test_snowflake_transform.py b/tests/operators/transform/test_snowflake_transform.py
--- a/tests/operators/transform/test_snowflake_transform.py
+++ b/tests/operators/transform/test_snowflake_transform.py
@@ -78,6 +78,7 @@ def validate_table(df: pd.DataFrame):
             path=str(CWD) + "/../../data/homes.csv",
             output_table=table,
         )
+
         f = sample_snow(
             input_table=loaded_table,
             output_table=TempTable(
@@ -105,3 +106,32 @@ def test_roles_failing(sql_server, sample_dag, test_table):
 def test_roles_passing(sql_server, sample_dag, test_table):
     test_table.role = os.getenv("SNOWFLAKE_ROLE")
     run_role_query(sample_dag, test_table, role=os.getenv("SNOWFLAKE_ROLE"))
+
+
+def run_simple_transform(dag, table, role):
+    """Reproduces a live issue #319"""
+
+    @aql.transform(conn_id="snowflake_conn")
+    def select_without_input_table():
+        """A transform without input_table"""
+        return """
+            select catalog_name as database,
+                schema_name,
+                schema_owner,
+                created,
+                last_altered
+            from information_schema.schemata
+            order by schema_name;
+        """
+
+    with dag:
+        select_without_input_table(output_table=table)
+
+    test_utils.run_dag(dag)
+
+
+@pytest.mark.parametrize("sql_server", ["snowflake"], indirect=True)
+def test_transform_without_input_table(sql_server, sample_dag, test_table):
+    """Reproduces issue #319"""
+    test_table.role = os.getenv("SNOWFLAKE_ROLE")
+    run_simple_transform(sample_dag, test_table, role=os.getenv("SNOWFLAKE_ROLE"))