Can I dynamically set a Postgres schema for a table with a plugin?

[andrewhare]

I am trying to create a plugin that allows me to dynamically prefix a Postgres schema to any tables in the query. The schema name is pulled out of whatever context.Context was supplied with the query.

type DynamicSchema struct {
	Schema func(context.Context) string
}

func (DynamicSchema) Name() string {
	return "dynamic_schema"
}

func (s DynamicSchema) Initialize(db *gorm.DB) error {
	db.Callback().Create().Before("gorm:create").Register("dynamic_schema:set", s.set)
	db.Callback().Query().Before("gorm:query").Register("dynamic_schema:set", s.set)
	db.Callback().Delete().Before("gorm:delete").Register("dynamic_schema:set", s.set)
	db.Callback().Update().Before("gorm:update").Register("dynamic_schema:set", s.set)

	return nil
}

func (s DynamicSchema) set(db *gorm.DB) {
	ctx := db.Statement.Context

	db.Statement.Table = s.prepend(ctx, db.Statement.Table)

	for i, j := range db.Statement.Joins {
		db.Statement.Joins[i].Name = s.prepend(ctx, j.Name)
	}
}

func (s DynamicSchema) prepend(ctx context.Context, v string) string {
	return s.Schema(ctx) + "." + v
}

This works for SELECT tables but not for tables referenced in a JOIN. Is it possible to set the name of the joined table?

[andrewhare]

I was able to do what I wanted by changing the for loop:

for _, j := range db.Statement.Joins {
	r := db.Statement.Schema.Relationships.Relations[j.Name]
	r.FieldSchema.Table = s.prepend(ctx, r.FieldSchema.Table)
}

Will this break anything else as the query is built?