Merge pull request #34 from gcp-kit/issue/33

refで投入したデータを保存しておくようにする #33

Author: moezakura

README.md

@@ -28,6 +28,18 @@ gcloud auth application-default login
 gcloud config set ProjectID
 してる時に該当のプロジェクトに対して接続される感じ
 
+### state_dir
+
+投入したデータのステータスを保存するディレクトリ。  
+空を指定した場合は、ステータスを保存しない。
+
+作成されるファイル
+
+```text
+state_dir/
+└── ref_ids.yaml   # ref_idのリスト
+```
+
 ## テストデータの作り方
 
 1. FirestoreのCollection名と同じディレクトリを作成する

cmd/fti/main.go

@@ -28,6 +28,10 @@ var (
 	versionFlag = flag.Bool("v", false, "show version")
 )
 
+const (
+	refIDsFileName = "ref_ids.yaml"
+)
+
 func main() {
 	flag.Parse()
 
@@ -60,11 +64,36 @@ func main() {
 		log.Fatalf("failed to init firestore: %+v", err)
 	}
 
-	is := inserter.NewInserter(client)
-	err = is.Execute(context.Background(), &cfg)
+	// refIDの一覧
+	refIDs := inserter.RefIDs{}
+	refFilePath := func() string {
+		if cfg.StateDir != "" {
+			return filepath.Join(cfg.StateDir, refIDsFileName)
+		}
+		return ""
+	}()
+
+	// stateDirが指定されている場合は、そこからrefIDを読み込む
+	if cfg.StateDir != "" {
+		err := refIDs.LoadFromFile(refFilePath)
+		if err != nil {
+			log.Fatalf("failed to load state file: %+v", err)
+		}
+	}
+
+	is := inserter.NewInserter(client, refIDs)
+	refIDs, err = is.Execute(context.Background(), &cfg)
 	if err != nil {
 		log.Fatalf("failed to execute insert: %+v", err)
 	}
+
+	// stateDirが指定されている場合は、そこにrefIDを書き込む
+	if cfg.StateDir != "" {
+		err := refIDs.SaveToFile(refFilePath)
+		if err != nil {
+			log.Fatalf("failed to save state file: %+v", err)
+		}
+	}
 }
 
 func initFirestore(ctx context.Context, cfg *config.Config) (*firestore.Client, error) {

pkg/config/config.go

@@ -6,4 +6,5 @@ type Config struct {
 	Targets                    []string `config:"targets,required" yaml:"targets"`
 	FirestoreProjectOnEmulator string   `config:"firestore_project_on_emulator" yaml:"firestore_project_on_emulator"`
 	FirestoreEmulatorHost      string   `config:"firestore_emulator_host" yaml:"firestore_emulator_host"`
+	StateDir                   string   `config:"state_dir" yaml:"state_dir"`
 }

pkg/inserter/common_inserter.go

@@ -15,17 +15,26 @@ import (
 // CommonInserter - Inserterの共通部分
 type CommonInserter struct {
 	client *firestore.Client
-	refIDs map[string]string
+	refIDs RefIDs
 }
 
 // NewCommonInserter - CommonInserter constructor
-func NewCommonInserter(client *firestore.Client) *CommonInserter {
+func NewCommonInserter(client *firestore.Client, refIDs RefIDs) *CommonInserter {
+	if refIDs == nil {
+		refIDs = RefIDs{}
+	}
+
 	return &CommonInserter{
 		client: client,
-		refIDs: map[string]string{},
+		refIDs: refIDs,
 	}
 }
 
+// RefIDs - ref id の一覧を返す
+func (c *CommonInserter) RefIDs() map[string]string {
+	return c.refIDs
+}
+
 // CreateItem - item を Firestore に作る
 func (c *CommonInserter) CreateItem(ctx context.Context, cn, ID, refID string, item map[string]interface{}) error {
 	item = c.tryParseDate(item)
@@ -93,7 +102,7 @@ func (c *CommonInserter) setRefs(item map[string]interface{}) map[string]interfa
 	for k, v := range item {
 		switch vt := v.(type) {
 		case []interface{}:
-			new := make([]string, len(vt))
+			ni := make([]string, len(vt))
 			isStr := true
 			for i, vtv := range vt {
 				vStr, ok := vtv.(string)
@@ -107,19 +116,19 @@ func (c *CommonInserter) setRefs(item map[string]interface{}) map[string]interfa
 					if !ok {
 						log.Printf("%s was not found", refID)
 					} else {
-						new[i] = rv
+						ni[i] = rv
 					}
 					continue
 				}
 				n := c.replaceMultiRefs(vStr, reg)
 				if n != "" {
-					new[i] = n
+					ni[i] = n
 					continue
 				}
-				new[i] = vStr
+				ni[i] = vStr
 			}
 			if isStr {
-				item[k] = new
+				item[k] = ni
 			}
 		case map[string]interface{}:
 			for vtk, vtv := range vt {

pkg/inserter/inserter.go

@@ -16,36 +16,41 @@ import (
 
 // Inserter - Inserter
 type Inserter struct {
+	commonInserter *CommonInserter
+
 	jsonInserter *JSONInserter
 	yamlInserter *YAMLInserter
 	jsInserter   *JSInserter
 }
 
 // NewInserter - Inserter constructor
-func NewInserter(client *firestore.Client) *Inserter {
-	ci := NewCommonInserter(client)
+func NewInserter(client *firestore.Client, refIDs RefIDs) *Inserter {
+	ci := NewCommonInserter(client, refIDs)
 	return &Inserter{
-		jsonInserter: NewJSONInserter(ci),
-		jsInserter:   NewJSInserter(ci),
-		yamlInserter: NewYAMLInserter(ci),
+		commonInserter: ci,
+		jsonInserter:   NewJSONInserter(ci),
+		jsInserter:     NewJSInserter(ci),
+		yamlInserter:   NewYAMLInserter(ci),
 	}
 }
 
 // Execute - ダミーデータソースの読み込みと実行処理呼び出しを行う
-func (i *Inserter) Execute(ctx context.Context, cfg *config.Config) error {
+func (i *Inserter) Execute(ctx context.Context, cfg *config.Config) (RefIDs, error) {
 	targetDir := cfg.Targets
 	for _, t := range targetDir {
 		if !files.Exists(t) {
-			return xerrors.Errorf("cannot find target directory: %s", targetDir)
+			return nil, xerrors.Errorf("cannot find target directory: %s", targetDir)
 		}
 
 		err := filepath.Walk(t, i.executeFile(ctx))
 		if err != nil {
-			return xerrors.Errorf("failed to execute file: %w", err)
+			return nil, xerrors.Errorf("failed to execute file: %w", err)
 		}
 	}
 
-	return nil
+	refIDs := i.commonInserter.RefIDs()
+
+	return refIDs, nil
 }
 
 func (i *Inserter) executeFile(ctx context.Context) func(path string, info os.FileInfo, _ error) error {

pkg/inserter/ref_ids.go

@@ -0,0 +1,50 @@
+package inserter
+
+import (
+	"os"
+
+	"github.com/goccy/go-yaml"
+	"golang.org/x/xerrors"
+)
+
+// RefIDs - ref id の一覧
+type RefIDs map[string]string
+
+// LoadFromFile - yamlファイルからrefIDを読み込む
+func (r RefIDs) LoadFromFile(path string) error {
+	// pathのyamlファイル読み込む
+	yb, err := os.ReadFile(path)
+	if err != nil {
+		return xerrors.Errorf("failed to read yaml file: %+v", err)
+	}
+
+	// yamlファイルを構造体に変換する
+	ym := new(RefIDs)
+	err = yaml.Unmarshal(yb, ym)
+	if err != nil {
+		return xerrors.Errorf("failed to unmarshal yaml: %w", err)
+	}
+	// 構造体の値をrにコピーする
+	for k, v := range *ym {
+		r[k] = v
+	}
+
+	return nil
+}
+
+// SaveToFile - yamlファイルにrefIDを書き込む
+func (r RefIDs) SaveToFile(path string) error {
+	// yamlに変換する
+	yb, err := yaml.Marshal(r)
+	if err != nil {
+		return xerrors.Errorf("failed to marshal yaml: %w", err)
+	}
+
+	// yamlファイルに書き込む
+	err = os.WriteFile(path, yb, 0666)
+	if err != nil {
+		return xerrors.Errorf("failed to write yaml file: %w", err)
+	}
+
+	return nil
+}