def execute(self, context: 'Context') -> None: self.log.info('Executing: %s', self.sql) hook = JdbcHook(jdbc_conn_id=self.jdbc_conn_id) return hook.run(self.sql, self.autocommit, parameters=self.parameters, handler=fetch_all_handler)
def execute(self, context: 'Context'): self.log.info('Executing: %s', self.sql) hook = JdbcHook(jdbc_conn_id=self.jdbc_conn_id) if self.do_xcom_push: return hook.run(self.sql, self.autocommit, parameters=self.parameters, handler=self.handler) else: return hook.run(self.sql, self.autocommit, parameters=self.parameters)
class JdbcOperator(BaseOperator): """ Executes sql code in a database using jdbc driver. Requires jaydebeapi. :param sql: the sql code to be executed. (templated) :type sql: Can receive a str representing a sql statement, a list of str (sql statements), or reference to a template file. Template reference are recognized by str ending in '.sql' :param jdbc_conn_id: reference to a predefined database :type jdbc_conn_id: str :param autocommit: if True, each command is automatically committed. (default value: False) :type autocommit: bool :param parameters: (optional) the parameters to render the SQL query with. :type parameters: mapping or iterable """ template_fields = ('sql', ) template_ext = ('.sql', ) ui_color = '#ededed' @apply_defaults def __init__(self, sql: str, jdbc_conn_id: str = 'jdbc_default', autocommit: bool = False, parameters: Optional[Union[Mapping, Iterable]] = None, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.parameters = parameters self.sql = sql self.jdbc_conn_id = jdbc_conn_id self.autocommit = autocommit def execute(self, context): self.log.info('Executing: %s', self.sql) self.hook = JdbcHook(jdbc_conn_id=self.jdbc_conn_id) self.hook.run(self.sql, self.autocommit, parameters=self.parameters)
def test_jdbc_conn_get_autocommit(self, _): jdbc_hook = JdbcHook() jdbc_conn = jdbc_hook.get_conn() jdbc_hook.get_autocommit(jdbc_conn) jdbc_conn.jconn.getAutoCommit.assert_called_once_with()
def test_jdbc_conn_connection(self, jdbc_mock): jdbc_hook = JdbcHook() jdbc_conn = jdbc_hook.get_conn() self.assertTrue(jdbc_mock.called) self.assertIsInstance(jdbc_conn, Mock) self.assertEqual(jdbc_conn.name, jdbc_mock.return_value.name) # pylint: disable=no-member
def execute(self, context): self.log.info('Executing: %s', self.sql) hook = JdbcHook(jdbc_conn_id=self.jdbc_conn_id) hook.run(self.sql, self.autocommit, parameters=self.parameters)
def get_hook(self): if self.conn_type == 'mysql': from airflow.providers.mysql.hooks.mysql import MySqlHook return MySqlHook(mysql_conn_id=self.conn_id) elif self.conn_type == 'google_cloud_platform': from airflow.gcp.hooks.bigquery import BigQueryHook return BigQueryHook(bigquery_conn_id=self.conn_id) elif self.conn_type == 'postgres': from airflow.providers.postgres.hooks.postgres import PostgresHook return PostgresHook(postgres_conn_id=self.conn_id) elif self.conn_type == 'pig_cli': from airflow.providers.apache.pig.hooks.pig import PigCliHook return PigCliHook(pig_cli_conn_id=self.conn_id) elif self.conn_type == 'hive_cli': from airflow.providers.apache.hive.hooks.hive import HiveCliHook return HiveCliHook(hive_cli_conn_id=self.conn_id) elif self.conn_type == 'presto': from airflow.providers.presto.hooks.presto import PrestoHook return PrestoHook(presto_conn_id=self.conn_id) elif self.conn_type == 'hiveserver2': from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook return HiveServer2Hook(hiveserver2_conn_id=self.conn_id) elif self.conn_type == 'sqlite': from airflow.providers.sqlite.hooks.sqlite import SqliteHook return SqliteHook(sqlite_conn_id=self.conn_id) elif self.conn_type == 'jdbc': from airflow.providers.jdbc.hooks.jdbc import JdbcHook return JdbcHook(jdbc_conn_id=self.conn_id) elif self.conn_type == 'mssql': from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook return MsSqlHook(mssql_conn_id=self.conn_id) elif self.conn_type == 'odbc': from airflow.providers.odbc.hooks.odbc import OdbcHook return OdbcHook(odbc_conn_id=self.conn_id) elif self.conn_type == 'oracle': from airflow.providers.oracle.hooks.oracle import OracleHook return OracleHook(oracle_conn_id=self.conn_id) elif self.conn_type == 'vertica': from airflow.providers.vertica.hooks.vertica import VerticaHook return VerticaHook(vertica_conn_id=self.conn_id) elif self.conn_type == 'cloudant': from airflow.providers.cloudant.hooks.cloudant import CloudantHook return CloudantHook(cloudant_conn_id=self.conn_id) elif self.conn_type == 'jira': from airflow.providers.jira.hooks.jira import JiraHook return JiraHook(jira_conn_id=self.conn_id) elif self.conn_type == 'redis': from airflow.providers.redis.hooks.redis import RedisHook return RedisHook(redis_conn_id=self.conn_id) elif self.conn_type == 'wasb': from airflow.providers.microsoft.azure.hooks.wasb import WasbHook return WasbHook(wasb_conn_id=self.conn_id) elif self.conn_type == 'docker': from airflow.providers.docker.hooks.docker import DockerHook return DockerHook(docker_conn_id=self.conn_id) elif self.conn_type == 'azure_data_lake': from airflow.providers.microsoft.azure.hooks.azure_data_lake import AzureDataLakeHook return AzureDataLakeHook(azure_data_lake_conn_id=self.conn_id) elif self.conn_type == 'azure_cosmos': from airflow.providers.microsoft.azure.hooks.azure_cosmos import AzureCosmosDBHook return AzureCosmosDBHook(azure_cosmos_conn_id=self.conn_id) elif self.conn_type == 'cassandra': from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook return CassandraHook(cassandra_conn_id=self.conn_id) elif self.conn_type == 'mongo': from airflow.providers.mongo.hooks.mongo import MongoHook return MongoHook(conn_id=self.conn_id) elif self.conn_type == 'gcpcloudsql': from airflow.gcp.hooks.cloud_sql import CloudSQLDatabaseHook return CloudSQLDatabaseHook(gcp_cloudsql_conn_id=self.conn_id) elif self.conn_type == 'grpc': from airflow.providers.grpc.hooks.grpc import GrpcHook return GrpcHook(grpc_conn_id=self.conn_id) raise AirflowException("Unknown hook type {}".format(self.conn_type))
def test_jdbc_conn_connection(self, jdbc_mock): jdbc_hook = JdbcHook() jdbc_conn = jdbc_hook.get_conn() assert jdbc_mock.called assert isinstance(jdbc_conn, Mock) assert jdbc_conn.name == jdbc_mock.return_value.name # pylint: disable=no-member