Go语言mgo：高效流式上传文件至MongoDB GridFS实践指南

package main

import (
    "fmt"
    "io/ioutil" // 废弃，但此处用于演示旧代码模式
    "log"
    "net/http"
    "time"

    "gopkg.in/mgo.v2"
    "gopkg.in/mgo.v2/bson"
)

// mongoSession 假设是已初始化并连接的mgo会话
var mongoSession *mgo.Session

func init() {
    // 实际应用中应更健壮地处理连接和错误
    var err error
    mongoSession, err = mgo.Dial("mongodb://localhost:27017/testdb") // 请替换为你的MongoDB连接字符串和数据库名
    if err != nil {
        log.Fatalf("Failed to connect to MongoDB: %v", err)
    }
    mongoSession.SetMode(mgo.Monotonic, true)
    log.Println("MongoDB connected successfully.")
}

func uploadFilePageHandler(w http.ResponseWriter, req *http.Request) {
    if req.Method != http.MethodPost {
        http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
        return
    }

    file, handler, err := req.FormFile("filename") // "filename" 是HTML表单中input type="file"的name属性
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to get file from form: %v", err), http.StatusBadRequest)
        return
    }
    defer file.Close() // 确保上传的文件句柄被关闭

    // 核心问题所在：将整个文件读入内存
    data, err := ioutil.ReadAll(file) // 对于大文件，这将消耗大量内存
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to read file into memory: %v", err), http.StatusInternalServerError)
        return
    }

    session := mongoSession.Copy()
    defer session.Close()
    db := session.DB("testdb") // 替换为你的数据库名

    // 为GridFS文件生成唯一文件名
    uniqueFilename := fmt.Sprintf("%d-%s", time.Now().UnixNano(), handler.Filename)

    gridFile, err := db.GridFS("fs").Create(uniqueFilename)
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to create GridFS file: %v", err), http.StatusInternalServerError)
        return
    }
    defer gridFile.Close() // 确保GridFS文件写入器被关闭

    // 从内存中的数据写入GridFS
    bytesWritten, err := gridFile.Write(data)
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to write file to GridFS: %v", err), http.StatusInternalServerError)
        return
    }

    fmt.Fprintf(w, "File uploaded successfully (inefficiently)! Bytes written: %d, GridFS ID: %s\n", bytesWritten, gridFile.Id().(bson.ObjectId).Hex())
    log.Printf("File '%s' uploaded to GridFS as '%s', %d bytes written (inefficiently).", handler.Filename, uniqueFilename, bytesWritten)
}

// func main() {
//  http.HandleFunc("/upload_inefficient", uploadFilePageHandler)
//  log.Println("Inefficient server started on :8081, upload endpoint: /upload_inefficient")
//  log.Fatal(http.ListenAndServe(":8081", nil))
// }

package main

import (
    "fmt"
    "io"
    "log"
    "net/http"
    "path/filepath"
    "time"

    "gopkg.in/mgo.v2"
    "gopkg.in/mgo.v2/bson"
)

// mongoSession 假设是已初始化并连接的mgo会话
var mongoSession *mgo.Session

func init() {
    // 实际应用中应更健壮地处理连接和错误
    var err error
    // 请替换为你的MongoDB连接字符串和数据库名
    mongoSession, err = mgo.Dial("mongodb://localhost:27017/testdb")
    if err != nil {
        log.Fatalf("Failed to connect to MongoDB: %v", err)
    }
    mongoSession.SetMode(mgo.Monotonic, true)
    log.Println("MongoDB connected successfully.")
}

// uploadFileHandler 处理文件上传请求，采用流式传输
func uploadFileHandler(w http.ResponseWriter, req *http.Request) {
    if req.Method != http.MethodPost {
        http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
        return
    }

    // 1. 从multipart表单中获取文件
    // "filename" 是HTML表单中input type="file"的name属性
    file, handler, err := req.FormFile("filename")
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to get file from form: %v", err), http.StatusBadRequest)
        return
    }
    defer file.Close() // 确保上传的文件句柄被关闭，释放系统资源

    // 获取MongoDB会话和数据库
    session := mongoSession.Copy() // 为每个请求复制会话，确保并发安全
    defer session.Close()
    db := session.DB("testdb") // 替换为你的数据库名

    // 2. 为GridFS文件生成唯一的文件名
    // 建议生成唯一文件名以避免冲突，并保留原始文件名作为元数据
    uniqueFilename := fmt.Sprintf("%d-%s", time.Now().UnixNano(), filepath.Base(handler.Filename))

    // 3. 在GridFS中创建文件写入器
    gridFile, err := db.GridFS("fs").Create(uniqueFilename) // "fs" 是GridFS的默认集合前缀
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to create GridFS file: %v", err), http.StatusInternalServerError)
        return
    }
    defer gridFile.Close() // 确保GridFS文件写入器被关闭，完成文件写入操作

    // 可选：设置GridFS文件的元数据
    // 这些元数据会存储在fs.files集合中，便于后续查询
    gridFile.SetMeta(bson.M{
        "original_filename": handler.Filename,
        "content_type":      handler.Header.Get("Content-Type"),
        "upload_date":       time.Now(),
    })

    // 4. 使用io.Copy直接从上传文件流写入到GridFS文件
    // 这是关键步骤，避免了将整个文件读入内存
    bytesWritten, err := io.Copy(gridFile, file)
    if err != nil {
        http.Error(w, fmt.Sprintf("Failed to write file to GridFS: %v", err), http.StatusInternalServerError)
        return
    }

    // 5. 响应客户端
    fmt.Fprintf(w, "File uploaded successfully! Bytes written: %d, GridFS ID: %s\n", bytesWritten, gridFile.Id().(bson.ObjectId).Hex())
    log.Printf("File '%s' uploaded to GridFS as '%s', %d bytes written.", handler.Filename, uniqueFilename, bytesWritten)
}

func main() {
    http.HandleFunc("/upload", uploadFileHandler)
    log.Println("Server started on :8080, upload endpoint: /upload")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

/*
一个简单的HTML表单用于测试上传：

<!DOCTYPE html>
<html>
<head>
    <title>Upload File</title>
</head>
<body>
    <h1>Upload File to GridFS</h1>
    <form action="/upload" method="post" enctype="multipart/form-data">
        <label for="file">Choose a file:</label>
        <input type="file" id="file" name="filename" required>
        <br><br>
        <input type="submit" value="Upload File">
    </form>
</body>
</html>
*/