Add cleanup func for returning dgo client to pool

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 {

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)