From 2039b4af8cb2907565c82f8cb2a81f96fa01c759 Mon Sep 17 00:00:00 2001
From: Ray Wang <2532542@qq.com>
Date: Thu, 22 Nov 2018 15:57:28 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=A0=B9=E6=8D=AE=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E6=BA=90=E7=94=9F=E6=88=90=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../web/controller/tool/GenController.java    | 25 ++++++++++++-------
 .../resources/templates/tool/gen/gen.html     | 11 +++++---
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/ruoyi-admin/src/main/java/com/ruoyi/web/controller/tool/GenController.java b/ruoyi-admin/src/main/java/com/ruoyi/web/controller/tool/GenController.java
index b5e68dc9e..eea2d6790 100644
--- a/ruoyi-admin/src/main/java/com/ruoyi/web/controller/tool/GenController.java
+++ b/ruoyi-admin/src/main/java/com/ruoyi/web/controller/tool/GenController.java
@@ -3,15 +3,15 @@ package com.ruoyi.web.controller.tool;
 import java.io.IOException;
 import java.util.List;
 import javax.servlet.http.HttpServletResponse;
+
+import com.ruoyi.common.enums.DataSourceType;
+import com.ruoyi.framework.datasource.DynamicDataSourceContextHolder;
 import org.apache.commons.io.IOUtils;
 import org.apache.shiro.authz.annotation.RequiresPermissions;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Controller;
-import org.springframework.web.bind.annotation.GetMapping;
-import org.springframework.web.bind.annotation.PathVariable;
-import org.springframework.web.bind.annotation.PostMapping;
-import org.springframework.web.bind.annotation.RequestMapping;
-import org.springframework.web.bind.annotation.ResponseBody;
+import org.springframework.ui.ModelMap;
+import org.springframework.web.bind.annotation.*;
 import com.ruoyi.common.annotation.Log;
 import com.ruoyi.common.enums.BusinessType;
 import com.ruoyi.common.support.Convert;
@@ -36,8 +36,9 @@ public class GenController extends BaseController
 
     @RequiresPermissions("tool:gen:view")
     @GetMapping()
-    public String gen()
+    public String gen(ModelMap mmap)
     {
+        mmap.put("dataSource", DataSourceType.MASTER.name());
         return prefix + "/gen";
     }
 
@@ -47,7 +48,9 @@ public class GenController extends BaseController
     public TableDataInfo list(TableInfo tableInfo)
     {
         startPage();
+        DynamicDataSourceContextHolder.setDateSoureType((String) tableInfo.getParams().get("dataSource"));
         List<TableInfo> list = genService.selectTableList(tableInfo);
+        DynamicDataSourceContextHolder.clearDateSoureType();
         return getDataTable(list);
     }
 
@@ -56,10 +59,12 @@ public class GenController extends BaseController
      */
     @RequiresPermissions("tool:gen:code")
     @Log(title = "代码生成", businessType = BusinessType.GENCODE)
-    @GetMapping("/genCode/{tableName}")
-    public void genCode(HttpServletResponse response, @PathVariable("tableName") String tableName) throws IOException
+    @GetMapping("/genCode/{tableName}/{dataSource}")
+    public void genCode(HttpServletResponse response, @PathVariable("tableName") String tableName, @PathVariable("dataSource") String dataSource) throws IOException
     {
+        DynamicDataSourceContextHolder.setDateSoureType(dataSource);
         byte[] data = genService.generatorCode(tableName);
+        DynamicDataSourceContextHolder.clearDateSoureType();
         response.reset();
         response.setHeader("Content-Disposition", "attachment; filename=\"ruoyi.zip\"");
         response.addHeader("Content-Length", "" + data.length);
@@ -75,10 +80,12 @@ public class GenController extends BaseController
     @Log(title = "代码生成", businessType = BusinessType.GENCODE)
     @GetMapping("/batchGenCode")
     @ResponseBody
-    public void batchGenCode(HttpServletResponse response, String tables) throws IOException
+    public void batchGenCode(HttpServletResponse response, String tables, String dataSource) throws IOException
     {
         String[] tableNames = Convert.toStrArray(tables);
+        DynamicDataSourceContextHolder.setDateSoureType(dataSource);
         byte[] data = genService.generatorCode(tableNames);
+        DynamicDataSourceContextHolder.clearDateSoureType();
         response.reset();
         response.setHeader("Content-Disposition", "attachment; filename=\"ruoyi.zip\"");
         response.addHeader("Content-Length", "" + data.length);
diff --git a/ruoyi-admin/src/main/resources/templates/tool/gen/gen.html b/ruoyi-admin/src/main/resources/templates/tool/gen/gen.html
index 42df3c328..716076eb9 100644
--- a/ruoyi-admin/src/main/resources/templates/tool/gen/gen.html
+++ b/ruoyi-admin/src/main/resources/templates/tool/gen/gen.html
@@ -22,6 +22,9 @@
 								<span>-</span>
 								<input type="text" class="time-input" id="endTime" placeholder="结束时间" name="params[endTime]"/>
 							</li>
+							<li>
+								数据源：<input type="text" id="dataSource" name="params[dataSource]" th:value="${dataSource}"/>
+							</li>
 							<li>
 								<a class="btn btn-primary btn-rounded btn-sm" onclick="$.table.search()"><i class="fa fa-search"></i>&nbsp;搜索</a>
 								<a class="btn btn-warning btn-rounded btn-sm" onclick="$.form.reset()"><i class="fa fa-refresh"></i>&nbsp;重置</a>
@@ -45,7 +48,7 @@
 	<div th:include="include :: footer"></div>
 	<script type="text/javascript">
 		var prefix = ctx + "tool/gen";
-	
+
 		$(function() {
 		    var options = {
 		        url: prefix + "/list",
@@ -95,7 +98,7 @@
 		// 生成代码
 		function genCode(tableName) {
 		    $.modal.confirm("确定要生成" + tableName + "表代码吗？", function() {
-		        location.href = prefix + "/genCode/" + tableName;
+		        location.href = prefix + "/genCode/" + tableName + "/" + $("#dataSource").val();
 		        layer.msg('执行成功,正在生成代码请稍后…', { icon: 1 });
 		    })
 		}
@@ -107,8 +110,8 @@
 		        $.modal.alertWarning("请选择要生成的数据");
 		        return;
 		    }
-		    $.modal.confirm("确认要生成选中的" + rows.length + "条数据吗?", function() {
-		        location.href = prefix + "/batchGenCode?tables=" + rows;
+			$.modal.confirm("确认要生成选中的" + rows.length + "条数据吗?", function() {
+		        location.href = prefix + "/batchGenCode?tables=" + rows + "&dataSource=" + $("#dataSource").val();
 		        layer.msg('执行成功,正在生成代码请稍后…', { icon: 1 });
 		    });
 		}