From ceadc67a2facbd875bc552f604d54d693f529937 Mon Sep 17 00:00:00 2001
From: Xin Hao <haoxinst@gmail.com>
Date: Sun, 4 Feb 2024 16:39:18 +0800
Subject: [PATCH] Add ShowSpaces and Execute for Session

Add Show spaces and for
---
 session.go      | 26 ++++++++++++++++++++++++++
 session_pool.go |  3 +--
 session_test.go | 21 +++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/session.go b/session.go
index 6a7da2af..58cf48bd 100644
--- a/session.go
+++ b/session.go
@@ -185,6 +185,32 @@ func (session *Session) ExecuteJsonWithParameter(stmt string, params map[string]
 	return resp.([]byte), err
 }
 
+func (session *Session) ExecuteAndCheck(stmt string) (*ResultSet, error) {
+	rs, err := session.Execute(stmt)
+	if err != nil {
+		return nil, err
+	}
+
+	if !rs.IsSucceed() {
+		errMsg := rs.GetErrorMsg()
+		return nil, fmt.Errorf("Fail to execute query. %s", errMsg)
+	}
+
+	return rs, nil
+}
+
+func (session *Session) ShowSpaces() ([]SpaceName, error) {
+	rs, err := session.ExecuteAndCheck("SHOW SPACES;")
+	if err != nil {
+		return nil, err
+	}
+
+	var names []SpaceName
+	rs.Scan(&names)
+
+	return names, nil
+}
+
 func (session *Session) reConnect() error {
 	newConnection, err := session.connPool.getIdleConn()
 	if err != nil {
diff --git a/session_pool.go b/session_pool.go
index 43f1aef9..a3e7a604 100644
--- a/session_pool.go
+++ b/session_pool.go
@@ -278,8 +278,7 @@ func (pool *SessionPool) ExecuteAndCheck(q string) (*ResultSet, error) {
 		return nil, err
 	}
 
-	errCode := rs.GetErrorCode()
-	if errCode != ErrorCode_SUCCEEDED {
+	if !rs.IsSucceed() {
 		errMsg := rs.GetErrorMsg()
 		return nil, fmt.Errorf("Fail to execute query. %s", errMsg)
 	}
diff --git a/session_test.go b/session_test.go
index 925463ef..214345c5 100644
--- a/session_test.go
+++ b/session_test.go
@@ -117,3 +117,24 @@ func TestSession_Recover(t *testing.T) {
 	}
 	assert.Equal(t, 1, pool.getActiveConnCount()+pool.getIdleConnCount())
 }
+
+func TestSession_ShowSpaces(t *testing.T) {
+	config := GetDefaultConf()
+	host := HostAddress{address, port}
+	pool, err := NewConnectionPool([]HostAddress{host}, config, DefaultLogger{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sess, err := pool.GetSession("root", "nebula")
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, 1, pool.getActiveConnCount()+pool.getIdleConnCount())
+
+	spaceNames, err := sess.ShowSpaces()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.LessOrEqual(t, 1, len(spaceNames))
+}