From bdf4bc5674bdf341e73be46ffa92bd03c5e14935 Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Fri, 11 Mar 2022 14:28:47 +0100
Subject: [PATCH] Remove Snowflake limits

The limits of snowflake can be removed as the conditions for
all limits are fulfilled now.

This is the last thing to remove before we can enable Python 3.10
---
 airflow/providers/snowflake/hooks/snowflake.py | 7 ++++---
 setup.py                                       | 9 ++-------
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py
index 5bf94fae6984f..82373b65861f9 100644
--- a/airflow/providers/snowflake/hooks/snowflake.py
+++ b/airflow/providers/snowflake/hooks/snowflake.py
@@ -18,7 +18,7 @@
 import os
 from contextlib import closing
 from io import StringIO
-from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Union
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
@@ -262,7 +262,7 @@ def run(
         self,
         sql: Union[str, list],
         autocommit: bool = False,
-        parameters: Optional[Union[dict, Iterable]] = None,
+        parameters: Optional[Union[Sequence[Any], Dict[Any, Any]]] = None,
         handler: Optional[Callable] = None,
     ):
         """
@@ -290,7 +290,8 @@ def run(
                 sql = [sql_string for sql_string, _ in split_statements_tuple if sql_string]
 
             self.log.debug("Executing %d statements against Snowflake DB", len(sql))
-            with closing(conn.cursor(DictCursor)) as cur:
+            # SnowflakeCursor does not extend ContextManager, so we have to ignore mypy error here
+            with closing(conn.cursor(DictCursor)) as cur:  # type: ignore[type-var]
 
                 for sql_statement in sql:
 
diff --git a/setup.py b/setup.py
index ab4b5018ea249..2e4f3ef72fd69 100644
--- a/setup.py
+++ b/setup.py
@@ -510,13 +510,8 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
     'slack_sdk>=3.0.0',
 ]
 snowflake = [
-    # Snowflake connector 2.7.2 requires pyarrow >=6.0.0 but apache-beam requires < 6.0.0
-    # We should remove the limitation when apache-beam upgrades pyarrow
-    'snowflake-connector-python>=2.4.1,<2.7.2',
-    # The snowflake-alchemy 1.2.5 introduces a hard dependency on sqlalchemy>=1.4.0, but they didn't define
-    # this requirements in setup.py, so pip cannot figure out the correct set of dependencies.
-    # See: https://github.com/snowflakedb/snowflake-sqlalchemy/issues/234
-    'snowflake-sqlalchemy>=1.1.0,!=1.2.5',
+    'snowflake-connector-python>=2.4.1',
+    'snowflake-sqlalchemy>=1.1.0',
 ]
 spark = [
     'pyspark',