Repository: astronomer/astro-sdk

Issue:
Release 0.8.3
**Context**
Release 0.8.3 by creating a branch from the last 0.8.x release and cherry-picking the fix for #329.

**Acceptance criteria**
* [ ] Release 0.8.3
* [ ] It does not contain any refactoring
* [ ] It contains the bug fix fr #329 
SQLDecorator attempting to create a schema when it already exists
**Describe the bug**
When a schema already exists, we should not attempt to create it. This is particularly important for users who may not have permission to create schemas.

Reported by: @jlaneve 

**Version**
* Astro: 0.8.1

**To Reproduce**
When running queries against the snowflake_conn_4 connection, everything logs the error:
```
(snowflake.connector.errors.ProgrammingError) 003001 (42501): 01a3ebcc-0404-8f39-0000-68210845688e: SQL access control error:
Insufficient privileges to operate on database 'SANDBOX'
[SQL: CREATE SCHEMA IF NOT EXISTS tmp_astro]
(Background on this error at: http://sqlalche.me/e/13/f405)
```
**Expected behavior**
If a schema already exists, Astro should not attempt to create it.

**Additional context**
The SQLDecorator (in theory) should not need permission to create schemas. It originally did, and [this commit](https://github.com/astro-projects/astro/commit/e8b8d076047522c9421c9a33a2e5f9b524d89e3a) fixed the issue. However, it was reverted in [this commit](https://github.com/astro-projects/astro/commit/eb25da9a9c9003a973e10defeb17914f5c00f0b6) 1 month ago. The specific line causing the issue is [here](https://github.com/astro-projects/astro/blame/be28e5ea36ac810d550e4f1b140cbd1306b2da96/src/astro/sql/operators/sql_decorator.py#L109).

**Acceptance criteria**
* [ ] The SQLDecorator does not attempt to create schema when it does not exist
* [ ] There is a test covering this use-case


Discussion:

Resolving patch:
diff --git a/conftest.py b/conftest.py
--- a/conftest.py
+++ b/conftest.py
@@ -195,14 +195,14 @@ def remote_file(request):
         raise ValueError(f"File location {provider} not supported")
 
     new_connection = Connection(conn_id=conn_id, conn_type=conn_type, extra=extra)
-    session = settings.Session()
+    session_ = settings.Session()
     if not (
-        session.query(Connection)
+        session_.query(Connection)
         .filter(Connection.conn_id == new_connection.conn_id)
         .first()
     ):
-        session.add(new_connection)
-        session.commit()
+        session_.add(new_connection)
+        session_.commit()
 
     filename = pathlib.Path(f"{CWD}/tests/data/sample.{file_extension}")
     object_paths = []
@@ -235,5 +235,22 @@ def remote_file(request):
             ):
                 hook.delete_objects(bucket_name, object_prefix)
 
-    session.delete(new_connection)
-    session.flush()
+    session_.delete(new_connection)
+    session_.flush()
+
+
+@pytest.fixture
+def schema_fixture(request, sql_server):
+    """
+    Given request.param in the format:
+        "someschema"  # name of the schema to be created
+
+    If the schema exists, it is deleted during the tests setup and tear down.
+    """
+    schema_name = request.param
+    _, hook = sql_server
+
+    hook.run(f"DROP SCHEMA IF EXISTS {schema_name}")
+    yield
+
+    hook.run(f"DROP SCHEMA IF EXISTS {schema_name}")
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.8.2"
+__version__ = "0.8.3"
 
 from astro.dataframe import dataframe  # noqa: F401
 
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
@@ -4,8 +4,6 @@
 import pandas as pd
 from airflow.decorators.base import DecoratedOperator, task_decorator_factory
 from airflow.hooks.base import BaseHook
-from airflow.utils.db import provide_session
-from sqlalchemy.exc import ProgrammingError
 from sqlalchemy.sql.functions import Function
 
 from astro.constants import Database
@@ -37,6 +35,7 @@ def __init__(
         sql="",
         **kwargs,
     ):
+        self.op_args = ()
         self.raw_sql = raw_sql
         self.autocommit = autocommit
         self.parameters = parameters
@@ -81,6 +80,7 @@ def hook(self):
         )
 
     def execute(self, context: Dict):
+
         if not isinstance(self.sql, str):
             cursor = self._run_sql(self.sql, self.parameters)
             if self.handler is not None:
@@ -105,18 +105,14 @@ def execute(self, context: Dict):
         self._process_params()
 
         output_table_name: str = ""
-        # Create DEFAULT SCHEMA if not created.
-        self._set_schema_if_needed(SCHEMA)
+
+        self._set_schema_if_needed()
 
         if not self.raw_sql:
-            # Create a table name for the temp table
-            if not schema_exists(
-                hook=self.hook, schema=self.schema, conn_type=self.conn_type
-            ):
-                self._set_schema_if_needed()
 
-            if not self.output_table or type(self.output_table) == TempTable:
+            if not self.output_table or isinstance(self.output_table, TempTable):
                 output_table_name = create_unique_table_name()
+                self._set_schema_if_needed(schema=SCHEMA)
                 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,
@@ -135,13 +131,13 @@ def execute(self, context: Dict):
         else:
             # If there's no SQL to run we simply return
             if self.sql == "" or not self.sql:
-                return
+                return None
 
         query_result = self._run_sql(self.sql, self.parameters)
         # Run execute function of subclassed Operator.
 
         if self.output_table:
-            if type(self.output_table) == TempTable:
+            if isinstance(self.output_table, TempTable):
                 self.output_table = self.output_table.to_table(
                     table_name=output_table_name, schema=self.output_table.schema
                 )
@@ -206,17 +202,26 @@ def handle_output_table_schema(
         return output_table_name
 
     def _set_schema_if_needed(self, schema=None):
-        if self.conn_type in ["postgres", "snowflake", "bigquery"]:
+        if schema is None:
+            schema = self.schema
+
+        database_expects_schema = self.conn_type in [
+            "postgres",
+            "snowflake",
+            "bigquery",
+        ]
+        schema_exist_ = schema_exists(
+            hook=self.hook, schema=schema, conn_type=self.conn_type
+        )
+
+        if database_expects_schema and not schema_exist_:
             schema_statement = create_schema_query(
                 conn_type=self.conn_type,
                 hook=self.hook,
-                schema_id=schema if schema else self.schema,
+                schema_id=schema,
                 user=self.user,
             )
-            try:
-                self._run_sql(schema_statement, {})
-            except ProgrammingError as e:
-                print(e)
+            self._run_sql(schema_statement, {})
 
     def get_sql_alchemy_engine(self):
         return get_sqlalchemy_engine(self.hook)
@@ -253,17 +258,6 @@ def clean_trailing_semicolon(query):
             statement = f"CREATE TABLE {output_table_name} AS ({clean_trailing_semicolon(query)});"
         return statement
 
-    @provide_session
-    def pre_execute(self, context, session=None):
-        """This hook is triggered right before self.execute() is called."""
-        pass
-
-    def post_execute(self, context, result=None):
-        """
-        This hook is triggered right after self.execute() is called.
-        """
-        pass
-
     def _process_params(self):
         if self.conn_type == "snowflake":
             self.parameters = snowflake_transform.process_params(
@@ -274,14 +268,14 @@ def _add_templates_to_context(self, context):
         database = get_database_from_conn_id(self.conn_id)
         if database in (Database.POSTGRES, Database.BIGQUERY):
             return postgres_transform.add_templates_to_context(self.parameters, context)
-        elif database == Database.SNOWFLAKE:
+        if database == Database.SNOWFLAKE:
             return snowflake_transform.add_templates_to_context(
                 self.parameters, context
             )
-        else:
-            return self.default_transform(self.parameters, context)
+        return self.default_transform(self.parameters, context)
 
-    def default_transform(self, parameters, context):
+    @staticmethod
+    def default_transform(parameters, context):
         for k, v in parameters.items():
             if isinstance(v, Table):
                 context[k] = v.table_name
@@ -289,15 +283,10 @@ def default_transform(self, parameters, context):
                 context[k] = ":" + k
         return context
 
-    def _cleanup(self):
-        """Remove DAG's objects from S3 and db."""
-        # To-do
-        pass
-
     def convert_op_arg_dataframes(self):
         final_args = []
-        for i, arg in enumerate(self.op_args):
-            if type(arg) == pd.DataFrame:
+        for arg in self.op_args:
+            if isinstance(arg, pd.DataFrame):
                 pandas_dataframe = arg
                 output_table_name = create_unique_table_name()
                 output_table = Table(
@@ -322,7 +311,7 @@ def convert_op_arg_dataframes(self):
     def convert_op_kwarg_dataframes(self):
         final_kwargs = {}
         for key, value in self.op_kwargs.items():
-            if type(value) == pd.DataFrame:
+            if isinstance(value, pd.DataFrame):
                 pandas_dataframe = value
                 output_table_name = create_unique_table_name()
                 output_table = Table(


Test patch:
diff --git a/tests/operators/test_sql_decorator.py b/tests/operators/test_sql_decorator.py
--- a/tests/operators/test_sql_decorator.py
+++ b/tests/operators/test_sql_decorator.py
@@ -7,6 +7,7 @@
 from tests.operators import utils as test_utils
 
 CWD = pathlib.Path(__file__).parent
+TEST_SCHEMA = test_utils.get_table_name("test")
 
 
 @pytest.mark.parametrize(
@@ -58,3 +59,97 @@ def null_function():  # skipcq: PTC-W0049
             sql=f"SELECT list FROM {test_table.qualified_name()} WHERE sell=232",
         )
     test_utils.run_dag(sample_dag)
+
+
+@pytest.mark.parametrize(
+    "sql_server",
+    ["snowflake", "bigquery", "postgres"],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "test_table",
+    [
+        {
+            "path": str(CWD) + "/../data/sample.csv",
+            "load_table": True,
+            "is_temp": False,
+            "param": {
+                "table_name": test_utils.get_table_name("test"),
+            },
+        }
+    ],
+    indirect=True,
+)
+def test_sql_decorator_does_not_create_schema_when_the_schema_exists(
+    sample_dag, sql_server, test_table
+):
+    """Test basic sql execution of SqlDecoratedOperator."""
+    _, hook = sql_server
+
+    sql_statement = f"SELECT * FROM {test_table.qualified_name()} WHERE id=4"
+    df = hook.get_pandas_df(sql_statement)
+    assert df.empty
+
+    with sample_dag:
+        SqlDecoratedOperator(
+            raw_sql=True,
+            parameters={},
+            task_id="SomeTask",
+            op_args=(),
+            conn_id=test_table.conn_id,
+            database=test_table.database,
+            python_callable=lambda: None,
+            sql=f"INSERT INTO {test_table.qualified_name()} (id, name) VALUES (4, 'New Person');",
+        )
+    test_utils.run_dag(sample_dag)
+
+    df = hook.get_pandas_df(sql_statement).rename(columns=str.lower)
+    assert df.to_dict("r") == [{"id": 4, "name": "New Person"}]
+
+
+@pytest.mark.parametrize(
+    "sql_server",
+    ["postgres", "bigquery"],
+    indirect=True,
+)
+@pytest.mark.parametrize("schema_fixture", [TEST_SCHEMA], indirect=True)
+@pytest.mark.parametrize(
+    "test_table",
+    [
+        {
+            "path": str(CWD) + "/../data/sample.csv",
+            "load_table": True,
+            "is_temp": False,
+            "param": {
+                "schema": TEST_SCHEMA,
+                "table_name": test_utils.get_table_name("test"),
+            },
+        }
+    ],
+    indirect=True,
+)
+def test_sql_decorator_creates_schema_when_it_does_not_exist(
+    sample_dag, sql_server, schema_fixture, test_table
+):
+    """Test basic sql execution of SqlDecoratedOperator."""
+    _, hook = sql_server
+
+    sql_statement = f"SELECT * FROM {test_table.qualified_name()} WHERE id=4"
+    df = hook.get_pandas_df(sql_statement)
+    assert df.empty
+
+    with sample_dag:
+        SqlDecoratedOperator(
+            raw_sql=True,
+            parameters={},
+            task_id="SomeTask",
+            op_args=(),
+            conn_id=test_table.conn_id,
+            database=test_table.database,
+            python_callable=lambda: None,
+            sql=f"INSERT INTO {test_table.qualified_name()} (id, name) VALUES (4, 'New Person');",
+        )
+    test_utils.run_dag(sample_dag)
+
+    df = hook.get_pandas_df(sql_statement)
+    assert df.to_dict("r") == [{"id": 4, "name": "New Person"}]