gosora/cmd/elasticsearch/setup.go

// Work in progress
package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"log"
	"os"
	"strconv"

	c "github.com/Azareal/Gosora/common"
	"github.com/Azareal/Gosora/query_gen"
	"gopkg.in/olivere/elastic.v6"
)

func main() {
	log.Print("Loading the configuration data")
	err := c.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	log.Print("Processing configuration data")
	err = c.ProcessConfig()
	if err != nil {
		log.Fatal(err)
	}

	if c.DbConfig.Adapter != "mysql" && c.DbConfig.Adapter != "" {
		log.Fatal("Only MySQL is supported for upgrades right now, please wait for a newer build of the patcher")
	}

	err = prepMySQL()
	if err != nil {
		log.Fatal(err)
	}

	client, err := elastic.NewClient(elastic.SetErrorLog(log.New(os.Stdout, "ES ", log.LstdFlags)))
	if err != nil {
		log.Fatal(err)
	}
	_, _, err = client.Ping("http://127.0.0.1:9200").Do(context.Background())
	if err != nil {
		log.Fatal(err)
	}

	err = setupIndices(client)
	if err != nil {
		log.Fatal(err)
	}

	err = setupData(client)
	if err != nil {
		log.Fatal(err)
	}
}

func prepMySQL() error {
	return qgen.Builder.Init("mysql", map[string]string{
		"host":      c.DbConfig.Host,
		"port":      c.DbConfig.Port,
		"name":      c.DbConfig.Dbname,
		"username":  c.DbConfig.Username,
		"password":  c.DbConfig.Password,
		"collation": "utf8mb4_general_ci",
	})
}

type ESIndexBase struct {
	Mappings ESIndexMappings `json:"mappings"`
}

type ESIndexMappings struct {
	Doc ESIndexDoc `json:"_doc"`
}

type ESIndexDoc struct {
	Properties map[string]map[string]string `json:"properties"`
}

type ESDocMap map[string]map[string]string

func (d ESDocMap) Add(column string, cType string) {
	d["column"] = map[string]string{"type": cType}
}

func setupIndices(client *elastic.Client) error {
	exists, err := client.IndexExists("topics").Do(context.Background())
	if err != nil {
		return err
	}
	if exists {
		deleteIndex, err := client.DeleteIndex("topics").Do(context.Background())
		if err != nil {
			return err
		}
		if !deleteIndex.Acknowledged {
			return errors.New("delete not acknowledged")
		}
	}

	docMap := make(ESDocMap)
	docMap.Add("tid", "integer")
	docMap.Add("title", "text")
	docMap.Add("content", "text")
	docMap.Add("createdBy", "integer")
	docMap.Add("ip", "ip")
	docMap.Add("suggest", "completion")
	indexBase := ESIndexBase{ESIndexMappings{ESIndexDoc{docMap}}}
	oBytes, err := json.Marshal(indexBase)
	if err != nil {
		return err
	}
	createIndex, err := client.CreateIndex("topics").Body(string(oBytes)).Do(context.Background())
	if err != nil {
		return err
	}
	if !createIndex.Acknowledged {
		return errors.New("not acknowledged")
	}

	exists, err = client.IndexExists("replies").Do(context.Background())
	if err != nil {
		return err
	}
	if exists {
		deleteIndex, err := client.DeleteIndex("replies").Do(context.Background())
		if err != nil {
			return err
		}
		if !deleteIndex.Acknowledged {
			return errors.New("delete not acknowledged")
		}
	}

	docMap = make(ESDocMap)
	docMap.Add("rid", "integer")
	docMap.Add("tid", "integer")
	docMap.Add("content", "text")
	docMap.Add("createdBy", "integer")
	docMap.Add("ip", "ip")
	docMap.Add("suggest", "completion")
	indexBase = ESIndexBase{ESIndexMappings{ESIndexDoc{docMap}}}
	oBytes, err = json.Marshal(indexBase)
	if err != nil {
		return err
	}
	createIndex, err = client.CreateIndex("replies").Body(string(oBytes)).Do(context.Background())
	if err != nil {
		return err
	}
	if !createIndex.Acknowledged {
		return errors.New("not acknowledged")
	}

	return nil
}

type ESTopic struct {
	ID        int    `json:"tid"`
	Title     string `json:"title"`
	Content   string `json:"content"`
	CreatedBy int    `json:"createdBy"`
	IP string `json:"ip"`
}

type ESReply struct {
	ID        int    `json:"rid"`
	TID       int    `json:"tid"`
	Content   string `json:"content"`
	CreatedBy int    `json:"createdBy"`
	IP string `json:"ip"`
}

func setupData(client *elastic.Client) error {
	tcount := 4
	errs := make(chan error)

	go func() {
		tin := make([]chan ESTopic, tcount)
		tf := func(tin chan ESTopic) {
			for {
				topic, more := <-tin
				if !more {
					break
				}
				_, err := client.Index().Index("topics").Type("_doc").Id(strconv.Itoa(topic.ID)).BodyJson(topic).Do(context.Background())
				if err != nil {
					errs <- err
				}
			}
		}
		for i := 0; i < 4; i++ {
			go tf(tin[i])
		}

		oi := 0
		err := qgen.NewAcc().Select("topics").Cols("tid,title,content,createdBy,ip").Each(func(rows *sql.Rows) error {
			t := ESTopic{}
			err := rows.Scan(&t.ID, &t.Title, &t.Content, &t.CreatedBy, &t.IP)
			if err != nil {
				return err
			}
			tin[oi] <- t
			if oi < 3 {
				oi++
			}
			return nil
		})
		for i := 0; i < 4; i++ {
			close(tin[i])
		}
		errs <- err
	}()

	go func() {
		rin := make([]chan ESReply, tcount)
		rf := func(rin chan ESReply) {
			for {
				reply, more := <-rin
				if !more {
					break
				}
				_, err := client.Index().Index("replies").Type("_doc").Id(strconv.Itoa(reply.ID)).BodyJson(reply).Do(context.Background())
				if err != nil {
					errs <- err
				}
			}
		}
		for i := 0; i < 4; i++ {
			rf(rin[i])
		}
		oi := 0
		err := qgen.NewAcc().Select("replies").Cols("rid,tid,content,createdBy,ip").Each(func(rows *sql.Rows) error {
			r := ESReply{}
			err := rows.Scan(&r.ID, &r.TID, &r.Content, &r.CreatedBy, &r.IP)
			if err != nil {
				return err
			}
			rin[oi] <- r
			if oi < 3 {
				oi++
			}
			return nil
		})
		for i := 0; i < 4; i++ {
			close(rin[i])
		}
		errs <- err
	}()

	fin := 0
	for {
		err := <-errs
		if err == nil {
			fin++
			if fin == 2 {
				return nil
			}
		} else {
			return err
		}
	}
}