diff --git a/database/gorm/event.go b/database/gorm/event.go index ebcaf5932..0df1e224f 100644 --- a/database/gorm/event.go +++ b/database/gorm/event.go @@ -74,7 +74,7 @@ func (e *Event) ColumnNamesWithDbColumnNames() map[string]string { } func (e *Event) Context() context.Context { - return e.query.instance.Statement.Context + return e.query.ctx } func (e *Event) DestOfMap() map[string]any { diff --git a/database/gorm/query.go b/database/gorm/query.go index e965d6ab5..09051f49c 100644 --- a/database/gorm/query.go +++ b/database/gorm/query.go @@ -635,6 +635,11 @@ func (r *QueryImpl) Select(query any, args ...any) ormcontract.Query { return r.setConditions(conditions) } +func (r *QueryImpl) SetContext(ctx context.Context) { + r.ctx = ctx + r.instance.Statement.Context = ctx +} + func (r *QueryImpl) SharedLock() ormcontract.Query { conditions := r.conditions conditions.sharedLock = true diff --git a/database/orm.go b/database/orm.go index 3a23e06c6..7cb6e96d1 100644 --- a/database/orm.go +++ b/database/orm.go @@ -67,9 +67,9 @@ func (r *OrmImpl) Connection(name string) ormcontract.Orm { } func (r *OrmImpl) DB() (*sql.DB, error) { - db := r.Query().(*databasegorm.QueryImpl) + query := r.Query().(*databasegorm.QueryImpl) - return db.Instance().DB() + return query.Instance().DB() } func (r *OrmImpl) Query() ormcontract.Query { @@ -105,7 +105,19 @@ func (r *OrmImpl) Transaction(txFunc func(tx ormcontract.Transaction) error) err } func (r *OrmImpl) WithContext(ctx context.Context) ormcontract.Orm { - instance, _ := NewOrmImpl(ctx, r.config, r.connection, r.query) + for _, query := range r.queries { + query := query.(*databasegorm.QueryImpl) + query.SetContext(ctx) + } + + query := r.query.(*databasegorm.QueryImpl) + query.SetContext(ctx) - return instance + return &OrmImpl{ + ctx: ctx, + config: r.config, + connection: r.connection, + query: query, + queries: r.queries, + } } diff --git a/database/orm_test.go b/database/orm_test.go index c593710ad..6068f6bc1 100644 --- a/database/orm_test.go +++ b/database/orm_test.go @@ -22,6 +22,10 @@ var connections = []contractsorm.Driver{ contractsorm.DriverSqlserver, } +type contextKey int + +const testContextKey contextKey = 0 + type User struct { orm.Model orm.SoftDeletes @@ -92,8 +96,9 @@ func TestOrmSuite(t *testing.T) { func (s *OrmSuite) SetupTest() { s.orm = &OrmImpl{ - ctx: context.Background(), - query: testMysqlQuery, + connection: contractsorm.DriverMysql.String(), + ctx: context.Background(), + query: testMysqlQuery, queries: map[string]contractsorm.Query{ contractsorm.DriverMysql.String(): testMysqlQuery, contractsorm.DriverPostgresql.String(): testPostgresqlQuery, @@ -194,6 +199,38 @@ func (s *OrmSuite) TestTransactionError() { } } +func (s *OrmSuite) TestWithContext() { + s.orm.Observe(User{}, &UserObserver{}) + ctx := context.WithValue(context.Background(), testContextKey, "with_context_goravel") + user := User{Name: "with_context_name"} + + // Call Query directly + err := s.orm.WithContext(ctx).Query().Create(&user) + s.Nil(err) + s.Equal("with_context_name", user.Name) + s.Equal("with_context_goravel", user.Avatar) + + // Call Connection, then call WithContext + for _, connection := range connections { + user.ID = 0 + user.Avatar = "" + err := s.orm.Connection(connection.String()).WithContext(ctx).Query().Create(&user) + s.Nil(err) + s.Equal("with_context_name", user.Name) + s.Equal("with_context_goravel", user.Avatar) + } + + // Call WithContext, then call Connection + for _, connection := range connections { + user.ID = 0 + user.Avatar = "" + err := s.orm.WithContext(ctx).Connection(connection.String()).Query().Create(&user) + s.Nil(err) + s.Equal("with_context_name", user.Name) + s.Equal("with_context_goravel", user.Avatar) + } +} + type UserObserver struct{} func (u *UserObserver) Retrieved(event contractsorm.Event) error { @@ -202,8 +239,15 @@ func (u *UserObserver) Retrieved(event contractsorm.Event) error { func (u *UserObserver) Creating(event contractsorm.Event) error { name := event.GetAttribute("name") - if name != nil && name.(string) == "observer_name" { - return errors.New("error") + if name != nil { + if name.(string) == "observer_name" { + return errors.New("error") + } + if name.(string) == "with_context_name" { + if avatar := event.Context().Value(testContextKey); avatar != nil { + event.SetAttribute("avatar", avatar.(string)) + } + } } return nil