Repository: apache/airflow

Issue:
S3ToSnowflakeTransfer enforces usage of schema parameter
Unlike all other Snowflake operators, S3ToSnowflakeTransfer requires the usage of schema instead of letting it be an optional variable that when passed would be used, and if not, then connection metadata is used. 

I can be assigned on this if needed. Keeping this up here so this doesn't go unnoticed. 


Discussion:
@davido912 assigned the issue to you :)
hi @davido912 since you didn't finish your pr https://github.com/apache/airflow/pull/13017 I unassign you from this issue. If you wish to complete it feel free to do so.

This issue is open for anyone who wants to pick it.

Resolving patch:
diff --git a/airflow/providers/snowflake/transfers/s3_to_snowflake.py b/airflow/providers/snowflake/transfers/s3_to_snowflake.py
--- a/airflow/providers/snowflake/transfers/s3_to_snowflake.py
+++ b/airflow/providers/snowflake/transfers/s3_to_snowflake.py
@@ -79,7 +79,7 @@ def __init__(
         stage: str,
         prefix: Optional[str] = None,
         file_format: str,
-        schema: str,  # TODO: shouldn't be required, rely on session/user defaults
+        schema: Optional[str] = None,
         columns_array: Optional[list] = None,
         warehouse: Optional[str] = None,
         database: Optional[str] = None,
@@ -117,33 +117,23 @@ def execute(self, context: Any) -> None:
             session_parameters=self.session_parameters,
         )
 
-        files = ""
-        if self.s3_keys:
-            files = "files=({})".format(", ".join(f"'{key}'" for key in self.s3_keys))
+        if self.schema:
+            into = f"{self.schema}.{self.table}"
+        else:
+            into = self.table
+        if self.columns_array:
+            into = f"{into}({','.join(self.columns_array)})"
 
-        # we can extend this based on stage
-        base_sql = """
-                    FROM @{stage}/{prefix}
-                    {files}
-                    file_format={file_format}
-                """.format(
-            stage=self.stage,
-            prefix=(self.prefix if self.prefix else ""),
-            files=files,
-            file_format=self.file_format,
-        )
+        sql_parts = [
+            f"COPY INTO {into}",
+            f"FROM @{self.stage}/{self.prefix or ''}",
+        ]
+        if self.s3_keys:
+            files = ", ".join(f"'{key}'" for key in self.s3_keys)
+            sql_parts.append(f"files=({files})")
+        sql_parts.append(f"file_format={self.file_format}")
 
-        if self.columns_array:
-            copy_query = """
-                COPY INTO {schema}.{table}({columns}) {base_sql}
-            """.format(
-                schema=self.schema, table=self.table, columns=",".join(self.columns_array), base_sql=base_sql
-            )
-        else:
-            copy_query = f"""
-                COPY INTO {self.schema}.{self.table} {base_sql}
-            """
-        copy_query = "\n".join(line.strip() for line in copy_query.splitlines())
+        copy_query = "\n".join(sql_parts)
 
         self.log.info('Executing COPY command...')
         snowflake_hook.run(copy_query, self.autocommit)


Test patch:
diff --git a/tests/providers/snowflake/transfers/test_s3_to_snowflake.py b/tests/providers/snowflake/transfers/test_s3_to_snowflake.py
--- a/tests/providers/snowflake/transfers/test_s3_to_snowflake.py
+++ b/tests/providers/snowflake/transfers/test_s3_to_snowflake.py
@@ -27,12 +27,12 @@ class TestS3ToSnowflakeTransfer:
     @pytest.mark.parametrize("columns_array", [None, ['col1', 'col2', 'col3']])
     @pytest.mark.parametrize("s3_keys", [None, ['1.csv', '2.csv']])
     @pytest.mark.parametrize("prefix", [None, 'prefix'])
+    @pytest.mark.parametrize("schema", [None, 'schema'])
     @mock.patch("airflow.providers.snowflake.hooks.snowflake.SnowflakeHook.run")
-    def test_execute(self, mock_run, prefix, s3_keys, columns_array):
+    def test_execute(self, mock_run, schema, prefix, s3_keys, columns_array):
         table = 'table'
         stage = 'stage'
         file_format = 'file_format'
-        schema = 'schema'
 
         S3ToSnowflakeOperator(
             s3_keys=s3_keys,
@@ -46,22 +46,21 @@ def test_execute(self, mock_run, prefix, s3_keys, columns_array):
             dag=None,
         ).execute(None)
 
-        files = None
+        copy_query = "COPY INTO "
+        if schema:
+            copy_query += f"{schema}.{table}"
+        else:
+            copy_query += table
+        if columns_array:
+            copy_query += f"({','.join(columns_array)})"
+
+        copy_query += f"\nFROM @{stage}/{prefix or ''}"
+
         if s3_keys:
-            files = "files=({})".format(", ".join(f"'{key}'" for key in s3_keys))
-        base_sql = f"""
-                FROM @{stage}/{prefix if prefix else ''}
-                {files if files else ''}
-                file_format={file_format}
-            """
+            files = ", ".join(f"'{key}'" for key in s3_keys)
+            copy_query += f"\nfiles=({files})"
 
-        columns = None
-        if columns_array:
-            columns = f"({','.join(columns_array)})"
-        copy_query = f"""
-                COPY INTO {schema}.{table}{columns if columns else ''} {base_sql}
-            """
-        copy_query = "\n".join(line.strip() for line in copy_query.splitlines())
+        copy_query += f"\nfile_format={file_format}"
 
         mock_run.assert_called_once()
         assert mock_run.call_args[0][0] == copy_query