Create & Init Project...

app/job/main/tv/dao/ftp/BUILD

@@ -0,0 +1,51 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_library",
+    "go_test",
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "dao.go",
+        "ftp.go",
+    ],
+    importpath = "go-common/app/job/main/tv/dao/ftp",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//app/job/main/tv/conf:go_default_library",
+        "//library/log:go_default_library",
+        "//vendor/github.com/ftp-master:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = [
+        "dao_test.go",
+        "ftp_test.go",
+    ],
+    embed = [":go_default_library"],
+    tags = ["automanaged"],
+    deps = [
+        "//app/job/main/tv/conf:go_default_library",
+        "//vendor/github.com/smartystreets/goconvey/convey:go_default_library",
+    ],
+)

app/job/main/tv/dao/ftp/dao.go

@@ -0,0 +1,18 @@
+package ftp
+
+import (
+	"go-common/app/job/main/tv/conf"
+)
+
+// Dao dao.
+type Dao struct {
+	conf *conf.Config
+}
+
+// New create a instance of Dao and return.
+func New(c *conf.Config) (d *Dao) {
+	d = &Dao{
+		conf: c,
+	}
+	return
+}

app/job/main/tv/dao/ftp/dao_test.go

@@ -0,0 +1,35 @@
+package ftp
+
+import (
+	"flag"
+	"os"
+	"testing"
+
+	"go-common/app/job/main/tv/conf"
+)
+
+var (
+	d *Dao
+)
+
+func TestMain(m *testing.M) {
+	if os.Getenv("DEPLOY_ENV") != "" {
+		flag.Set("app_id", "main.web-svr.tv-job")
+		flag.Set("conf_token", "ab3e9801a77c076b997de0ac5cb21775")
+		flag.Set("tree_id", "15260")
+		flag.Set("conf_version", "docker-1")
+		flag.Set("deploy_env", "uat")
+		flag.Set("conf_host", "config.bilibili.co")
+		flag.Set("conf_path", "/tmp")
+		flag.Set("region", "sh")
+		flag.Set("zone", "sh001")
+	} else {
+		flag.Set("conf", "../../cmd/tv-job-test.toml")
+	}
+	flag.Parse()
+	if err := conf.Init(); err != nil {
+		panic(err)
+	}
+	d = New(conf.Conf)
+	os.Exit(m.Run())
+}

app/job/main/tv/dao/ftp/ftp.go

@@ -0,0 +1,112 @@
+package ftp
+
+import (
+	"bytes"
+	"crypto/md5"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"time"
+
+	"go-common/library/log"
+
+	"github.com/ftp-master"
+)
+
+const (
+	_ftpRetry = 3
+	errFormat = "Func:[%s] - Step:[%s] - Error:[%v]"
+	_sleep    = 100 * time.Millisecond
+)
+
+// Retry . retry one function until no error
+func Retry(callback func() error, retry int, sleep time.Duration) (err error) {
+	for i := 0; i < retry; i++ {
+		if err = callback(); err == nil {
+			return
+		}
+		time.Sleep(sleep)
+	}
+	return
+}
+
+// FileMd5 calculates the local file's md5 and store it in a file
+func (d *Dao) FileMd5(path string, md5Path string) (err error) {
+	var (
+		content []byte
+	)
+	if content, err = ioutil.ReadFile(path); err != nil {
+		log.Error(errFormat+" FilePath: %s", "fileMd5", "ReadFile", err, path)
+		return
+	}
+	md5hash := md5.New()
+	if _, err = io.Copy(md5hash, bytes.NewReader(content)); err != nil {
+		log.Error(errFormat, "fileMd5", "CopyContent", err)
+		return
+	}
+	md5 := md5hash.Sum(nil)
+	fMd5 := hex.EncodeToString(md5[:])
+	file, error := os.OpenFile(md5Path, os.O_RDWR|os.O_CREATE, 0766)
+	if error != nil {
+		log.Error(errFormat, "fileMd5", "OpenFile", err)
+		return
+	}
+	file.WriteString(fMd5)
+	file.Close()
+	return
+}
+
+// UploadFile the file to remote frp server and update the md5 file
+func (d *Dao) UploadFile(localPath string, remotePath string, url string) (err error) {
+	var (
+		ftpInfo  = d.conf.Search.FTP
+		c        *ftp.ServerConn
+		content  []byte // file's content
+		fileSize int64
+	)
+	// Dial
+	if c, err = ftp.DialTimeout(ftpInfo.Host, time.Duration(ftpInfo.Timeout)); err != nil {
+		log.Error(errFormat, "uploadFile", "DialTimeout", err)
+		return
+	}
+	// use EPSV or not
+	if !ftpInfo.UseEPSV {
+		c.DisableEPSV = true
+	}
+	// Login
+	if err = c.Login(ftpInfo.User, ftpInfo.Pass); err != nil {
+		log.Error(errFormat, "uploadFile", "Login", err)
+		return
+	}
+	// Change dir
+	if err = c.ChangeDir(url); err != nil {
+		log.Error(errFormat, "uploadFile", "ChangeDir", err)
+		return
+	}
+	// Upload the file
+	if content, err = ioutil.ReadFile(localPath); err != nil {
+		log.Error(errFormat, "uploadFile", "ReadFile", err)
+		return
+	}
+	data := bytes.NewBuffer(content)
+	if err = Retry(func() (err error) {
+		return c.Stor(remotePath, data)
+	}, _ftpRetry, _sleep); err != nil {
+		log.Error("upArchives Error %+v", err)
+		return
+	}
+	// Calculate the file size to check it's ok
+	if fileSize, err = c.FileSize(remotePath); err != nil {
+		log.Error(errFormat, "uploadFile", "FileSize", err)
+		return
+	}
+	if localSize := int64(len(content)); localSize != fileSize {
+		err = fmt.Errorf("LocalSize is %d, RemoteSize is %d", localSize, fileSize)
+		log.Error(errFormat, "uploadFile", "FileSize", err)
+		return
+	}
+	log.Info("File %s is uploaded successfully, size: %d", remotePath, fileSize)
+	return
+}

app/job/main/tv/dao/ftp/ftp_test.go

@@ -0,0 +1,78 @@
+package ftp
+
+import (
+	"fmt"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/smartystreets/goconvey/convey"
+)
+
+func fileExist(path string) bool {
+	_, err := os.Stat(path)
+	if err != nil {
+		if os.IsNotExist(err) {
+
+			return false
+		}
+	}
+	return true
+}
+
+func createFile(path string) {
+	if !fileExist(path) {
+		// If the file doesn't exist, create it, or append to the file
+		f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+		if err != nil {
+			fmt.Println(err)
+		}
+		_, err = f.Write([]byte("Hello"))
+		if err != nil {
+			fmt.Println(err)
+		}
+		f.Close()
+	}
+}
+
+func TestFtpRetry(t *testing.T) {
+	var (
+		callback func() error
+		retry    = int(0)
+		sleep    time.Duration
+	)
+	convey.Convey("Retry", t, func(ctx convey.C) {
+		err := Retry(callback, retry, sleep)
+		ctx.Convey("Then err should be nil.", func(ctx convey.C) {
+			ctx.So(err, convey.ShouldBeNil)
+		})
+	})
+}
+
+func TestFtpFileMd5(t *testing.T) {
+	var (
+		path    = "/tmp/testMd5.source"
+		md5Path = "/tmp/testMd5.target"
+	)
+	convey.Convey("FileMd5", t, func(ctx convey.C) {
+		createFile(path)
+		err := d.FileMd5(path, md5Path)
+		ctx.Convey("Then err should be nil.", func(ctx convey.C) {
+			ctx.So(err, convey.ShouldBeNil)
+		})
+	})
+}
+
+func TestFtpUploadFile(t *testing.T) {
+	var (
+		cfg  = d.conf.Search
+		path = "/tmp/testMd5.source"
+	)
+	convey.Convey("UploadFile", t, func(ctx convey.C) {
+		createFile(path)
+		err := d.UploadFile(path, "testMd5.remote", cfg.FTP.RemotePgcURL)
+		ctx.Convey("Then err should be nil.", func(ctx convey.C) {
+			ctx.So(err, convey.ShouldBeNil)
+		})
+	})
+}