Implements pre/post-load routine in the new ETL

Author: cuducos

transformnext/transform.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/cuducos/minha-receita/download"
 	"github.com/schollz/progressbar/v3"
 	"golang.org/x/sync/errgroup"
 )
@@ -24,8 +25,22 @@ const (
 	MaxParallelDBQueries = 8
 )
 
+var extraIndexes = [...]string{
+	"cnae_fiscal",
+	"cnaes_secundarios.codigo",
+	"codigo_municipio",
+	"codigo_municipio_ibge",
+	"codigo_natureza_juridica",
+	"qsa.cnpj_cpf_do_socio",
+	"uf",
+}
+
 type database interface {
+	PreLoad() error
 	CreateCompanies([][]string) error
+	PostLoad() error
+	CreateExtraIndexes([]string) error
+	MetaSave(string, string) error
 }
 
 func sources() map[string]*source { // all but Estabelecimentos (this one is loaded later on)
@@ -65,27 +80,35 @@ func newProgressBar(label string, srcs int) (*progressbar.ProgressBar, error) {
 	return bar, bar.RenderBlank()
 }
 
-func Cleanup() error {
-	return filepath.WalkDir(os.TempDir(), func(pth string, d fs.DirEntry, err error) error {
-		if !d.IsDir() {
-			return nil
-		}
-		if !strings.HasPrefix(d.Name(), "minha-receita-") {
-			return nil
-		}
-		part := strings.Split(d.Name(), "-")
-		if len(part) != 4 {
-			return nil
-		}
-		if _, err := time.Parse("20060102150405", part[2]); err != nil {
-			return nil
-		}
-		fmt.Printf("Removing %s\n", pth)
-		return os.RemoveAll(pth)
-	})
+func saveUpdatedAt(db database, dir string) error {
+	slog.Info("Saving the updated at date to the database…")
+	p := filepath.Join(dir, download.FederalRevenueUpdatedAt)
+	v, err := os.ReadFile(p)
+	if err != nil {
+		return fmt.Errorf("error reading %s: %w", p, err)
+
+	}
+	return db.MetaSave("updated-at", string(v))
+}
+
+func postLoad(db database) error {
+	slog.Info("Consolidating the database…")
+	if err := db.PostLoad(); err != nil {
+		return err
+	}
+	slog.Info("Database consolidated!")
+	slog.Info("Creating indexes…")
+	if err := db.CreateExtraIndexes(extraIndexes[:]); err != nil {
+		return err
+	}
+	slog.Info("Indexes created!")
+	return nil
 }
 
 func Transform(dir string, db database, batch, maxDB int, privacy bool) error {
+	if err := db.PreLoad(); err != nil {
+		return err
+	}
 	srcs := sources()
 	tmp, err := os.MkdirTemp("", fmt.Sprintf("minha-receita-%s-*", time.Now().Format("20060102150405")))
 	if err != nil {
@@ -121,5 +144,31 @@ func Transform(dir string, db database, batch, maxDB int, privacy bool) error {
 	if err := g.Wait(); err != nil {
 		return err
 	}
-	return writeJSONs(ctx, srcs, kv, db, maxDB, batch, dir, privacy)
+	if err := writeJSONs(ctx, srcs, kv, db, maxDB, batch, dir, privacy); err != nil {
+		return err
+	}
+	if err := postLoad(db); err != nil {
+		return err
+	}
+	return saveUpdatedAt(db, dir)
+}
+
+func Cleanup() error {
+	return filepath.WalkDir(os.TempDir(), func(pth string, d fs.DirEntry, err error) error {
+		if !d.IsDir() {
+			return nil
+		}
+		if !strings.HasPrefix(d.Name(), "minha-receita-") {
+			return nil
+		}
+		part := strings.Split(d.Name(), "-")
+		if len(part) != 4 {
+			return nil
+		}
+		if _, err := time.Parse("20060102150405", part[2]); err != nil {
+			return nil
+		}
+		fmt.Printf("Removing %s\n", pth)
+		return os.RemoveAll(pth)
+	})
 }