Prepared statements with multishard DDL

pkg/datashard/datashard.go

@@ -8,6 +8,7 @@ import (
 	"github.com/pg-sharding/spqr/pkg/config"
 	"github.com/pg-sharding/spqr/pkg/conn"
 	"github.com/pg-sharding/spqr/pkg/models/kr"
+	"github.com/pg-sharding/spqr/pkg/models/spqrerror"
 	"github.com/pg-sharding/spqr/pkg/prepstatement"
 	"github.com/pg-sharding/spqr/pkg/shard"
 	"github.com/pg-sharding/spqr/pkg/spqrlog"
@@ -556,7 +557,10 @@ func (sh *Conn) TxStatus() txstatus.TXStatus {
 // Returns:
 // - bool: true if the prepared statement exists, false otherwise.
 // - *shard.PreparedStatementDescriptor: the prepared statement descriptor, or nil if it does not exist.
-func (srv *Conn) HasPrepareStatement(hash uint64) (bool, *prepstatement.PreparedStatementDescriptor) {
+func (srv *Conn) HasPrepareStatement(hash uint64, shardId uint) (bool, *prepstatement.PreparedStatementDescriptor) {
+	if shardId != srv.ID() {
+		return false, nil
+	}
 	rd, ok := srv.stmtDesc[hash]
 	return ok, rd
 }
@@ -571,9 +575,14 @@ func (srv *Conn) HasPrepareStatement(hash uint64) (bool, *prepstatement.Prepared
 //
 // Returns:
 // - None.
-func (srv *Conn) StorePrepareStatement(hash uint64, def *prepstatement.PreparedStatementDefinition, rd *prepstatement.PreparedStatementDescriptor) {
+func (srv *Conn) StorePrepareStatement(hash uint64, shardId uint, def *prepstatement.PreparedStatementDefinition, rd *prepstatement.PreparedStatementDescriptor) error {
+	id := srv.ID()
+	if shardId != id {
+		return spqrerror.Newf(spqrerror.SPQR_ROUTING_ERROR, "Cannot store stmt for shard \"%d\" in shard \"%d\"", shardId, id)
+	}
 	srv.stmtDef[hash] = def
 	srv.stmtDesc[hash] = rd
+	return nil
 }
 
 // AuthRule returns the backend auth configuration of the Conn object.

pkg/prepstatement/prepstatement.go

@@ -17,8 +17,8 @@ type PreparedStatementDescriptor struct {
 }
 
 type PreparedStatementHolder interface {
-	HasPrepareStatement(hash uint64) (bool, *PreparedStatementDescriptor)
-	StorePrepareStatement(hash uint64, d *PreparedStatementDefinition, rd *PreparedStatementDescriptor)
+	HasPrepareStatement(hash uint64, shardId uint) (bool, *PreparedStatementDescriptor)
+	StorePrepareStatement(hash uint64, shardId uint, d *PreparedStatementDefinition, rd *PreparedStatementDescriptor) error
 }
 
 type PreparedStatementMapper interface {

router/frontend/frontend_test.go

@@ -191,11 +191,7 @@ func TestFrontendXProto(t *testing.T) {
 
 	/* query Router */
 
-	qr.EXPECT().DataShardsRoutes().AnyTimes().Return([]*kr.ShardKey{
-		&kr.ShardKey{
-			Name: "sh1",
-		},
-	})
+	qr.EXPECT().DataShardsRoutes().AnyTimes().Return([]*kr.ShardKey{{Name: "sh1"}})
 
 	cl.EXPECT().Server().AnyTimes().Return(srv)
 	cl.EXPECT().MaintainParams().AnyTimes().Return(false)
@@ -271,8 +267,8 @@ func TestFrontendXProto(t *testing.T) {
 	res := false
 	rd := &prepstatement.PreparedStatementDescriptor{}
 
-	srv.EXPECT().HasPrepareStatement(gomock.Any()).DoAndReturn(func(interface{}) (interface{}, interface{}) { return res, rd }).AnyTimes()
-	srv.EXPECT().StorePrepareStatement(gomock.Any(), gomock.Any(), gomock.Any()).Do(func(interface{}, interface{}, interface{}) {
+	srv.EXPECT().HasPrepareStatement(gomock.Any(), gomock.Any()).DoAndReturn(func(interface{}, interface{}) (interface{}, interface{}) { return res, rd }).AnyTimes()
+	srv.EXPECT().StorePrepareStatement(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Do(func(interface{}, interface{}, interface{}, interface{}) {
 		res = true
 		rd.ParamDesc = &pgproto3.ParameterDescription{}
 		rd.RowDesc = &pgproto3.RowDescription{}