def execute(self, context): self.log.info('Executing: %s', self.sql) hook = ExasolHook(exasol_conn_id=self.exasol_conn_id, schema=self.schema) hook.run(self.sql, autocommit=self.autocommit, parameters=self.parameters)
def setUp(self): super().setUp() self.connection = models.Connection( login='******', password='******', host='host', port=1234, schema='schema', ) self.db_hook = ExasolHook() self.db_hook.get_connection = mock.Mock() self.db_hook.get_connection.return_value = self.connection
class TestExasolHookConn(unittest.TestCase): def setUp(self): super().setUp() self.connection = models.Connection( login='******', password='******', host='host', port=1234, schema='schema', ) self.db_hook = ExasolHook() self.db_hook.get_connection = mock.Mock() self.db_hook.get_connection.return_value = self.connection @mock.patch('airflow.providers.exasol.hooks.exasol.pyexasol') def test_get_conn(self, mock_pyexasol): self.db_hook.get_conn() mock_connect = mock_pyexasol.connect mock_connect.assert_called_once() args, kwargs = mock_connect.call_args self.assertEqual(args, ()) self.assertEqual(kwargs['user'], 'login') self.assertEqual(kwargs['password'], 'password') self.assertEqual(kwargs['dsn'], 'host:1234') self.assertEqual(kwargs['schema'], 'schema') @mock.patch('airflow.providers.exasol.hooks.exasol.pyexasol') def test_get_conn_extra_args(self, mock_pyexasol): self.connection.extra = json.dumps({'encryption': True}) self.db_hook.get_conn() mock_connect = mock_pyexasol.connect mock_connect.assert_called_once() args, kwargs = mock_connect.call_args self.assertEqual(args, ()) self.assertEqual(kwargs['encryption'], True)
def execute(self, context: 'Context'): exasol_hook = ExasolHook(exasol_conn_id=self.exasol_conn_id) s3_hook = S3Hook(aws_conn_id=self.aws_conn_id) with NamedTemporaryFile("w+") as file: exasol_hook.export_to_file( filename=file.name, query_or_table=self.query_or_table, export_params=self.export_params, query_params=self.query_params, ) file.flush() self.log.info("Uploading the data as %s", self.key) s3_hook.load_file( filename=file.name, key=self.key, bucket_name=self.bucket_name, replace=self.replace, encrypt=self.encrypt, gzip=self.gzip, acl_policy=self.acl_policy, ) self.log.info("Data uploaded") return self.key