From 1af4d3f020035601cfd5d6361e8deee3e66f01b1 Mon Sep 17 00:00:00 2001
From: SwallowGG <1558143046@qq.com>
Date: Wed, 6 Sep 2023 16:37:11 +0800
Subject: [PATCH] fix dm function bug

---
 .../plugin/sqlserver/SqlServerMetaData.java   | 38 +++++++++++++++++--
 .../ai/chat2db/spi/sql/IDriverManager.java    |  7 +++-
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java
index 35af5786..d83b3135 100644
--- a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java
+++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java
@@ -68,6 +68,10 @@ public class SqlServerMetaData extends DefaultMetaService implements MetaData {
         = "SELECT type_desc, OBJECT_NAME(object_id) AS FunctionName, OBJECT_DEFINITION(object_id) AS "
         + "definition FROM sys.objects WHERE type_desc IN(%s) and name = '%s' ;";
 
+
+    private static String OBJECT_SQL
+            = "SELECT name FROM sys.objects WHERE type = '%s' and SCHEMA_ID = SCHEMA_ID('%s');";
+
     @Override
     public Function function(Connection connection, @NotEmpty String databaseName, String schemaName,
         String functionName) {
@@ -87,8 +91,22 @@ public class SqlServerMetaData extends DefaultMetaService implements MetaData {
 
     @Override
     public List<Function> functions(Connection connection, String databaseName, String schemaName) {
-        List<Function> functions = SQLExecutor.getInstance().functions(connection, databaseName, schemaName);
-        return functions.stream().map(function -> removeVersion(function)).collect(Collectors.toList());
+        List<Function> functions = new ArrayList<>();
+//        List<Function> functions = SQLExecutor.getInstance().functions(connection, databaseName, schemaName);
+//        return functions.stream().map(function -> removeVersion(function)).collect(Collectors.toList());
+        String sql = String.format(OBJECT_SQL,"FN", schemaName);
+        return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
+            while (resultSet.next()) {
+                Function function = new Function();
+                function.setDatabaseName(databaseName);
+                function.setSchemaName(schemaName);
+                if (resultSet.next()) {
+                    function.setFunctionName(resultSet.getString("name"));
+                }
+                functions.add(function);
+            }
+            return functions;
+        });
     }
 
     private Function removeVersion(Function function) {
@@ -103,8 +121,20 @@ public class SqlServerMetaData extends DefaultMetaService implements MetaData {
 
     @Override
     public List<Procedure> procedures(Connection connection, String databaseName, String schemaName) {
-        List<Procedure> procedures = SQLExecutor.getInstance().procedures(connection, databaseName, schemaName);
-        return procedures.stream().map(procedure -> removeVersion(procedure)).collect(Collectors.toList());
+        List<Procedure> procedures = new ArrayList<>();
+        String sql = String.format(OBJECT_SQL,"P", schemaName);
+        return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
+            while (resultSet.next()) {
+                Procedure procedure = new Procedure();
+                procedure.setDatabaseName(databaseName);
+                procedure.setSchemaName(schemaName);
+                if (resultSet.next()) {
+                    procedure.setProcedureName(resultSet.getString("name"));
+                }
+                procedures.add(procedure);
+            }
+            return procedures;
+        });
     }
 
     private Procedure removeVersion(Procedure procedure) {
diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/IDriverManager.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/IDriverManager.java
index 83ae4e2b..78f65c3f 100644
--- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/IDriverManager.java
+++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/IDriverManager.java
@@ -78,13 +78,16 @@ public class IDriverManager {
         if (url == null) {
             throw new SQLException("The url cannot be null", "08001");
         }
-        SQLException reason = null;
         DriverEntry driverEntry = DRIVER_ENTRY_MAP.get(driver.getJdbcDriver());
         if (driverEntry == null) {
             driverEntry = getJDBCDriver(driver);
         }
         try {
-            return driverEntry.getDriver().connect(url, info);
+            Connection connection = driverEntry.getDriver().connect(url, info);
+            if(connection == null){
+                throw new ConnectionException("driverEntry.getDriver().connect return null",null);
+            }
+            return connection;
         } catch (SQLException var7) {
             Connection con = tryConnectionAgain(driverEntry, url, info);
             if (con != null) {