Data transfer

cmd/mover/main.go

@@ -38,7 +38,6 @@ func (p *ProxyW) Write(bt []byte) (int, error) {
 	spqrlog.Zero.Debug().
 		Bytes("bytes", bt).
 		Msg("got bytes")
-
 	return p.w.Write(bt)
 }
 
@@ -69,11 +68,10 @@ func moveData(ctx context.Context, from, to *pgx.Conn, keyRange kr.KeyRange, key
 SELECT table_schema, table_name
 FROM information_schema.columns
 WHERE column_name=$1;
-`, key.Entries()[0])
+`, key.Entries()[0].Column)
 	if err != nil {
 		return err
 	}
-
 	var ress []MoveTableRes
 
 	for rows.Next() {
@@ -103,21 +101,8 @@ WHERE column_name=$1;
 			w: w,
 		}
 
-		ch := make(chan struct{})
-
-		go func() {
-			spqrlog.Zero.Debug().Msg("sending rows to dest shard")
-			_, err := txTo.Conn().PgConn().CopyFrom(ctx,
-				r, fmt.Sprintf("COPY %s.%s FROM STDIN", v.TableSchema, v.TableName))
-			if err != nil {
-				spqrlog.Zero.Error().Err(err).Msg("copy in failed")
-			}
-
-			ch <- struct{}{}
-		}()
-
 		qry := fmt.Sprintf("copy (delete from %s.%s WHERE %s >= %s and %s <= %s returning *) to stdout", v.TableSchema, v.TableName,
-			key.Entries()[0], keyRange.LowerBound, key.Entries()[0], keyRange.UpperBound)
+			key.Entries()[0].Column, keyRange.LowerBound, key.Entries()[0].Column, keyRange.UpperBound)
 
 		spqrlog.Zero.Debug().
 			Str("query", qry).
@@ -132,9 +117,14 @@ WHERE column_name=$1;
 			spqrlog.Zero.Error().Err(err).Msg("error closing pipe")
 		}
 
-		spqrlog.Zero.Debug().Msg("copy cmd executed")
+		_, err = txTo.Conn().PgConn().CopyFrom(ctx,
+			r, fmt.Sprintf("COPY %s.%s FROM STDIN", v.TableSchema, v.TableName))
+		if err != nil {
+			spqrlog.Zero.Debug().Msg("copy in failed")
+			return err
+		}
 
-		<-ch
+		spqrlog.Zero.Debug().Msg("copy cmd executed")
 	}
 
 	_ = txTo.Commit(ctx)
@@ -165,6 +155,10 @@ func main() {
 		return
 	}
 
+	//entrys := []shrule.ShardingRuleEntry{*shrule.NewShardingRuleEntry("id", "nohash")}
+	//my_rule := shrule.NewShardingRule("r1", "fast", entrys)
+	//db.AddShardingRule(context.TODO(), shrule.ShardingRuleToDB(my_rule))
+
 	shRule, err := db.GetShardingRule(context.TODO(), *shkey)
 	if err != nil {
 		spqrlog.Zero.Error().Err(err).Msg("")

coordinator/provider/coordinator.go

@@ -6,6 +6,7 @@ import (
 	"net"
 	"time"
 
+	"github.com/pg-sharding/spqr/pkg/datatransfers"
 	"github.com/pg-sharding/spqr/pkg/meta"
 	"github.com/pg-sharding/spqr/pkg/models/topology"
 	"github.com/pg-sharding/spqr/pkg/shard"
@@ -227,6 +228,38 @@ func NewCoordinator(db qdb.QDB) *qdbCoordinator {
 		db: db,
 	}
 
+	ranges, err := db.ListKeyRanges(context.TODO())
+	if err != nil {
+		spqrlog.Zero.Error().
+			Err(err).
+			Msg("faild to list key ranges")
+	}
+
+	for _, r := range ranges {
+		tx, err := db.GetTransferTx(context.TODO(), r.KeyRangeID)
+		if tx == nil || err != nil {
+			continue
+		}
+		if tx.ToStatus == "commit" {
+			datatransfers.ResolvePreparedTransaction(context.TODO(), tx.FromShardId, tx.FromTxName, true)
+			tem := kr.MoveKeyRange{
+				ShardId: tx.ToShardId,
+				Krid:    r.KeyRangeID,
+			}
+			err = cc.Move(context.TODO(), &tem)
+			if err != nil {
+				spqrlog.Zero.Error().Err(err).Msg("failed to move key range")
+			}
+		} else {
+			datatransfers.ResolvePreparedTransaction(context.TODO(), tx.ToShardId, tx.ToTxName, false)
+			datatransfers.ResolvePreparedTransaction(context.TODO(), tx.FromShardId, tx.FromTxName, false)
+		}
+		err = db.RemoveTransferTx(context.TODO(), r.KeyRangeID)
+		if err != nil {
+			spqrlog.Zero.Error().Err(err).Msg("error removing from qdb")
+		}
+	}
+
 	go cc.watchRouters(context.TODO())
 	return cc
 }
@@ -593,6 +626,7 @@ func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error
 		Str("shard-id", req.ShardId).
 		Msg("qdb coordinator move key range")
 
+	//move i qdb
 	if err := qc.traverseRouters(ctx, func(cc *grpc.ClientConn) error {
 		cl := routerproto.NewKeyRangeServiceClient(cc)
 		lockResp, err := cl.LockKeyRange(ctx, &routerproto.LockKeyRangeRequest{
@@ -634,6 +668,15 @@ func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error
 		}
 	}()
 
+	//move between shards
+	keyRange, _ := qc.db.GetKeyRange(ctx, req.Krid)
+	shardingRules, _ := qc.ListShardingRules(ctx)
+	err = datatransfers.MoveKeys(ctx, keyRange.ShardID, req.ShardId, *keyRange, shardingRules, &qc.db)
+	if err != nil {
+		spqrlog.Zero.Error().Msg("failed to move rows")
+		return err
+	}
+
 	krmv.ShardID = req.ShardId
 	if err := ops.ModifyKeyRangeWithChecks(ctx, qc.db, kr.KeyRangeFromDB(krmv)); err != nil {
 		// TODO: check if unlock here is ok
@@ -663,7 +706,6 @@ func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error
 				Interface("response", unlockResp).
 				Msg("unlock key range response")
 		}()
-
 		moveResp, err := cl.MoveKeyRange(ctx, &routerproto.MoveKeyRangeRequest{
 			KeyRange: kr.KeyRangeFromDB(krmv).ToProto(),
 		})
@@ -674,7 +716,6 @@ func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error
 	}); err != nil {
 		return err
 	}
-	// do phyc keys move
 	return nil
 }
 

docker/coordinator/cfg.yaml

@@ -2,3 +2,4 @@ addr: '[spqr_coordinator]:7002'
 http_addr: 'spqr_coordinator:7003'
 qdb_addr: '[spqr_qdb_0_1]:2379'
 log_level: INFO
+shard_data: '/spqr/docker/coordinator/shard_data.yaml'

docker/coordinator/shard_data.yaml

@@ -0,0 +1,13 @@
+shards:
+  sh1:
+    db: db1
+    user: user1
+    Password: 12345678
+    host: 'spqr_shard_1'
+    port: '6432'
+  sh2:
+    db: db1
+    user: user1
+    Password: 12345678
+    host: 'spqr_shard_2'
+    port: '6432'

docker/shard/bin/setup

@@ -15,6 +15,7 @@ EOF
 cat >> /var/lib/postgresql/13/main/postgresql.conf <<-EOF
 listen_addresses = '*'
 port = 6432
+max_prepared_transactions = 5
 EOF
 
 sudo -u postgres /usr/lib/postgresql/13/bin/pg_ctl -D /var/lib/postgresql/13/main/ start

docker/tests/bin/move.sh

@@ -4,6 +4,21 @@ sleep 20
 
 set -ex
 
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c 'ADD SHARDING RULE r1 COLUMNS w_id;' || {
+    echo "ERROR: tests failed"
+    exit 1
+}
+
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c 'ADD KEY RANGE krid1 FROM 1 TO 10 ROUTE TO sh1;' || {
+    echo "ERROR: tests failed"
+    exit 1
+}
+
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c 'ADD KEY RANGE krid2 FROM 11 TO 20 ROUTE TO sh2;' || {
+    echo "ERROR: tests failed"
+    exit 1
+}
+
 psql "host=spqr_router_1_1 sslmode=disable user=user1 dbname=db1 port=6432" -c 'CREATE TABLE xMove(w_id INT, s TEXT)' || {
 	echo "ERROR: tests failed"
 	exit 1
@@ -24,23 +39,44 @@ psql "host=spqr_router_1_1 sslmode=disable user=user1 dbname=db1 port=6432" -c "
 	exit 1
 }
 
-psql "host=spqr_router_1_1 sslmode=disable user=user1 dbname=db1 port=7432" -c "LOCK KEY RANGE krid2;" || {
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c "LOCK KEY RANGE krid2;" || {
 	echo "ERROR: tests failed"
 	exit 1
 }
 
-psql "host=spqr_router_1_1 sslmode=disable user=user1 dbname=db1 port=7432" -c "MOVE KEY RANGE krid2 to sh1;" || {
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c "MOVE KEY RANGE krid2 to sh1;" || {
 	echo "ERROR: tests failed"
 	exit 1
 }
 
-psql "host=spqr_router_1_1 sslmode=disable user=user1 dbname=db1 port=7432" -c "UNLOCK KEY RANGE krid2;" || {
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c "UNLOCK KEY RANGE krid2;" || {
 	echo "ERROR: tests failed"
 	exit 1
 }
 
+out=$(psql "host=spqr_shard_1 sslmode=disable user=user1 dbname=db1 port=6432" -c "select * from xMove")
+test "$out" = " w_id |  s  
+------+-----
+    1 | 001
+   11 | 002
+(2 rows)"
+
+out=$(psql "host=spqr_shard_2 sslmode=disable user=user1 dbname=db1 port=6432" -c "select * from xMove")
+test "$out" = " w_id | s 
+------+---
+(0 rows)"
+
 psql "host=spqr_router_1_1 sslmode=disable user=user1 dbname=db1 port=6432" -c "select * from xMove where w_id = 11" || {
 	echo "ERROR: tests failed"
 	exit 1
 }
 
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c "DROP KEY RANGE ALL;" || {
+	echo "ERROR: tests failed"
+	exit 1
+}
+psql "host=spqr_coordinator sslmode=disable user=user1 dbname=db1 port=7002" -c 'DROP SHARDING RULE ALL;' || {
+    echo "ERROR: tests failed"
+    exit 1
+}
+

examples/coordinator.yaml

@@ -2,4 +2,5 @@ addr: 'localhost:7002'
 http_addr: 'localhost:7003'
 qdb_addr: 'localhost:2379'
 log_level: INFO
+shard_data: '/spqr/docker/coordinator/shard_data.yaml'
 

pkg/config/coordinator.go

@@ -14,10 +14,11 @@ import (
 var cfgCoordinator Coordinator
 
 type Coordinator struct {
-	LogLevel string `json:"log_level" toml:"log_level" yaml:"log_level"`
-	QdbAddr  string `json:"qdb_addr" toml:"qdb_addr" yaml:"qdb_addr"`
-	HttpAddr string `json:"http_addr" toml:"http_addr" yaml:"http_addr"`
-	Addr     string `json:"addr" toml:"addr" yaml:"addr"`
+	LogLevel     string `json:"log_level" toml:"log_level" yaml:"log_level"`
+	QdbAddr      string `json:"qdb_addr" toml:"qdb_addr" yaml:"qdb_addr"`
+	HttpAddr     string `json:"http_addr" toml:"http_addr" yaml:"http_addr"`
+	Addr         string `json:"addr" toml:"addr" yaml:"addr"`
+	ShardDataCfg string `json:"shard_data" toml:"shard_data" yaml:"shard_data"`
 }
 
 func LoadCoordinatorCfg(cfgPath string) error {

pkg/config/data_transfers.go

@@ -0,0 +1,66 @@
+package config
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"os"
+	"strings"
+
+	"github.com/BurntSushi/toml"
+	"gopkg.in/yaml.v2"
+)
+
+type DatatransferConnections struct {
+	ShardsData map[string]*ShardConnect `json:"shards" toml:"shards" yaml:"shards"`
+}
+
+type ShardConnect struct {
+	Host     string `json:"host" toml:"host" yaml:"host"`
+	Port     string `json:"port" toml:"port" yaml:"port"`
+	DB       string `json:"db" toml:"db" yaml:"db"`
+	User     string `json:"user" toml:"user" yaml:"user"`
+	Password string `json:"password" toml:"password" yaml:"password"`
+}
+
+func LoadShardDataCfg(cfgPath string) (*DatatransferConnections, error) {
+	var cfg DatatransferConnections
+	file, err := os.Open(cfgPath)
+	if err != nil {
+		return &cfg, err
+	}
+
+	defer func(file *os.File) {
+		err := file.Close()
+		if err != nil {
+			log.Fatalf("failed to close config file: %v", err)
+		}
+	}(file)
+
+	if err := initShardDataConfig(file, &cfg); err != nil {
+		return &cfg, err
+	}
+
+	configBytes, err := json.MarshalIndent(cfg, "", "  ")
+	if err != nil {
+		return &cfg, err
+	}
+
+	fmt.Println("Running config:", string(configBytes))
+
+	return &cfg, nil
+}
+
+func initShardDataConfig(file *os.File, cfg *DatatransferConnections) error {
+	if strings.HasSuffix(file.Name(), ".toml") {
+		_, err := toml.NewDecoder(file).Decode(&cfg)
+		return err
+	}
+	if strings.HasSuffix(file.Name(), ".yaml") {
+		return yaml.NewDecoder(file).Decode(&cfg)
+	}
+	if strings.HasSuffix(file.Name(), ".json") {
+		return json.NewDecoder(file).Decode(&cfg)
+	}
+	return fmt.Errorf("unknown config format type: %s. Use .toml, .yaml or .json suffix in filename", file.Name())
+}