From fc539738f126ddc2d1b1465decb2a01283ce3871 Mon Sep 17 00:00:00 2001
From: Windfarer <windfarer@gmail.com>
Date: Mon, 19 Apr 2021 11:55:02 +0800
Subject: [PATCH] examples/blog: add cleanup (#845)

* add cleanup
---
 examples/blog/cmd/blog/main.go      |  6 ++++--
 examples/blog/cmd/blog/wire.go      |  2 +-
 examples/blog/cmd/blog/wire_gen.go  | 10 ++++++----
 examples/blog/internal/data/data.go | 18 ++++++++++++++----
 4 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/examples/blog/cmd/blog/main.go b/examples/blog/cmd/blog/main.go
index 51e90e156..f862d22f1 100644
--- a/examples/blog/cmd/blog/main.go
+++ b/examples/blog/cmd/blog/main.go
@@ -2,9 +2,10 @@ package main
 
 import (
 	"flag"
+	"os"
+
 	"go.opentelemetry.io/otel/exporters/trace/jaeger"
 	"go.opentelemetry.io/otel/sdk/trace"
-	"os"
 
 	"github.com/go-kratos/kratos/examples/blog/internal/conf"
 	"github.com/go-kratos/kratos/v2"
@@ -75,10 +76,11 @@ func main() {
 	}
 	defer flush()
 
-	app, err := initApp(bc.Server, bc.Data, tp, logger)
+	app, cleanup, err := initApp(bc.Server, bc.Data, tp, logger)
 	if err != nil {
 		panic(err)
 	}
+	defer cleanup()
 
 	// start and wait for stop signal
 	if err := app.Run(); err != nil {
diff --git a/examples/blog/cmd/blog/wire.go b/examples/blog/cmd/blog/wire.go
index e4c3f6464..de51b789c 100644
--- a/examples/blog/cmd/blog/wire.go
+++ b/examples/blog/cmd/blog/wire.go
@@ -17,6 +17,6 @@ import (
 )
 
 // initApp init kratos application.
-func initApp(*conf.Server, *conf.Data, trace.TracerProvider, log.Logger) (*kratos.App, error) {
+func initApp(*conf.Server, *conf.Data, trace.TracerProvider, log.Logger) (*kratos.App, func(), error) {
 	panic(wire.Build(server.ProviderSet, data.ProviderSet, biz.ProviderSet, service.ProviderSet, newApp))
 }
diff --git a/examples/blog/cmd/blog/wire_gen.go b/examples/blog/cmd/blog/wire_gen.go
index ca4f500aa..5156de64b 100644
--- a/examples/blog/cmd/blog/wire_gen.go
+++ b/examples/blog/cmd/blog/wire_gen.go
@@ -19,10 +19,10 @@ import (
 // Injectors from wire.go:
 
 // initApp init kratos application.
-func initApp(confServer *conf.Server, confData *conf.Data, tracerProvider trace.TracerProvider, logger log.Logger) (*kratos.App, error) {
-	dataData, err := data.NewData(confData, logger)
+func initApp(confServer *conf.Server, confData *conf.Data, tracerProvider trace.TracerProvider, logger log.Logger) (*kratos.App, func(), error) {
+	dataData, cleanup, err := data.NewData(confData, logger)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	articleRepo := data.NewArticleRepo(dataData, logger)
 	articleUsecase := biz.NewArticleUsecase(articleRepo, logger)
@@ -30,5 +30,7 @@ func initApp(confServer *conf.Server, confData *conf.Data, tracerProvider trace.
 	httpServer := server.NewHTTPServer(confServer, tracerProvider, blogService)
 	grpcServer := server.NewGRPCServer(confServer, tracerProvider, blogService)
 	app := newApp(logger, httpServer, grpcServer)
-	return app, nil
+	return app, func() {
+		cleanup()
+	}, nil
 }
diff --git a/examples/blog/internal/data/data.go b/examples/blog/internal/data/data.go
index 9acdd58b6..725cc26f8 100644
--- a/examples/blog/internal/data/data.go
+++ b/examples/blog/internal/data/data.go
@@ -2,6 +2,7 @@ package data
 
 import (
 	"context"
+
 	"github.com/go-kratos/kratos/examples/blog/internal/conf"
 	"github.com/go-kratos/kratos/examples/blog/internal/data/ent"
 	"github.com/go-kratos/kratos/v2/log"
@@ -23,7 +24,7 @@ type Data struct {
 }
 
 // NewData .
-func NewData(conf *conf.Data, logger log.Logger) (*Data, error) {
+func NewData(conf *conf.Data, logger log.Logger) (*Data, func(), error) {
 	log := log.NewHelper("data", logger)
 	client, err := ent.Open(
 		conf.Database.Driver,
@@ -31,12 +32,12 @@ func NewData(conf *conf.Data, logger log.Logger) (*Data, error) {
 	)
 	if err != nil {
 		log.Errorf("failed opening connection to sqlite: %v", err)
-		return nil, err
+		return nil, nil, err
 	}
 	// Run the auto migration tool.
 	if err := client.Schema.Create(context.Background()); err != nil {
 		log.Errorf("failed creating schema resources: %v", err)
-		return nil, err
+		return nil, nil, err
 	}
 
 	rdb := redis.NewClient(&redis.Options{
@@ -48,8 +49,17 @@ func NewData(conf *conf.Data, logger log.Logger) (*Data, error) {
 		ReadTimeout:  conf.Redis.ReadTimeout.AsDuration(),
 	})
 	rdb.AddHook(redisotel.TracingHook{})
-	return &Data{
+	d := &Data{
 		db:  client,
 		rdb: rdb,
+	}
+	return d, func() {
+		log.Info("message", "closing the data resources")
+		if err := d.db.Close(); err != nil {
+			log.Error(err)
+		}
+		if err := d.rdb.Close(); err != nil {
+			log.Error(err)
+		}
 	}, nil
 }