Possibilidade de criar os índices extras via CLI

cmd/cmd.go

@@ -67,6 +67,20 @@ var dropCmd = &cobra.Command{
 	},
 }
 
+var extraIndexesCmd = &cobra.Command{
+	Use:   "extra-indexes index-1 [index-2 …]",
+	Short: "Creates extra indexes in the JSON fields",
+	Args:  cobra.MinimumNArgs(1),
+	RunE: func(_ *cobra.Command, args []string) error {
+		db, err := loadDatabase()
+		if err != nil {
+			return fmt.Errorf("could not find database: %w", err)
+		}
+		defer db.Close()
+		return db.ExtraIndexes(args)
+	},
+}
+
 func addDataDir(c *cobra.Command) *cobra.Command {
 	c.Flags().StringVarP(&dir, "directory", "d", defaultDataDir, "directory of the downloaded files")
 	return c
@@ -80,7 +94,7 @@ func addDatabase(c *cobra.Command) *cobra.Command {
 
 // CLI returns the root command from Cobra CLI tool.
 func CLI() *cobra.Command {
-	for _, c := range []*cobra.Command{createCmd, dropCmd} {
+	for _, c := range []*cobra.Command{createCmd, dropCmd, extraIndexesCmd} {
 		addDatabase(c)
 	}
 	for _, c := range []*cobra.Command{
@@ -92,6 +106,7 @@ func CLI() *cobra.Command {
 		checkCLI(),
 		createCmd,
 		dropCmd,
+		extraIndexesCmd,
 		transformCLI(),
 		sampleCLI(),
 		mirrorCLI(),

cmd/db.go

@@ -22,6 +22,7 @@ type database interface {
 	CreateCompanies([][]string) error
 	PostLoad() error
 	MetaSave(string, string) error
+	ExtraIndexes(idxs []string) error
 	// api
 	GetCompany(string) (string, error)
 	MetaRead(string) (string, error)

db/mongodb.go

@@ -223,3 +223,28 @@ func (m *MongoDB) GetCompany(cnpj string) (string, error) {
 	}
 	return string(b), nil
 }
+
+func (m *MongoDB) ExtraIndexes(idxs []string) error {
+	log.Output(1, "Creating the indexes…")
+	c := m.db.Collection(companyTableName)
+	var i []mongo.IndexModel
+	for _, v := range idxs {
+		if strings.Contains(v, "qsa_") {
+			v = strings.ReplaceAll(strings.Replace(v, "qsa_", "", 1), "_", "")
+			v = fmt.Sprintf("%s_%s", "qsa", v)
+		}
+		if strings.Contains(v, "secundario") && strings.Contains(v, "cnae") {
+			v = fmt.Sprintf("cnaesecundarios_%s", strings.Split(v, "secundarios_")[1])
+		}
+		i = append(i, mongo.IndexModel{
+			Keys: bson.D{{Key: v, Value: 1}},
+		})
+	}
+	r, err := c.Indexes().CreateMany(m.ctx, i)
+	if err != nil {
+		return fmt.Errorf("error creating indexes: %w", err)
+	}
+	log.Output(1, fmt.Sprintf("%d indexes successfully created in the collection %s", len(r), companyTableName))
+	return nil
+
+}

db/postgres.go

@@ -212,3 +212,35 @@ func NewPostgreSQL(uri, schema string) (PostgreSQL, error) {
 	}
 	return p, nil
 }
+
+func (p *PostgreSQL) ExtraIndexes(idxs []string) error {
+	c := 0
+	for _, v := range idxs {
+		t := fmt.Sprintf("((json->'%s'))", v)
+		name := "json_"
+
+		if strings.Contains(v, "qsa") && !strings.Contains(v, "cnae") {
+			v = strings.Split(v, ".")[1]
+			t = fmt.Sprintf("(jsonb_extract_path(json, 'qsa', '%s') jsonb_ops)", v)
+			name += "qsa_"
+		}
+
+		if !strings.Contains(v, "qsa") && strings.Contains(v, "cnae") {
+			v = strings.Split(v, ".")[1]
+			t = fmt.Sprintf("(jsonb_extract_path(json, 'cnaes_secundarios', '%s') jsonb_ops)", v)
+			name += "cnaes_secundarios_"
+		}
+		q := fmt.Sprintf(
+			"CREATE INDEX IF NOT EXISTS idx_%s%s ON %s USING GIN %s;",
+			name, v, p.CompanyTableName, t,
+		)
+		if _, err := p.pool.Exec(context.Background(), q); err != nil {
+			return fmt.Errorf("error to create indexe %s: %w", v, err)
+		}
+		c += 1
+	}
+
+	log.Output(1, fmt.Sprintf("%d Indexes successfully created in the table %s", c, p.CompanyTableName))
+
+	return nil
+}