This code can be simplified using models. You only have to change the initial code to start with a model class instead of a plain dataset.