How to implement Custom Policy for DBResolver with Zone or other customized Information

[spencercjh]

I'm currently working on a project where I need to implement a custom policy for DBResolver that takes into account the zone information of each database connection. The goal is to preferentially select a connection from a specific zone when resolving the connection pool.

I've created a custom struct DBWithZone that embeds *gorm.DB, includes a Zone field and implements the gorm.ConnPool interface. I've also implemented a custom policy NearZonePolicy that attempts to select a DBWithZone from the connection pool based on the preferred zone.

However, I've encountered an issue where the gorm.ConnPool in the Resolve method of my custom policy is actually of type *sql.DB, and I can't directly convert it to my DBWithZone type.

Here's a simplified version of my code:

type DBWithZone struct {
 *gorm.DB
 Zone string
}

type NearZonePolicy struct {
 PreferredZone string
}

func (n *NearZonePolicy) Resolve(connPools []gorm.ConnPool) gorm.ConnPool {
 for _, pool := range connPools {
  if dbWithZone, ok := pool.(*DBWithZone); ok {
   if dbWithZone.Zone == n.PreferredZone {
    return dbWithZone.DB
   }
  }
 }
 return connPools[0]
}

In the Resolve method, the type assertion pool.(*DBWithZone) fails because pool is of type *sql.DB.

I'm looking for a way to associate each gorm.ConnPool (or *sql.DB) with its corresponding zone information so that I can implement my custom policy. Is there a recommended way to achieve this with GORM DBResolver? Any guidance would be greatly appreciated.

Maybe we need a factory func to let Gorm know how to create a customized gorm.ConnPool implementation instead of *sql.DB.

Thank you.