def execute(self, context: Dict[str, str]): mssql = MsSqlHook(mssql_conn_id=self.mssql_conn_id) self.log.info( "Dumping Microsoft SQL Server query results to local file") with mssql.get_conn() as conn: with conn.cursor() as cursor: cursor.execute(self.sql) with NamedTemporaryFile("w") as tmp_file: csv_writer = csv.writer(tmp_file, delimiter=self.delimiter, encoding='utf-8') field_dict = OrderedDict() col_count = 0 for field in cursor.description: col_count += 1 col_position = "Column{position}".format( position=col_count) field_dict[col_position if field[0] == '' else field[0]] = self.type_map(field[1]) csv_writer.writerows(cursor) tmp_file.flush() hive = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id) self.log.info("Loading file into Hive") hive.load_file(tmp_file.name, self.hive_table, field_dict=field_dict, create=self.create, partition=self.partition, delimiter=self.delimiter, recreate=self.recreate, tblproperties=self.tblproperties)
def test_get_conn_should_return_connection(self, get_connection, mssql_get_conn): get_connection.return_value = PYMSSQL_CONN mssql_get_conn.return_value = mock.Mock() hook = MsSqlHook() conn = hook.get_conn() self.assertEqual(mssql_get_conn.return_value, conn) mssql_get_conn.assert_called_once()
def test_get_autocommit_should_return_autocommit_state(self, get_connection, mssql_get_conn): get_connection.return_value = PYMSSQL_CONN mssql_get_conn.return_value = mock.Mock() mssql_get_conn.return_value.autocommit_state = 'autocommit_state' hook = MsSqlHook() conn = hook.get_conn() mssql_get_conn.assert_called_once() self.assertEqual(hook.get_autocommit(conn), 'autocommit_state')
def query(self): """ Queries MSSQL and returns a cursor of results. :return: mssql cursor """ mssql = MsSqlHook(mssql_conn_id=self.mssql_conn_id) conn = mssql.get_conn() cursor = conn.cursor() cursor.execute(self.sql) return cursor
def test_set_autocommit_should_invoke_autocommit(self, get_connection, mssql_get_conn): get_connection.return_value = PYMSSQL_CONN mssql_get_conn.return_value = mock.Mock() autocommit_value = mock.Mock() hook = MsSqlHook() conn = hook.get_conn() hook.set_autocommit(conn, autocommit_value) mssql_get_conn.assert_called_once() mssql_get_conn.return_value.autocommit.assert_called_once_with(autocommit_value)