From 3ab2be6a3dd71c138cb386e6c8ece4644c2a2b70 Mon Sep 17 00:00:00 2001
From: Matthew McNeely <matthew.mcneely@gmail.com>
Date: Tue, 6 May 2025 21:40:56 -0400
Subject: [PATCH] Add cleanup func for returning dgo client to pool

---
 client.go     | 23 ++++++++++++++++++-----
 query_test.go |  3 ++-
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/client.go b/client.go
index f567ba5..3013c0c 100644
--- a/client.go
+++ b/client.go
@@ -29,7 +29,7 @@ type Client interface {
 	DropData(context.Context) error
 	QueryRaw(context.Context, string) ([]byte, error)
 
-	DgraphClient() (*dgo.Dgraph, error)
+	DgraphClient() (*dgo.Dgraph, func(), error)
 }
 
 const (
@@ -373,10 +373,23 @@ func (c client) Close() {
 	c.pool.close()
 }
 
-// DgraphClient returns the underlying Dgraph client.
-func (c client) DgraphClient() (*dgo.Dgraph, error) {
-	client, err := c.pool.get()
-	return client, err
+// DgraphClient returns a Dgraph client from the pool and a cleanup function to put it back.
+//
+// Usage:
+//
+//	client, cleanup, err := c.DgraphClient()
+//	if err != nil { ... }
+//	defer cleanup()
+//
+// The cleanup function is safe to call even if client is nil or err is not nil.
+func (c client) DgraphClient() (client *dgo.Dgraph, cleanup func(), err error) {
+	client, err = c.pool.get()
+	cleanup = func() {
+		if client != nil {
+			c.pool.put(client)
+		}
+	}
+	return client, cleanup, err
 }
 
 type clientPool struct {
diff --git a/query_test.go b/query_test.go
index a0db4de..b8279a0 100644
--- a/query_test.go
+++ b/query_test.go
@@ -261,8 +261,9 @@ func TestVectorSimilaritySearch(t *testing.T) {
 			vectorVar := "[0.51, 0.39, 0.29, 0.19, 0.09]"
 			query := dg.NewQuery().Model(&testItem).RootFunc("similar_to(vector, 1, $vec)")
 
-			dgo, err := client.DgraphClient()
+			dgo, cleanup, err := client.DgraphClient()
 			require.NoError(t, err)
+			defer cleanup()
 			tx := dg.NewReadOnlyTxn(dgo)
 			err = tx.Query(query).Vars("similar_to($vec: string)", map[string]string{"$vec": vectorVar}).Scan()
 			require.NoError(t, err)