groovy脚本连接mongodb导出数据到excel

import com.mongodb.BasicDBObject;
import com.mongodb.MongoClient;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import org.bson.Document;
import org.apache.poi.hssf.usermodel.HSSFCell;
import org.apache.poi.hssf.usermodel.HSSFCellStyle;
import org.apache.poi.hssf.usermodel.HSSFFont;
import org.apache.poi.hssf.usermodel.HSSFRow;
import org.apache.poi.hssf.usermodel.HSSFSheet;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.apache.poi.hssf.util.HSSFColor;
import org.apache.poi.ss.usermodel.BorderStyle;
import org.apache.poi.ss.usermodel.FillPatternType;
import org.apache.poi.ss.usermodel.HorizontalAlignment;
import org.apache.poi.ss.usermodel.VerticalAlignment;


try{
    //startTime、endTime、testId为自定义参数，根据实际情况自行调整
	Date startDate = parseStringToDate(startTime);
	if(startDate == null)
		return false;
	println "查询开始时间： "+startDate;
	Date endDate = parseStringToDate(endTime);
	if(endDate == null)
		return false;
	println "查询结束时间： " +endDate;
	println "开始连接mongodb数据库...";
	MongoClient mongoClient = connectMongoDB();
	if(mongoClient == null){
		println "mongodb数据库连接失败！";
		return false;
	}
	println "mongodb数据库连接成功！";
        MongoDatabase mongoDatabase = mongoClient.getDatabase("test_db");
        MongoCollection<Document> collection =  mongoDatabase.getCollection("test_coll");
        BasicDBObject timeCondition = new BasicDBObject("\$gte", startDate);
        timeCondition.put("\$lte", endDate);
        BasicDBObject condition = new BasicDBObject("occurTime",timeCondition);
        condition.put("testId", testId);
        println "查询条件： "+condition;
	def dataNum = collection.count(condition);
	if(dataNum > 100000){
		println "数据量过大可能会导致数据写入崩溃，建议缩小时间区间。"
		return false;
	}	
        FindIterable<Document> find = collection.find(condition);
	println "开始导出数据...";
        exportExcel(find);
	println "数据导出成功！";
	return true;
} catch (Exception e) {
	println "Export alert infos error: "+e.getMessage();
}

//连接mongodb数据库
def connectMongoDB(){
	String usr = "root";
	String pwd = "root_123";
	String src = "admin";
	MongoCredential au = MongoCredential.createCredential(usr, src, pwd.toCharArray());
	//mongodb的地址信息存在主备情况
	String mongoAddr = "0.0.0.1:1001,0.0.0.2:1002,0.0.0.3:1003";
	String[] mongoAddrs = mongoAddr.split(",");
	for(String ip_port : mongoAddrs){
	    String mongoDB_Ip = ip_port.split(":")[0].trim();
		Integer mongoDB_Port = Integer.valueOf(ip_port.split(":")[1]).trim();
		MongoClient client = new MongoClient(new ServerAddress(mongoDB_Ip, mongoDB_Port), Arrays.asList(au));
		List<String> dbs = null;
		try {
        	dbs = client.getDatabaseNames(); //验证mongodb是否连接上
    	} catch (Exception e) {
        	println "mongodb数据库连接不上，正在尝试连接到备用数据库......错误信息："+e.getMessage();
		dbs = null;
    	}
		if(dbs != null){
			return client;
		}
	}
	return null;
}


//Date类型日期转化为日期字符串
def parseDate2String(Date date){
    Calendar calendar = Calendar.getInstance();
    calendar.setTime(date);
    calendar.add(Calendar.HOUR, 8);
    SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    return sdf.format(date);
}

//日期字符串转换为Date类型日期
def parseStringToDate(String dateStr){
    SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    try {
        return sdf.parse(dateStr);
    } catch (ParseException e) {
        println "日期格式错误："+dateStr+", 正确格式样例：1970-01-01 08:00:00";
    }
}

def exportExcel(FindIterable<Document> find){
    HSSFWorkbook wb = new HSSFWorkbook();
    HSSFSheet sheet = wb.createSheet("sheetName");
    HSSFFont font = wb.createFont();
    HSSFCellStyle headerStyle = wb.createCellStyle();
    sheet.setDefaultColumnWidth((short) 18);
    //设置字体样式
    font.setFontHeightInPoints((short) 11.2);
    font.setBold(true);
    font.setFontName("宋体");
					 
    //设置表头
    headerStyle.setWrapText(true);//自动换行
    headerStyle.setFont(font);
    //设置居中
    headerStyle.setAlignment(HorizontalAlignment.CENTER);
    headerStyle.setVerticalAlignment(VerticalAlignment.CENTER);
    //设置单元格背景颜色
    headerStyle.setFillForegroundColor(HSSFColor.GREY_25_PERCENT.index);//  LIGHT_GREEN  LIGHT_TURQUOISE   GREY_25_PERCENT
    headerStyle.setFillPattern(FillPatternType.SOLID_FOREGROUND);
    //style.setFillBackgroundColor(HSSFColor.GREY_25_PERCENT.index);
    //设置单元格边框titleStyle.setFillPattern(FillPatternType.SOLID_FOREGROUND);
    headerStyle.setBorderTop(BorderStyle.THIN);
    headerStyle.setBorderBottom(BorderStyle.THIN);
    headerStyle.setBorderLeft(BorderStyle.THIN);
    headerStyle.setBorderRight(BorderStyle.THIN);
					 
    // 创建[表中数据]样式
    HSSFCellStyle dataSetStyle = wb.createCellStyle();
    // 设置[表中数据]样式
    dataSetStyle.setFillPattern(FillPatternType.SOLID_FOREGROUND);
    dataSetStyle.setBorderBottom(BorderStyle.THIN);
    dataSetStyle.setBorderLeft(BorderStyle.THIN);
    dataSetStyle.setBorderRight(BorderStyle.THIN);
    dataSetStyle.setBorderTop(BorderStyle.THIN);
    dataSetStyle.setAlignment(HorizontalAlignment.CENTER);
    dataSetStyle.setVerticalAlignment(VerticalAlignment.CENTER);
					 
    //设置首行标题列
    String[] headers = ["header1","header2","header3", "header4", "header5", "header6", "header7"];
    String[] dataKeys = ["key1", "key2", "key3", "key4", "key5","key6", "key7"];
    HSSFRow row0 = sheet.createRow(0);
    HSSFCell cell;
    for(int i=0; i<headers.length; i++){
        cell = row0.createCell((short) i);
        cell.setCellValue(headers[i]);
        cell.setCellStyle(headerStyle);
    }
	
    //设置数据
    int j = 1;
    for(Document doc : find){
        row0 = sheet.createRow(j);
        for(int k=0; k<dataKeys.length; k++){
            String dateKey = dataKeys[k];
            cell = row0.createCell(k);
            cell.setCellValue(doc.getString(dateKey));
            cell.setCellStyle(dataSetStyle);
        }
        j++;
    }
    println "excel中数据写入成功！"
    //将数据写入指定文件中
    try {
        String fileName = "导出数据-" + new Date().format("yyyyMMddHHmmss")  + ".xlsx";
	String baseDir = new File("").getCanonicalPath();
        String tempPath = baseDir + File.separator + "temp";
        new File(tempPath).mkdirs();
        File file = new File(tempPath+ File.separator+fileName);
	if(!file.exists()){
		file.createNewFile();
	}
	wb.write(file);
    } catch (IOException e) {
        e.printStackTrace();
    } finally {
        try {
            wb.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}