From d44819f1a2de40826fdc33079e7c3a3986537362 Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Fri, 14 Jun 2024 23:29:20 +0800
Subject: [PATCH] fix: panic in CloseStream

---
 client.go      | 4 ++++
 client_test.go | 8 ++++++++
 2 files changed, 12 insertions(+)

diff --git a/client.go b/client.go
index 96c8825..91e7071 100644
--- a/client.go
+++ b/client.go
@@ -272,6 +272,10 @@ func (c *Client) StreamDeleteObject(ctx context.Context, body any) error {
 // using client.StreamWrite or client.StreamWriteObject, we need to call client.CloseStream to let
 // GreptimeDB know that we’ve finished writing and are expecting to receive a response.
 func (c *Client) CloseStream(ctx context.Context) (*gpb.AffectedRows, error) {
+	if c.stream == nil {
+		return &gpb.AffectedRows{}, nil
+	}
+
 	resp, err := c.stream.CloseAndRecv()
 	if err != nil {
 		return nil, err
diff --git a/client_test.go b/client_test.go
index b30f187..c75823c 100644
--- a/client_test.go
+++ b/client_test.go
@@ -873,6 +873,14 @@ func TestStreamWrite(t *testing.T) {
 	}
 }
 
+func TestStreamClose(t *testing.T) {
+	lc := newClient()
+
+	affected, err := lc.CloseStream(context.Background())
+	assert.EqualValues(t, 0, affected.GetValue())
+	assert.Nil(t, err)
+}
+
 func TestStreamUpdate(t *testing.T) {
 	loc, err := time.LoadLocation(timezone)
 	assert.Nil(t, err)