dataflow_example.go

dataflow_example.go

// note this stays outside of updateDatabaseFn struct. 
// fields in pardo has to be json-serielizable so you cannot have db instance as field
var database *gorm.DB
type updateDatabaseFn struct {
 Env string `json:"env"`
 Project string `json:"project"`
}

func init() {
	beam.RegisterDoFn(reflect.TypeOf((*updateDatabaseFn)(nil)).Elem())
}

// initialize connection here
func (f *updateDatabaseFn) StartBundle(_ctx context.Context) error {
 var err error
 database, err = initializeDatabase(context.Background(), f.Env, f.Project)
 if err != nil {
   return errors.Wrap(err, "failed to connect to database")
 }
 return nil
}

// use your db connection here
func (f *updateDatabaseFn) ProcessElement(ctx context.Context, data SomeData) {
 // some data processing here
 err := insertRecords(database, data)
 // .. more logic here