Uploader/main.go

package main

import (
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/jaswdr/faker"
)

var dbpool *pgxpool.Pool
var uploadFolder string

func initUpload(c *gin.Context) {
	fileName, err := io.ReadAll(c.Request.Body)
	if err != nil {
		c.JSON(http.StatusBadRequest, "Couldn't read html request body")
		return
	}

	rows, _ := dbpool.Query(context.Background(), "select filepath from videos")
	for rows.Next() {
		var filepath string
		err = rows.Scan(&filepath)
		if err != nil {
			c.JSON(http.StatusInternalServerError, "")
			log.Panicf("initUpload: %v\n", err)
		}
		if filepath == fmt.Sprintf("%s/%s", uploadFolder, fileName) {
			c.JSON(http.StatusForbidden, "File already exists")
			return
		}
	}
	err = rows.Err()
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("initUpload: %v\n", err)
	}

	sqlStmt := `
	DROP TABLE IF EXISTS chunks;

	CREATE TABLE IF NOT EXISTS chunks
	(
		id integer,
		fileName text,
		chunk bytea
	)
	`
	_, err = dbpool.Exec(context.Background(), sqlStmt)
	if err != nil {
		log.Panicf("initUpload: %v\n", err)
	}

	c.JSON(http.StatusOK, "Upload can proceed")
}

func receiveChunk(c *gin.Context) {
	chunk, err := io.ReadAll(c.Request.Body)
	if err != nil {
		c.JSON(http.StatusBadRequest, "Couldn't read html request body")
		return
	}

	_, err = dbpool.Exec(context.Background(), "insert into chunks(id, fileName, chunk) values($1, $2, $3)", c.GetHeader("chunk-id"), c.GetHeader("file-name"), chunk)
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("receiveChunk: %v\n", err)
	}

	c.JSON(http.StatusOK, "Received chunk")
}

func finishUpload(c *gin.Context) {
	allChunks := map[int][]byte{}

	fileName, err := io.ReadAll(c.Request.Body)
	if err != nil {
		c.JSON(http.StatusBadRequest, "Couldn't read html request body")
		return
	}

	time.Sleep(time.Second)

	rows, _ := dbpool.Query(context.Background(), "select id, chunk from chunks where fileName like $1", fileName)
	for rows.Next() {
		var id int
		var chunk []byte
		err := rows.Scan(&id, &chunk)
		if err != nil {
			c.JSON(http.StatusInternalServerError, "")
			log.Panicf("finishUpload: %v\n", err)
		}

		allChunks[id] = chunk
	}
	err = rows.Err()
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("finishUpload: %v\n", err)
	}

	for i := 0; i < len(allChunks); i++ {
		f, err := os.OpenFile(fmt.Sprintf("%s/%s", uploadFolder, fileName), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0640)
		if err != nil {
			c.JSON(http.StatusInternalServerError, "")
			log.Panicf("finishUpload: %v\n", err)
		}
		if _, err := f.Write(allChunks[i]); err != nil {
			c.JSON(http.StatusInternalServerError, "")
			log.Panicf("finishUpload: %v\n", err)
		}
		if err := f.Close(); err != nil {
			c.JSON(http.StatusInternalServerError, "")
			log.Panicf("finishUpload: %v\n", err)
		}
	}

	_, err = dbpool.Exec(context.Background(), "insert into videos(filepath) values($1)", fmt.Sprintf("%s/%s", uploadFolder, fileName))
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("finishUpload: %v\n", err)
	}

	c.JSON(http.StatusOK, "File uploaded successfully")
}

func listVideos(c *gin.Context) {
	allVideos := map[int]string{}

	rows, _ := dbpool.Query(context.Background(), "select * from videos")
	for rows.Next() {
		var id int
		var filepath string
		err := rows.Scan(&id, &filepath)
		if err != nil {
			c.JSON(http.StatusInternalServerError, "")
			log.Panicf("listVideos: %v\n", err)
		}
		allVideos[id] = filepath
	}
	err := rows.Err()
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("listVideos: %v\n", err)
	}

	c.JSON(http.StatusOK, allVideos)
}

func getVideo(c *gin.Context) {
	rows, _ := dbpool.Query(context.Background(), "select filepath from videos where id = $1", c.Param("id"))
	rows.Next()
	err := rows.Err()
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("getVideo: %v\n", err)
	}

	var filepath string
	err = rows.Scan(&filepath)
	if err != nil {
		c.JSON(http.StatusBadRequest, "Video does not exist")
		log.Panicf("getVideo: %v\n", err)
	}

	c.JSON(http.StatusOK, filepath)
}

func deleteVideo(c *gin.Context) {
	rows, _ := dbpool.Query(context.Background(), "select filepath from videos where id = $1", c.Param("id"))
	rows.Next()
	err := rows.Err()
	if err != nil {
		c.JSON(http.StatusInternalServerError, "")
		log.Panicf("deleteVideo: %v\n", err)
	}

	var filepath string
	err = rows.Scan(&filepath)
	if err != nil {
		c.JSON(http.StatusBadRequest, "Video does not exist")
		log.Panicf("deleteVideo: %v\n", err)
	}

	_, err = dbpool.Exec(context.Background(), "delete from videos where id = $1", c.Param("id"))
	if err != nil {
		c.JSON(http.StatusInternalServerError, "Id was likely invalid")
		log.Panicf("deleteVideo: %v\n", err)
	}

	if err = os.Remove(filepath); err != nil {
		c.JSON(http.StatusInternalServerError, "DB entry was deleted, but file likely doesn't exist")
		log.Panicf("deleteVideo: %v\n", err)
	}

	c.JSON(http.StatusOK, "File deleted successfully")
}

func main() {
	var err error
	dbpool, err = pgxpool.New(context.Background(), "postgresql://postgres:postgres@172.18.0.3:5432/postgres")
	if err != nil {
		log.Fatal(err)
	}
	defer dbpool.Close()

	sqlStmt := `
	DROP TABLE IF EXISTS videos;

	CREATE TABLE IF NOT EXISTS videos
	(
		id serial NOT NULL,
		filepath text,
		CONSTRAINT videos_pkey PRIMARY KEY (id)
	)
	`
	_, err = dbpool.Exec(context.Background(), sqlStmt)
	if err != nil {
		log.Panicf("main: %v\n", err)
	}

	faker := faker.New()

	for i := 0; i < 10; i++ {
		_, err = dbpool.Exec(context.Background(), "insert into videos(filepath) values($1)", faker.File().AbsoluteFilePathForUnix(2))
		if err != nil {
			log.Panicf("main: %v\n", err)
		}
	}

	currentDir, err := os.Getwd()
	if err != nil {
		log.Panicf("main: %v\n", err)
	}
	uploadFolder = fmt.Sprintf("%s/upload", currentDir)

	if _, err := os.Stat(uploadFolder); err != nil {
		if !os.IsNotExist(err) {
			log.Panicf("main: %v\n", err)
		}

		if err := os.Mkdir(uploadFolder, 0750); err != nil {
			log.Panicf("main: %v\n", err)
		}
	}

	files, err := os.ReadDir(uploadFolder)
	if err != nil {
		panic(err)
	}

	if len(files) != 0 {
		for _, file := range files {
			if file.IsDir() {
				continue
			}

			_, err = dbpool.Exec(context.Background(), "insert into videos(filepath) values($1)", fmt.Sprintf("%s/%s", uploadFolder, file.Name()))
			if err != nil {
				log.Panicf("main: %v\n", err)
			}
		}
	}

	router := gin.Default()
	router.SetTrustedProxies(nil)

	router.POST("/video/init", initUpload)
	router.POST("/video/chunk", receiveChunk)
	router.POST("/video/completed", finishUpload)
	router.GET("/videos", listVideos)
	router.GET("/videos/:id", getVideo)
	router.DELETE("/videos/:id", deleteVideo)

	router.Run("localhost:8080")
}