Fix check if shard exists in qdb (#776)

* add shards to qdb from config, check if shard exists in qdb * fix regress error msg * check if cfg is specified
pg-sharding · Sep 23, 2024 · 1a2a55d · 1a2a55d
1 parent 2989f55
commit 1a2a55d
Show file tree

Hide file tree

Showing 8 changed files with 115 additions and 5 deletions.
diff --git a/cmd/coordinator/main.go b/cmd/coordinator/main.go
@@ -49,7 +49,11 @@ var rootCmd = &cobra.Command{
 			return fmt.Errorf("init frontend TLS: %w", err)
 		}
 
-		coordinator := provider.NewCoordinator(frTLS, db)
+		coordinator, err := provider.NewCoordinator(frTLS, db)
+		if err != nil {
+			return err
+		}
+
 		app := app.NewApp(coordinator)
 		return app.Run(true)
 	},

diff --git a/cmd/router/main.go b/cmd/router/main.go
@@ -242,7 +242,11 @@ var runCmd = &cobra.Command{
 						return fmt.Errorf("init frontend TLS: %w", err)
 					}
 
-					coordinator := provider.NewCoordinator(frTLS, db)
+					coordinator, err := provider.NewCoordinator(frTLS, db)
+					if err != nil {
+						return err
+					}
+
 					app := coordApp.NewApp(coordinator)
 					return app.Run(false)
 				}(); err != nil {

diff --git a/coordinator/provider/coordinator.go b/coordinator/provider/coordinator.go
@@ -276,11 +276,26 @@ func (qc *qdbCoordinator) watchRouters(ctx context.Context) {
 	}
 }
 
-func NewCoordinator(tlsconfig *tls.Config, db qdb.XQDB) *qdbCoordinator {
+func NewCoordinator(tlsconfig *tls.Config, db qdb.XQDB) (*qdbCoordinator, error) {
+	if config.CoordinatorConfig().ShardDataCfg != "" {
+		shards, err := config.LoadShardDataCfg(config.CoordinatorConfig().ShardDataCfg)
+		if err != nil {
+			return nil, err
+		}
+
+		if shards != nil {
+			for id, cfg := range shards.ShardsData {
+				if err := db.AddShard(context.TODO(), qdb.NewShard(id, cfg.Hosts)); err != nil {
+					return nil, err
+				}
+			}
+		}
+	}
+
 	return &qdbCoordinator{
 		db:        db,
 		tlsconfig: tlsconfig,
-	}
+	}, nil
 }
 
 // TODO : unit tests

diff --git a/qdb/etcdqdb.go b/qdb/etcdqdb.go
@@ -847,6 +847,10 @@ func (q *EtcdQDB) GetShard(ctx context.Context, id string) (*Shard, error) {
 		return nil, err
 	}
 
+	if len(resp.Kvs) == 0 {
+		return nil, spqrerror.Newf(spqrerror.SPQR_NO_DATASHARD, "shard \"%s\" not found", id)
+	}
+
 	shardInfo := &Shard{
 		ID: id,
 	}

diff --git a/test/regress/conf/coordinator.yaml b/test/regress/conf/coordinator.yaml
@@ -3,3 +3,4 @@ host: regress_coordinator
 coordinator_port: 7002
 grpc_api_port: 7003
 qdb_addr: '[regress_qdb_0_1]:2379'
+shard_data: '/spqr/test/feature/conf/shard_data.yaml'
diff --git a/test/regress/conf/shard_data.yaml b/test/regress/conf/shard_data.yaml
@@ -0,0 +1,13 @@
+shards:
+  sh1:
+    db: regress
+    usr: regress
+    pwd: 12345678
+    hosts: 
+    - 'spqr_shard_1:6432'
+  sh2:
+    db: regress
+    usr: regress
+    pwd: 12345678
+    hosts: 
+    - 'spqr_shard_2:6432'
diff --git a/test/regress/tests/coordinator/expected/coordinator.out b/test/regress/tests/coordinator/expected/coordinator.out
@@ -5,3 +5,55 @@ REGISTER ROUTER r1 ADDRESS regress_router:7000;
  router -> r1-regress_router:7000
 (1 row)
 
+CREATE DISTRIBUTION ds1 COLUMN TYPES integer;
+    add distribution    
+------------------------
+ distribution id -> ds1
+(1 row)
+
+CREATE KEY RANGE krid1 FROM 1 ROUTE TO sh1 FOR DISTRIBUTION ds1;
+ add key range 
+---------------
+ bound -> 1
+(1 row)
+
+CREATE KEY RANGE krid2 FROM 11 ROUTE TO sh1 FOR DISTRIBUTION ds1;
+ add key range 
+---------------
+ bound -> 11
+(1 row)
+
+SHOW key_ranges;
+ Key range ID | Shard ID | Distribution ID | Lower bound 
+--------------+----------+-----------------+-------------
+ krid1        | sh1      | ds1             | 1
+ krid2        | sh1      | ds1             | 11
+(2 rows)
+
+DROP KEY RANGE krid1;
+    drop key range     
+-----------------------
+ key range id -> krid1
+(1 row)
+
+CREATE KEY RANGE krid2 FROM 11 ROUTE TO sh2 FOR DISTRIBUTION ds1;
+ERROR:  key range krid2 already present in qdb.
+SHOW key_ranges;
+ Key range ID | Shard ID | Distribution ID | Lower bound 
+--------------+----------+-----------------+-------------
+ krid2        | sh1      | ds1             | 11
+(1 row)
+
+CREATE KEY RANGE krid2 FROM 33 ROUTE TO nonexistentshard FOR DISTRIBUTION ds1;
+ERROR:  shard "nonexistentshard" not found.
+DROP DISTRIBUTION ALL CASCADE;
+   drop distribution    
+------------------------
+ distribution id -> ds1
+(1 row)
+
+DROP KEY RANGE ALL;
+ drop key range 
+----------------
+(0 rows)
+
diff --git a/test/regress/tests/coordinator/sql/coordinator.sql b/test/regress/tests/coordinator/sql/coordinator.sql
@@ -1,2 +1,19 @@
 -- UNREGISTER ROUTER ALL
-REGISTER ROUTER r1 ADDRESS regress_router:7000;
+REGISTER ROUTER r1 ADDRESS regress_router:7000;
+
+CREATE DISTRIBUTION ds1 COLUMN TYPES integer;
+CREATE KEY RANGE krid1 FROM 1 ROUTE TO sh1 FOR DISTRIBUTION ds1;
+CREATE KEY RANGE krid2 FROM 11 ROUTE TO sh1 FOR DISTRIBUTION ds1;
+
+SHOW key_ranges;
+
+DROP KEY RANGE krid1;
+
+CREATE KEY RANGE krid2 FROM 11 ROUTE TO sh2 FOR DISTRIBUTION ds1;
+
+SHOW key_ranges;
+
+CREATE KEY RANGE krid2 FROM 33 ROUTE TO nonexistentshard FOR DISTRIBUTION ds1;
+
+DROP DISTRIBUTION ALL CASCADE;
+DROP KEY RANGE ALL;