Fix drop distrib in etcd (#565)

* Fix drop distrib in etcd * Add feature test scenario
pg-sharding · Mar 21, 2024 · 1f360b1 · 1f360b1
1 parent 0e77564
commit 1f360b1
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 2 deletions.
diff --git a/coordinator/provider/coordinator.go b/coordinator/provider/coordinator.go
@@ -4,10 +4,11 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
-	"github.com/pg-sharding/spqr/pkg/models/distributions"
 	"net"
 	"time"
 
+	"github.com/pg-sharding/spqr/pkg/models/distributions"
+
 	"github.com/pg-sharding/spqr/pkg/models/spqrerror"
 
 	"github.com/google/uuid"

diff --git a/qdb/etcdqdb.go b/qdb/etcdqdb.go
@@ -919,13 +919,30 @@ func (q *EtcdQDB) DropDistribution(ctx context.Context, id string) error {
 	case 0:
 		return spqrerror.New(spqrerror.SPQR_SHARDING_RULE_ERROR, "no such distribution present in qdb")
 	case 1:
+
+		var distrib *Distribution
+		if err := json.Unmarshal(resp.Kvs[0].Value, &distrib); err != nil {
+			return err
+		}
+
 		resp, err := q.cli.Delete(ctx, distributionNodePath(id))
 
 		spqrlog.Zero.Debug().
 			Interface("response", resp).
 			Msg("etcdqdb: drop distribution")
 
-		return err
+		if err != nil {
+			return err
+		}
+
+		for _, r := range distrib.Relations {
+			_, err := q.cli.Delete(ctx, relationMappingNodePath(r.Name))
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
 	default:
 		return spqrerror.Newf(spqrerror.SPQR_SHARDING_RULE_ERROR, "too much distributions matched: %d", len(resp.Kvs))
 	}

diff --git a/test/feature/features/coordinator.feature b/test/feature/features/coordinator.feature
@@ -28,6 +28,72 @@ Feature: Coordinator test
     """
     Then command return code should be "0"
 
+
+  Scenario: Add/Remove distribution works
+    When I run SQL on host "coordinator"
+    """
+    CREATE DISTRIBUTION ds1_test COLUMN TYPES integer; 
+    CREATE KEY RANGE krid11 FROM 0 ROUTE TO sh1 FOR DISTRIBUTION ds1_test;
+    CREATE KEY RANGE krid22 FROM 11 ROUTE TO sh2 FOR DISTRIBUTION ds1_test;
+    ALTER DISTRIBUTION ds1_test ATTACH RELATION test1 DISTRIBUTION KEY id;
+    """
+    Then command return code should be "0"
+
+    When I run SQL on host "coordinator"
+    """
+    SHOW distributions;
+    """
+    Then SQL result should match json
+    """
+    [
+        {
+            "Distribution ID":"ds1_test",
+            "Column types":"integer"
+        }
+    ]
+    """
+
+    When I run SQL on host "coordinator"
+    """
+    DROP DISTRIBUTION ds1_test CASCADE
+    """
+    Then command return code should be "0"
+
+
+    When I run SQL on host "coordinator"
+    """
+    SHOW distributions;
+    """
+    Then SQL result should match json
+    """
+    []
+    """
+
+    When I run SQL on host "coordinator"
+    """
+    CREATE DISTRIBUTION ds1_test COLUMN TYPES integer; 
+    CREATE KEY RANGE krid11 FROM 0 ROUTE TO sh1 FOR DISTRIBUTION ds1_test;
+    CREATE KEY RANGE krid22 FROM 11 ROUTE TO sh2 FOR DISTRIBUTION ds1_test;
+    ALTER DISTRIBUTION ds1_test ATTACH RELATION test1 DISTRIBUTION KEY id;
+    """
+    Then command return code should be "0"
+
+
+    When I run SQL on host "coordinator"
+    """
+    SHOW distributions;
+    """
+    Then SQL result should match json
+    """
+    [
+        {
+            "Distribution ID":"ds1_test",
+            "Column types":"integer"
+        }
+    ]
+    """
+
+
   Scenario: Add/Remove router works
     When I run SQL on host "coordinator"
     """