From 2841e91c0dde793364bad84c95becfd32810c8a1 Mon Sep 17 00:00:00 2001
From: lanchengx <lanchengx@qq.com>
Date: Thu, 26 Aug 2021 17:23:06 +0800
Subject: [PATCH] Filter incomplete schema & Set the schema to the connection.

---
 .../impl/UnicastDatabaseBackendHandler.java   | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/shardingsphere-proxy/shardingsphere-proxy-backend/src/main/java/org/apache/shardingsphere/proxy/backend/text/data/impl/UnicastDatabaseBackendHandler.java b/shardingsphere-proxy/shardingsphere-proxy-backend/src/main/java/org/apache/shardingsphere/proxy/backend/text/data/impl/UnicastDatabaseBackendHandler.java
index 632769a2f1e25..68307c5fb004a 100644
--- a/shardingsphere-proxy/shardingsphere-proxy-backend/src/main/java/org/apache/shardingsphere/proxy/backend/text/data/impl/UnicastDatabaseBackendHandler.java
+++ b/shardingsphere-proxy/shardingsphere-proxy-backend/src/main/java/org/apache/shardingsphere/proxy/backend/text/data/impl/UnicastDatabaseBackendHandler.java
@@ -30,6 +30,7 @@
 
 import java.sql.SQLException;
 import java.util.Collection;
+import java.util.Optional;
 
 /**
  * Database backend handler with unicast schema.
@@ -49,12 +50,18 @@ public final class UnicastDatabaseBackendHandler implements DatabaseBackendHandl
     
     @Override
     public ResponseHeader execute() throws SQLException {
-        String schemaName = null == backendConnection.getSchemaName() ? getFirstSchemaName() : backendConnection.getSchemaName();
+        String originSchema = backendConnection.getSchemaName();
+        String schemaName = null == originSchema ? getFirstSchemaName() : originSchema;
         if (!ProxyContext.getInstance().getMetaData(schemaName).isComplete()) {
             throw new RuleNotExistedException();
         }
-        databaseCommunicationEngine = databaseCommunicationEngineFactory.newTextProtocolInstance(sqlStatementContext, sql, backendConnection);
-        return databaseCommunicationEngine.execute();
+        try {
+            backendConnection.setCurrentSchema(schemaName);
+            databaseCommunicationEngine = databaseCommunicationEngineFactory.newTextProtocolInstance(sqlStatementContext, sql, backendConnection);
+            return databaseCommunicationEngine.execute();
+        } finally {
+            backendConnection.setCurrentSchema(originSchema);
+        }
     }
     
     private String getFirstSchemaName() {
@@ -62,7 +69,11 @@ private String getFirstSchemaName() {
         if (schemaNames.isEmpty()) {
             throw new NoDatabaseSelectedException();
         }
-        return schemaNames.iterator().next();
+        Optional<String> result = schemaNames.stream().filter(each -> ProxyContext.getInstance().getMetaData(each).isComplete()).findFirst();
+        if (!result.isPresent()) {
+            throw new RuleNotExistedException();
+        }
+        return result.get();
     }
     
     @Override