Repository: astronomer/astro-sdk

Issue:
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
@@ -171,6 +171,23 @@ def sql_server(request):
         hook.run(f"DROP TABLE IF EXISTS {schema}.{OUTPUT_TABLE_NAME}")
 
 
+@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}")
+
+
 @pytest.fixture
 def database_table_fixture(request):
     """
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
@@ -5,7 +5,6 @@
 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
@@ -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:
                 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,
@@ -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)


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"}]