Add replicated relations to constraints check (#1556)

Author: EinKrebs

Makefile

@@ -95,7 +95,7 @@ pooler_d_run:
 ####################### TESTS #######################
 
 unittest:
-	go test ./cmd/... ./pkg/... ./router/... ./coordinator/... ./yacc/console...
+	go test -timeout 120s ./cmd/... ./pkg/... ./router/... ./coordinator/... ./yacc/console...
 	go test -race -count 20 -timeout 30s ./qdb/...
 
 regress_local: proxy_2sh_run

pkg/coord/clustered_coord.go

@@ -946,15 +946,31 @@ func (qc *ClusteredCoordinator) checkKeyRangeMove(ctx context.Context, req *kr.B
 		}
 	}
 
-	deferrable, constraintName, err := datatransfers.CheckConstraints(ctx, sourceConn, rels)
+	exists, err := qc.qdb.CheckDistribution(ctx, distributions.REPLICATED)
+	if err != nil {
+		return fmt.Errorf("error checking for replicated distribution: %s", err)
+	}
+	replRels := []string{}
+	if exists {
+		replDs, err := qc.GetDistribution(ctx, distributions.REPLICATED)
+		if err != nil {
+			return fmt.Errorf("error getting replicated distribution: %s", err)
+		}
+		replRels = make([]string, 0, len(replDs.Relations))
+		for _, r := range replDs.Relations {
+			replRels = append(replRels, r.GetFullName())
+		}
+	}
+
+	deferrable, constraintName, err := datatransfers.CheckConstraints(ctx, sourceConn, rels, replRels)
 	if err != nil {
 		return spqrerror.Newf(spqrerror.SPQR_TRANSFER_ERROR, "error checking table constraints on source shard: %s", err)
 	}
 	if !deferrable {
 		return spqrerror.Newf(spqrerror.SPQR_TRANSFER_ERROR, "found non-deferrable constraint or constraint referencing non-distributed table on source shard: \"%s\"", constraintName)
 	}
 
-	deferrable, constraintName, err = datatransfers.CheckConstraints(ctx, destConn, rels)
+	deferrable, constraintName, err = datatransfers.CheckConstraints(ctx, destConn, rels, replRels)
 	if err != nil {
 		return spqrerror.Newf(spqrerror.SPQR_TRANSFER_ERROR, "error checking table constraints on destination shard: %s", err)
 	}

pkg/datatransfers/data_transfers.go

@@ -574,22 +574,32 @@ func CheckColumnExists(ctx context.Context, conn *pgx.Conn, relName, schema, col
 // Parameters:
 // - ctx (context.Context): the context for database operations;
 // - conn (*pgx.Conn): the connection to the database;
-// - relNames (string): the list of relations, for which to check for constraints;
+// - dsRels (string): the list of distributed relations, for which to check for constraints;
+// - rpRels (string): the list of replicated relations, which can also be referred to by distributed relations;
 //
 // Returns:
 // - bool: true if no such constraints are found, false otherwise;
 // - string: name of a constraint, if found;
 // - error: an error if there was a problem executing the query.
-func CheckConstraints(ctx context.Context, conn *pgx.Conn, relNames []string) (bool, string, error) {
-	if len(relNames) == 0 {
+func CheckConstraints(ctx context.Context, conn *pgx.Conn, dsRels []string, rpRels []string) (bool, string, error) {
+	if len(dsRels) == 0 {
 		return true, "", nil
 	}
-	relOidList := make([]string, len(relNames))
-	for i, relName := range relNames {
-		relOidList[i] = fmt.Sprintf("'%s'::regclass::oid", relName)
+	dsRelOidList := make([]string, len(dsRels))
+	for i, relName := range dsRels {
+		dsRelOidList[i] = fmt.Sprintf("'%s'::regclass::oid", relName)
 	}
-	relOids := strings.Join(relOidList, ", ")
-	rows, err := conn.Query(ctx, fmt.Sprintf(`SELECT conname FROM pg_constraint WHERE conrelid IN (%s) and confrelid != 0 and (condeferrable=false or not (confrelid IN (%s))) LIMIT 1`, relOids, relOids))
+	dsRelOids := strings.Join(dsRelOidList, ", ")
+	rpRelsClause := ""
+	if len(rpRels) > 0 {
+		refRelOidList := make([]string, len(rpRels))
+		for i, relName := range rpRels {
+			refRelOidList[i] = fmt.Sprintf("'%s'::regclass::oid", relName)
+		}
+		rpRelOids := strings.Join(refRelOidList, ", ")
+		rpRelsClause = fmt.Sprintf(" and not (confrelid IN (%s))", rpRelOids)
+	}
+	rows, err := conn.Query(ctx, fmt.Sprintf(`SELECT conname FROM pg_constraint WHERE conrelid IN (%s) and confrelid != 0 and (condeferrable=false or not (confrelid IN (%s)))%s LIMIT 1`, dsRelOids, dsRelOids, rpRelsClause))
 	if err != nil {
 		return false, "", err
 	}

pkg/mock/pgx/mock_pgxconn_iface.go

@@ -1319,6 +1319,20 @@ func (q *EtcdQDB) GetDistribution(ctx context.Context, id string) (*Distribution
 	return distrib, nil
 }
 
+// TODO : unit tests
+func (q *EtcdQDB) CheckDistribution(ctx context.Context, id string) (bool, error) {
+	spqrlog.Zero.Debug().
+		Str("id", id).
+		Msg("etcdqdb: get distribution by id")
+
+	resp, err := q.cli.Get(ctx, distributionNodePath(id), clientv3.WithCountOnly())
+	if err != nil {
+		return false, err
+	}
+
+	return resp.Count == 1, nil
+}
+
 // TODO : unit tests
 func (q *EtcdQDB) GetRelationDistribution(ctx context.Context, relName *rfqn.RelationFQN) (*Distribution, error) {
 	spqrlog.Zero.Debug().

qdb/etcdqdb.go

@@ -845,6 +845,16 @@ func (q *MemQDB) GetDistribution(_ context.Context, id string) (*Distribution, e
 	}
 }
 
+// TODO : unit tests
+func (q *MemQDB) CheckDistribution(_ context.Context, id string) (bool, error) {
+	spqrlog.Zero.Debug().Str("id", id).Msg("memqdb: check distribution")
+	q.mu.RLock()
+	defer q.mu.RUnlock()
+
+	_, ok := q.Distributions[id]
+	return ok, nil
+}
+
 func (q *MemQDB) GetRelationDistribution(_ context.Context, relation *rfqn.RelationFQN) (*Distribution, error) {
 	spqrlog.Zero.Debug().Str("relation", relation.RelationName).Msg("memqdb: get distribution for table")
 	q.mu.RLock()

qdb/memqdb.go

@@ -70,7 +70,7 @@ type QDB interface {
 	ListDistributions(ctx context.Context) ([]*Distribution, error)
 	DropDistribution(ctx context.Context, id string) error
 	GetDistribution(ctx context.Context, id string) (*Distribution, error)
-	// TODO: fix this by passing FQRN (fully qualified relation name (+schema))
+	CheckDistribution(ctx context.Context, id string) (bool, error)
 	GetRelationDistribution(ctx context.Context, relation *rfqn.RelationFQN) (*Distribution, error)
 
 	// Reference relations

qdb/mock/qdb.go

@@ -310,3 +310,31 @@ Feature: Redistribution test
     """
     found non-deferrable constraint or constraint referencing non-distributed table on destination shard: "xmove_pkey_fkey"
     """
+  
+Scenario: REDISTRIBUTE KEY RANGE allows constraints on reference tables
+    When I execute SQL on host "coordinator"
+    """
+    CREATE KEY RANGE kr1 FROM 0 ROUTE TO sh1 FOR DISTRIBUTION ds1;
+    CREATE REFERENCE TABLE ref;
+    """
+    Then command return code should be "0"
+    
+    When I run SQL on host "shard1"
+    """
+    CREATE TABLE ref(id INT PRIMARY KEY);
+    CREATE TABLE xMove(w_id INT, ext_id INT REFERENCES ref(id));
+    """
+    Then command return code should be "0"
+
+    When I run SQL on host "shard2"
+    """
+    CREATE TABLE ref(id INT PRIMARY KEY);
+    CREATE TABLE xMove(w_id INT, ext_id INT REFERENCES ref(id));
+    """
+    Then command return code should be "0"
+
+    When I run SQL on host "coordinator" with timeout "150" seconds
+    """
+    REDISTRIBUTE KEY RANGE kr1 TO sh2 CHECK;
+    """
+    Then command return code should be "0"