renaming repository to shard, collapsing some tables into columnsjc-praefect-datamodel-with-pg

3 files changed, 99 insertions, 52 deletions

diff --git a/internal/praefect/database/migrations/1_initial.up.sql b/internal/praefect/database/migrations/1_initial.up.sql
index 077baf03b..cc2138979 100644
--- a/internal/praefect/database/migrations/1_initial.up.sql
+++ b/internal/praefect/database/migrations/1_initial.up.sql
@@ -1,44 +1,29 @@
-CREATE TABLE storages (
-  name TEXT PRIMARY KEY
-);
-
-CREATE TABLE nodes (
-  address TEXT PRIMARY KEY
-);
-
-CREATE TABLE storage_nodes (
+CREATE TABLE IF NOT EXISTS node_storages (
   id SERIAL PRIMARY KEY,
-  storage_name TEXT REFERENCES storages(name),
-  node_address TEXT REFERENCES nodes(address)
+  address TEXT NOT NULL,
+  storage_name TEXT NOT NULL
 );
 
-CREATE TABLE repositories (
+CREATE TABLE IF NOT EXISTS shards (
   relative_path TEXT PRIMARY KEY,
-  "primary" INTEGER REFERENCES storage_nodes(id)
+  "primary" INTEGER REFERENCES node_storages(id)
 );
 
-CREATE TABLE repository_secondaries (
-  repository_relative_path TEXT REFERENCES repositories(relative_path),
-  storage_node_id INTEGER REFERENCES storage_nodes(id),
-  PRIMARY KEY(repository_relative_path, storage_node_id)
+CREATE TABLE IF NOT EXISTS shard_secondaries (
+  shard_relative_path TEXT REFERENCES shards(relative_path),
+  node_storage_id INTEGER REFERENCES node_storages(id),
+  PRIMARY KEY(shard_relative_path, node_storage_id)
 );
 
-INSERT INTO storages ("name") VALUES ('default');
-INSERT INTO storages ("name") VALUES ('backup1');
-INSERT INTO storages ("name") VALUES ('backup2');
-
-INSERT INTO nodes (address) VALUES('tcp://127.0.0.1:9999');
-INSERT INTO nodes (address) VALUES('tcp://127.0.0.1:9998');
-
-INSERT INTO storage_nodes (id, node_address, storage_name) VALUES(1, 'tcp://127.0.0.1:9999', 'default');
-INSERT INTO storage_nodes (id, node_address, storage_name) VALUES(2, 'tcp://127.0.0.1:9998', 'default');
-INSERT INTO storage_nodes (id, node_address, storage_name) VALUES(3, 'tcp://127.0.0.1:9999', 'backup1');
-INSERT INTO storage_nodes (id, node_address, storage_name) VALUES(4, 'tcp://127.0.0.1:9998', 'backup2');
+INSERT INTO node_storages (id, address, storage_name) VALUES(1, 'tcp://127.0.0.1:9999', 'default');
+INSERT INTO node_storages (id, address, storage_name) VALUES(2, 'tcp://127.0.0.1:9998', 'default');
+INSERT INTO node_storages (id, address, storage_name) VALUES(3, 'tcp://127.0.0.1:9999', 'backup1');
+INSERT INTO node_storages (id, address, storage_name) VALUES(4, 'tcp://127.0.0.1:9998', 'backup2');
 
-INSERT INTO repositories(relative_path, "primary")
+INSERT INTO shards(relative_path, "primary")
   VALUES('@hashed/2c/62/2c624232cdd221771294dfbb310aca000a0df6ac8b66b696d90ef06fdefb64a3.git', 1);
 
-INSERT INTO repository_secondaries(repository_relative_path, storage_node_id) VALUES
+INSERT INTO shard_secondaries(shard_relative_path, storage_node_id) VALUES
   ('@hashed/2c/62/2c624232cdd221771294dfbb310aca000a0df6ac8b66b696d90ef06fdefb64a3.git', 2),
   ('@hashed/2c/62/2c624232cdd221771294dfbb310aca000a0df6ac8b66b696d90ef06fdefb64a3.git', 3),
   ('@hashed/2c/62/2c624232cdd221771294dfbb310aca000a0df6ac8b66b696d90ef06fdefb64a3.git', 4);
\ No newline at end of file
diff --git a/internal/praefect/database/sql_datastore.go b/internal/praefect/database/sql_datastore.go
index 70f134d62..1a2ccac8f 100644
--- a/internal/praefect/database/sql_datastore.go
+++ b/internal/praefect/database/sql_datastore.go
@@ -1,12 +1,12 @@
 package database
 
 import (
+	"errors"
 	"fmt"
 	"os"
 
 	"database/sql"
 
-	_ "github.com/lib/pq"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/models"
 )
 
@@ -28,7 +28,7 @@ func (sd *SQLDatastore) GetSecondaries(relativePath string) ([]models.StorageNod
 	var secondaries []models.StorageNode
 
 	rows, err := sd.db.Query(`
-		SELECT storage_nodes.* FROM repositories
+		SELECT storage_nodes.id, storage_nodes.address, storage_nodes.storage_name FROM repositories
 			INNER JOIN repository_secondaries ON repositories.relative_path = repository_secondaries.repository_relative_path
 			INNER JOIN storage_nodes ON storage_nodes.id = repository_secondaries.storage_node_id WHERE repositories.relative_path = $1
 	`, relativePath)
@@ -39,7 +39,7 @@ func (sd *SQLDatastore) GetSecondaries(relativePath string) ([]models.StorageNod
 
 	for rows.Next() {
 		var s models.StorageNode
-		err = rows.Scan(&s.ID, &s.Name, &s.Address)
+		err = rows.Scan(&s.ID, &s.StorageName, &s.Address)
 		if err != nil {
 			return nil, err
 		}
@@ -49,18 +49,88 @@ func (sd *SQLDatastore) GetSecondaries(relativePath string) ([]models.StorageNod
 	return secondaries, nil
 }
 
-func (sd *SQLDatastore) GetPrimary(relativePath string) (models.StorageNode, error) {
+func (sd *SQLDatastore) GetPrimary(relativePath string) (*models.StorageNode, error) {
 
 	row := sd.db.QueryRow(`
-		SELECT storage_nodes.* FROM repositories
+		SELECT storage_nodes.id, storage_nodes.adddress, storage_nodes.storage_name FROM repositories
 			INNER JOIN storage_nodes ON repositories.primary = storage_nodes.id
 			WHERE repositories.relative_path = $1
 	`, relativePath)
 
 	var s models.StorageNode
-	if err := row.Scan(&s.ID, &s.Name, &s.Address); err != nil {
+	if err := row.Scan(&s.ID, &s.Address, &s.StorageName); err != nil {
 		return nil, err
 	}
 
-	return s, nil
+	return &s, nil
+}
+
+func (sd *SQLDatastore) SetPrimary(relativePath, storageNodeID int) error {
+	res, err := sd.db.Exec(`UPDATE repositories SET "primary" = ? WHERE relative_path= ?`, storageNodeID, relativePath)
+	if err != nil {
+		return err
+	}
+
+	if n, err := res.RowsAffected(); err != nil {
+		return err
+	} else if n == 0 {
+		return errors.New("repository does not exist")
+	}
+
+	return nil
+}
+
+func (sd *SQLDatastore) GetDefaultPrimary(storage string) (*models.StorageNode, error) {
+	row := sd.db.QueryRow("SELECT storage_nodes.id, storage_nodes.address, storage_nodes.storage_name from stroage_nodes where storage_name = ?", storage)
+
+	var s models.StorageNode
+	if err := row.Scan(&s.ID, &s.Address, &s.StorageName); err != nil {
+		return nil, err
+	}
+
+	return &s, nil
+}
+
+func (sd *SQLDatastore) AddSecondary(relativePath string, storageNodeID int) error {
+	res, err := sd.db.Exec("INSERT INTO repository_secondaries (repository_relative_path, node_storage_id) VALUES(?, ?)", relativePath, storageNodeID)
+	if err != nil {
+		return err
+	}
+
+	if n, err := res.RowsAffected(); err != nil {
+		return err
+	} else if n == 0 {
+		return errors.New("secondary already exists")
+	}
+
+	return nil
+}
+
+func (sd *SQLDatastore) RemoveSecondary(relativePath string, storageNodeID int) error {
+	res, err := sd.db.Exec("DELETE FROM repository_secondaries (repository_relative_path, node_storage_id) VALUES(?, ?)", relativePath, storageNodeID)
+	if err != nil {
+		return err
+	}
+
+	if n, err := res.RowsAffected(); err != nil {
+		return err
+	} else if n == 0 {
+		return errors.New("secondary did not exist")
+	}
+
+	return nil
+}
+
+func (sd *SQLDatastore) GetShard(relativePath string) (*models.Shard, error) {
+	primary, err := sd.GetPrimary(relativePath)
+	if err != nil {
+		return nil, fmt.Errorf("getting primary: %v", err)
+	}
+
+	secondaries, err := sd.GetSecondaries(relativePath)
+	if err != nil {
+		return nil, fmt.Errorf("getting secondaries: %v", err)
+	}
+
+	return &models.Shard{RelativePath: relativePath, Primary: *primary, Secondaries: secondaries}, nil
 }
diff --git a/internal/praefect/models/node.go b/internal/praefect/models/node.go
index 64b72337b..851a28a4e 100644
--- a/internal/praefect/models/node.go
+++ b/internal/praefect/models/node.go
@@ -1,21 +1,13 @@
 package models
 
-type Storage struct {
-	Name string
-}
-
-type Node struct {
-	Address string
-}
-
 type StorageNode struct {
-	ID int
-	Name string
-	Address string
+	ID          int
+	StorageName string
+	Address     string
 }
 
-type Repository struct {
+type Shard struct {
 	RelativePath string
-	Primary *StorageNode
-	Secondaries []*StorageNode
-}
\ No newline at end of file
+	Primary      StorageNode
+	Secondaries  []StorageNode
+}