Exasol hook sql sensor

airflow/contrib/hooks/exasol_hook.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+
+from contextlib import closing
+
+from pyexasol import connect
+
+from airflow.hooks.dbapi_hook import DbApiHook
+
+
+class ExasolHook(DbApiHook):
+    """
+    Interact with Vertica.
+    """
+
+    conn_name_attr = 'exasol_conn_id'
+    default_conn_name = 'exasol_default'
+    supports_autocommit = True
+
+    def get_conn(self):
+        """
+        Returns exasol connection object
+        """
+        conn = self.get_connection(self.exasol_conn_id)
+        conn_config = {
+            "user": conn.login,
+            "password": conn.password or '',
+            "schema": conn.schema,
+            "dsn": (conn.host or 'localhost') + ':' + str(conn.port)
+        }
+
+        conn = connect(**conn_config)
+        return conn
+
+    def get_records(self, sql, parameters=None):
+        """
+        Executes the sql and returns a set of records.
+
+        :param sql: the sql statement to be executed (str) or a list of
+            sql statements to execute
+        :type sql: str or list
+        :param parameters: The parameters to render the SQL query with.
+        :type parameters: mapping or iterable
+        """
+        with closing(self.get_conn()) as conn:
+            if parameters is not None:
+                cur = conn.execute(sql, parameters)
+            else:
+                cur = conn.execute(sql)
+            return cur.fetchall()

airflow/models/connection.py

@@ -105,6 +105,7 @@ class Connection(Base, LoggingMixin):
         ('qubole', 'Qubole'),
         ('mongo', 'MongoDB'),
         ('gcpcloudsql', 'Google Cloud SQL'),
+        ('exasol', 'Exasol'),
     ]
 
     def __init__(
@@ -198,72 +199,73 @@ def rotate_fernet_key(self):
             self._extra = fernet.rotate(self._extra.encode('utf-8')).decode()
 
     def get_hook(self):
-        try:
-            if self.conn_type == 'mysql':
-                from airflow.hooks.mysql_hook import MySqlHook
-                return MySqlHook(mysql_conn_id=self.conn_id)
-            elif self.conn_type == 'google_cloud_platform':
-                from airflow.contrib.hooks.bigquery_hook import BigQueryHook
-                return BigQueryHook(bigquery_conn_id=self.conn_id)
-            elif self.conn_type == 'postgres':
-                from airflow.hooks.postgres_hook import PostgresHook
-                return PostgresHook(postgres_conn_id=self.conn_id)
-            elif self.conn_type == 'hive_cli':
-                from airflow.hooks.hive_hooks import HiveCliHook
-                return HiveCliHook(hive_cli_conn_id=self.conn_id)
-            elif self.conn_type == 'presto':
-                from airflow.hooks.presto_hook import PrestoHook
-                return PrestoHook(presto_conn_id=self.conn_id)
-            elif self.conn_type == 'hiveserver2':
-                from airflow.hooks.hive_hooks import HiveServer2Hook
-                return HiveServer2Hook(hiveserver2_conn_id=self.conn_id)
-            elif self.conn_type == 'sqlite':
-                from airflow.hooks.sqlite_hook import SqliteHook
-                return SqliteHook(sqlite_conn_id=self.conn_id)
-            elif self.conn_type == 'jdbc':
-                from airflow.hooks.jdbc_hook import JdbcHook
-                return JdbcHook(jdbc_conn_id=self.conn_id)
-            elif self.conn_type == 'mssql':
-                from airflow.hooks.mssql_hook import MsSqlHook
-                return MsSqlHook(mssql_conn_id=self.conn_id)
-            elif self.conn_type == 'oracle':
-                from airflow.hooks.oracle_hook import OracleHook
-                return OracleHook(oracle_conn_id=self.conn_id)
-            elif self.conn_type == 'vertica':
-                from airflow.contrib.hooks.vertica_hook import VerticaHook
-                return VerticaHook(vertica_conn_id=self.conn_id)
-            elif self.conn_type == 'cloudant':
-                from airflow.contrib.hooks.cloudant_hook import CloudantHook
-                return CloudantHook(cloudant_conn_id=self.conn_id)
-            elif self.conn_type == 'jira':
-                from airflow.contrib.hooks.jira_hook import JiraHook
-                return JiraHook(jira_conn_id=self.conn_id)
-            elif self.conn_type == 'redis':
-                from airflow.contrib.hooks.redis_hook import RedisHook
-                return RedisHook(redis_conn_id=self.conn_id)
-            elif self.conn_type == 'wasb':
-                from airflow.contrib.hooks.wasb_hook import WasbHook
-                return WasbHook(wasb_conn_id=self.conn_id)
-            elif self.conn_type == 'docker':
-                from airflow.hooks.docker_hook import DockerHook
-                return DockerHook(docker_conn_id=self.conn_id)
-            elif self.conn_type == 'azure_data_lake':
-                from airflow.contrib.hooks.azure_data_lake_hook import AzureDataLakeHook
-                return AzureDataLakeHook(azure_data_lake_conn_id=self.conn_id)
-            elif self.conn_type == 'azure_cosmos':
-                from airflow.contrib.hooks.azure_cosmos_hook import AzureCosmosDBHook
-                return AzureCosmosDBHook(azure_cosmos_conn_id=self.conn_id)
-            elif self.conn_type == 'cassandra':
-                from airflow.contrib.hooks.cassandra_hook import CassandraHook
-                return CassandraHook(cassandra_conn_id=self.conn_id)
-            elif self.conn_type == 'mongo':
-                from airflow.contrib.hooks.mongo_hook import MongoHook
-                return MongoHook(conn_id=self.conn_id)
-            elif self.conn_type == 'gcpcloudsql':
-                from airflow.contrib.hooks.gcp_sql_hook import CloudSqlDatabaseHook
-                return CloudSqlDatabaseHook(gcp_cloudsql_conn_id=self.conn_id)
-        except Exception:
-            pass
+        if self.conn_type == 'mysql':
+            from airflow.hooks.mysql_hook import MySqlHook
+            return MySqlHook(mysql_conn_id=self.conn_id)
+        elif self.conn_type == 'google_cloud_platform':
+            from airflow.contrib.hooks.bigquery_hook import BigQueryHook
+            return BigQueryHook(bigquery_conn_id=self.conn_id)
+        elif self.conn_type == 'postgres':
+            from airflow.hooks.postgres_hook import PostgresHook
+            return PostgresHook(postgres_conn_id=self.conn_id)
+        elif self.conn_type == 'hive_cli':
+            from airflow.hooks.hive_hooks import HiveCliHook
+            return HiveCliHook(hive_cli_conn_id=self.conn_id)
+        elif self.conn_type == 'presto':
+            from airflow.hooks.presto_hook import PrestoHook
+            return PrestoHook(presto_conn_id=self.conn_id)
+        elif self.conn_type == 'hiveserver2':
+            from airflow.hooks.hive_hooks import HiveServer2Hook
+            return HiveServer2Hook(hiveserver2_conn_id=self.conn_id)
+        elif self.conn_type == 'sqlite':
+            from airflow.hooks.sqlite_hook import SqliteHook
+            return SqliteHook(sqlite_conn_id=self.conn_id)
+        elif self.conn_type == 'jdbc':
+            from airflow.hooks.jdbc_hook import JdbcHook
+            return JdbcHook(jdbc_conn_id=self.conn_id)
+        elif self.conn_type == 'mssql':
+            from airflow.hooks.mssql_hook import MsSqlHook
+            return MsSqlHook(mssql_conn_id=self.conn_id)
+        elif self.conn_type == 'oracle':
+            from airflow.hooks.oracle_hook import OracleHook
+            return OracleHook(oracle_conn_id=self.conn_id)
+        elif self.conn_type == 'vertica':
+            from airflow.contrib.hooks.vertica_hook import VerticaHook
+            return VerticaHook(vertica_conn_id=self.conn_id)
+        elif self.conn_type == 'cloudant':
+            from airflow.contrib.hooks.cloudant_hook import CloudantHook
+            return CloudantHook(cloudant_conn_id=self.conn_id)
+        elif self.conn_type == 'jira':
+            from airflow.contrib.hooks.jira_hook import JiraHook
+            return JiraHook(jira_conn_id=self.conn_id)
+        elif self.conn_type == 'redis':
+            from airflow.contrib.hooks.redis_hook import RedisHook
+            return RedisHook(redis_conn_id=self.conn_id)
+        elif self.conn_type == 'wasb':
+            from airflow.contrib.hooks.wasb_hook import WasbHook
+            return WasbHook(wasb_conn_id=self.conn_id)
+        elif self.conn_type == 'docker':
+            from airflow.hooks.docker_hook import DockerHook
+            return DockerHook(docker_conn_id=self.conn_id)
+        elif self.conn_type == 'azure_data_lake':
+            from airflow.contrib.hooks.azure_data_lake_hook import AzureDataLakeHook
+            return AzureDataLakeHook(azure_data_lake_conn_id=self.conn_id)
+        elif self.conn_type == 'azure_cosmos':
+            from airflow.contrib.hooks.azure_cosmos_hook import AzureCosmosDBHook
+            return AzureCosmosDBHook(azure_cosmos_conn_id=self.conn_id)
+        elif self.conn_type == 'cassandra':
+            from airflow.contrib.hooks.cassandra_hook import CassandraHook
+            return CassandraHook(cassandra_conn_id=self.conn_id)
+        elif self.conn_type == 'mongo':
+            from airflow.contrib.hooks.mongo_hook import MongoHook
+            return MongoHook(conn_id=self.conn_id)
+        elif self.conn_type == 'gcpcloudsql':
+            from airflow.contrib.hooks.gcp_sql_hook import CloudSqlDatabaseHook
+            return CloudSqlDatabaseHook(gcp_cloudsql_conn_id=self.conn_id)
+        elif self.conn_type == 'exasol':
+            from airflow.contrib.hooks.exasol_hook import ExasolHook
+            return ExasolHook(exasol_conn_id=self.conn_id)
+        raise AirflowException("Unknown hook type {}".format(self.conn_type))
 
     def __repr__(self):
         return self.conn_id

airflow/sensors/sql_sensor.py

@@ -55,7 +55,7 @@ def poke(self, context):
 
         allowed_conn_type = {'google_cloud_platform', 'jdbc', 'mssql',
                              'mysql', 'oracle', 'postgres',
-                             'presto', 'sqlite', 'vertica'}
+                             'presto', 'sqlite', 'vertica', 'exasol'}
         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)))