From 4e51b4b36cc51de04dbcdb1f95ac2f4bf5bf6c64 Mon Sep 17 00:00:00 2001
From: sylwiaszunejko <sylwia.szunejko@gmail.com>
Date: Wed, 19 Jul 2023 12:10:36 +0200
Subject: [PATCH] Provide tablets metadata

---
 host_source.go     |  2 ++
 metadata_scylla.go | 70 ++++++++++++++++++++++++++++++++++++++++++----
 session.go         | 13 +++++++++
 3 files changed, 80 insertions(+), 5 deletions(-)

diff --git a/host_source.go b/host_source.go
index 6daee05d2..a87dd8e34 100644
--- a/host_source.go
+++ b/host_source.go
@@ -1100,6 +1100,8 @@ func refreshTablets(r *ringDescriber) error {
 	r.session.ring.setTablets(tablets)
 	r.session.policy.SetTablets(tablets)
 
+	r.session.schemaDescriber.refreshTabletsSchema()
+
 	return nil
 }
 
diff --git a/metadata_scylla.go b/metadata_scylla.go
index 7efdeb9b4..66f05961c 100644
--- a/metadata_scylla.go
+++ b/metadata_scylla.go
@@ -1,3 +1,4 @@
+//go:build !cassandra || scylla
 // +build !cassandra scylla
 
 // Copyright (c) 2015 The gocql Authors. All rights reserved.
@@ -132,6 +133,29 @@ type IndexMetadata struct {
 	Options      map[string]string
 }
 
+// TabletsMetadata holds metadata for tablet list
+type TabletsMetadata struct {
+	Tablets []*TabletMetadata
+}
+
+// TabletMetadata holds metadata for single tablet
+type TabletMetadata struct {
+	KeyspaceName string
+	TableId      UUID
+	LastToken    int64
+	TableName    string
+	TabletCount  int
+	NewReplicas  []ReplicaMetadata
+	Replicas     []ReplicaMetadata
+	Stage        string
+}
+
+// TabletMetadata holds metadata for single replica
+type ReplicaMetadata struct {
+	HostId  UUID
+	ShardId int
+}
+
 const (
 	IndexKindCustom = "CUSTOM"
 )
@@ -215,20 +239,22 @@ func columnKindFromSchema(kind string) (ColumnKind, error) {
 	}
 }
 
-// queries the cluster for schema information for a specific keyspace
+// queries the cluster for schema information for a specific keyspace and for tablets
 type schemaDescriber struct {
 	session *Session
 	mu      sync.Mutex
 
-	cache map[string]*KeyspaceMetadata
+	cache        map[string]*KeyspaceMetadata
+	tabletsCache *TabletsMetadata
 }
 
 // creates a session bound schema describer which will query and cache
-// keyspace metadata
+// keyspace metadata and tablets metadata
 func newSchemaDescriber(session *Session) *schemaDescriber {
 	return &schemaDescriber{
-		session: session,
-		cache:   map[string]*KeyspaceMetadata{},
+		session:      session,
+		cache:        map[string]*KeyspaceMetadata{},
+		tabletsCache: &TabletsMetadata{},
 	}
 }
 
@@ -252,6 +278,40 @@ func (s *schemaDescriber) getSchema(keyspaceName string) (*KeyspaceMetadata, err
 	return metadata, nil
 }
 
+func (s *schemaDescriber) getTabletsSchema() *TabletsMetadata {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	metadata := s.tabletsCache
+
+	return metadata
+}
+
+func (s *schemaDescriber) refreshTabletsSchema() {
+	tablets := s.session.getTablets()
+	s.tabletsCache.Tablets = []*TabletMetadata{}
+
+	for _, tablet := range tablets {
+		t := &TabletMetadata{}
+		t.KeyspaceName = tablet.KeyspaceName()
+		t.TableId = tablet.TableId()
+		t.TableName = tablet.TableName()
+		t.Stage = tablet.Stage()
+		t.LastToken = tablet.LastToken()
+		t.TabletCount = tablet.TabletCount()
+		t.Replicas = []ReplicaMetadata{}
+		for _, replica := range tablet.Replicas() {
+			t.Replicas = append(t.Replicas, ReplicaMetadata{replica.hostId, replica.shardId})
+		}
+		t.NewReplicas = []ReplicaMetadata{}
+		for _, replica := range tablet.NewReplicas() {
+			t.NewReplicas = append(t.NewReplicas, ReplicaMetadata{replica.hostId, replica.shardId})
+		}
+
+		s.tabletsCache.Tablets = append(s.tabletsCache.Tablets, t)
+	}
+}
+
 // clears the already cached keyspace metadata
 func (s *schemaDescriber) clearSchema(keyspaceName string) {
 	s.mu.Lock()
diff --git a/session.go b/session.go
index 2e64a5c2c..e0099c7e7 100644
--- a/session.go
+++ b/session.go
@@ -582,6 +582,18 @@ func (s *Session) KeyspaceMetadata(keyspace string) (*KeyspaceMetadata, error) {
 	return s.schemaDescriber.getSchema(keyspace)
 }
 
+// TabletsMetadata returns the metadata about tablets
+func (s *Session) TabletsMetadata() (*TabletsMetadata, error) {
+	// fail fast
+	if s.Closed() {
+		return nil, ErrSessionClosed
+	} else if !s.hostSource.UsesTablets() {
+		return nil, ErrTabletsNotUsed
+	}
+
+	return s.schemaDescriber.getTabletsSchema(), nil
+}
+
 func (s *Session) getConn() *Conn {
 	hosts := s.ring.allHosts()
 	for _, host := range hosts {
@@ -2370,6 +2382,7 @@ var (
 	ErrNoKeyspace           = errors.New("no keyspace provided")
 	ErrKeyspaceDoesNotExist = errors.New("keyspace does not exist")
 	ErrNoMetadata           = errors.New("no metadata available")
+	ErrTabletsNotUsed       = errors.New("tablets not used")
 )
 
 type ErrProtocol struct{ error }