Repository: apache/airflow

Issue:
Add support for snowflake on the sql_sensor
**Description**

Enable the support for Snowflake on the SqlSensor.

**Use case**

Currently, the `SqlSensor` supports only the following connection types,

- google_cloud_platform
- jdbc
- mssql
- mysql
- oracle
- postgres
- presto
- sqlite
- vertica

At the moment, it doesn't have support for Snowflake. This ticket is to add support for Snowflake on the `SqlSensor`. 




Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

Hi, I would like to pick up this issue if it is a good feature to come in.
@KumarManoj-S I assigned you. to this ticket. This change seems sensible. We should not have lists of approved classes in the code, but we should rely on interfaces.
@KumarManoj-S  note that if this works for SqlSensor there is no reason it won't work for SqlOperators
https://github.com/apache/airflow/blob/1d36b0303b8632fce6de78ca4e782ae26ee06fea/airflow/sensors/sql_sensor.py#L72
https://github.com/apache/airflow/blob/0b9bf4a285a074bbde270839a90fb53c257340be/airflow/operators/sql.py#L26
 (not sure why two different lists are needed as lists are the same and there shouldn't be a reason why sensor would work on specific db but the operators won't or vice versa)
Might worth checking that if you can address both
@eladkal Sure. That's a good point. I will check if I can address both of them.

Resolving patch:
diff --git a/airflow/models/connection.py b/airflow/models/connection.py
--- a/airflow/models/connection.py
+++ b/airflow/models/connection.py
@@ -72,7 +72,7 @@
     "jira": ("airflow.providers.jira.hooks.jira.JiraHook", "jira_conn_id"),
     "kubernetes": ("airflow.providers.cncf.kubernetes.hooks.kubernetes.KubernetesHook", "kubernetes_conn_id"),
     "mongo": ("airflow.providers.mongo.hooks.mongo.MongoHook", "conn_id"),
-    "mssql": ("airflow.providers.microsoft.mssql.hooks.mssql.MsSqlHook", "mssql_conn_id"),
+    "mssql": ("airflow.providers.odbc.hooks.odbc.OdbcHook", "odbc_conn_id"),
     "mysql": ("airflow.providers.mysql.hooks.mysql.MySqlHook", "mysql_conn_id"),
     "odbc": ("airflow.providers.odbc.hooks.odbc.OdbcHook", "odbc_conn_id"),
     "oracle": ("airflow.providers.oracle.hooks.oracle.OracleHook", "oracle_conn_id"),
@@ -80,6 +80,7 @@
     "postgres": ("airflow.providers.postgres.hooks.postgres.PostgresHook", "postgres_conn_id"),
     "presto": ("airflow.providers.presto.hooks.presto.PrestoHook", "presto_conn_id"),
     "redis": ("airflow.providers.redis.hooks.redis.RedisHook", "redis_conn_id"),
+    "snowflake": ("airflow.providers.snowflake.hooks.snowflake.SnowflakeHook", "snowflake_conn_id"),
     "sqlite": ("airflow.providers.sqlite.hooks.sqlite.SqliteHook", "sqlite_conn_id"),
     "tableau": ("airflow.providers.salesforce.hooks.tableau.TableauHook", "tableau_conn_id"),
     "vertica": ("airflow.providers.vertica.hooks.vertica.VerticaHook", "vertica_conn_id"),
diff --git a/airflow/operators/sql.py b/airflow/operators/sql.py
--- a/airflow/operators/sql.py
+++ b/airflow/operators/sql.py
@@ -32,6 +32,7 @@
     "oracle",
     "postgres",
     "presto",
+    "snowflake",
     "sqlite",
     "vertica",
 }
diff --git a/airflow/sensors/sql_sensor.py b/airflow/sensors/sql_sensor.py
--- a/airflow/sensors/sql_sensor.py
+++ b/airflow/sensors/sql_sensor.py
@@ -71,7 +71,7 @@ def _get_hook(self):
 
         allowed_conn_type = {'google_cloud_platform', 'jdbc', 'mssql',
                              'mysql', 'odbc', 'oracle', 'postgres',
-                             'presto', 'sqlite', 'vertica'}
+                             'presto', 'snowflake', 'sqlite', 'vertica'}
         if conn.conn_type not in allowed_conn_type:
             raise AirflowException("The connection type is not supported by SqlSensor. " +
                                    "Supported connection types: {}".format(list(allowed_conn_type)))


Test patch:
diff --git a/tests/models/test_connection.py b/tests/models/test_connection.py
--- a/tests/models/test_connection.py
+++ b/tests/models/test_connection.py
@@ -30,6 +30,7 @@
 from airflow.models import Connection, crypto
 from airflow.models.connection import CONN_TYPE_TO_HOOK
 from airflow.providers.sqlite.hooks.sqlite import SqliteHook
+from airflow.utils.module_loading import import_string
 from tests.test_utils.config import conf_vars
 
 ConnectionParts = namedtuple("ConnectionParts", ["conn_type", "login", "password", "host", "port", "schema"])
@@ -533,3 +534,7 @@ def test_enforce_alphabetical_order(self):
         expected_keys = sorted(current_keys)
 
         self.assertEqual(expected_keys, current_keys)
+
+    def test_hooks_importable(self):
+        for hook_path, _ in CONN_TYPE_TO_HOOK.values():
+            self.assertTrue(issubclass(import_string(hook_path), BaseHook))