zurg/pkg/repo/mysql.go

package repo

import (
	"bytes"
	"database/sql"
	"encoding/gob"
	"fmt"
	"log"
	"net/url"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/debridmediamanager.com/zurg/pkg/davextra"
	"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
	_ "github.com/go-sql-driver/mysql"
	"github.com/qianbin/directcache"
	"github.com/zeebo/xxh3"
)

type Database struct {
	Connection *sql.DB
	Cache      *directcache.Cache
}

func GenerateID(segment1, segment2, segment3 string) string {
	fullPath := path.Join(segment1, segment2, segment3)
	hash := xxh3.HashString(fullPath)
	return fmt.Sprintf("%016x", hash)
}

func NewDatabase(dsn string) (*Database, error) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}

	cache := directcache.New(10 << 20) // This initializes a cache with 10 MB

	return &Database{Connection: db, Cache: cache}, nil
}

func (db *Database) Insert(parentHash, torrentName string, resp realdebrid.UnrestrictResponse) {
	// Generate the ID for the link
	var id string
	if resp.Filename == "" {
		// alternative ID for 404 links
		id = GenerateID(parentHash, resp.Link, "")
	} else {
		id = GenerateID(parentHash, resp.Filename, davextra.GetLinkFragment(resp.Link))
	}
	// Check if the link already exists in the database
	var exists int
	err := db.Connection.QueryRow("SELECT COUNT(*) FROM Links WHERE ID = ?", id).Scan(&exists)
	if err != nil {
		log.Printf("failed to check existence: %v", err)
	}

	// If link does not exist in the database, insert the new record
	if exists == 0 {
		_, err = db.Connection.Exec(`
			INSERT INTO Links (ID, ParentHash, Directory, Filename, Filesize, Link, Host)
			VALUES (?, ?, ?, ?, ?, ?, ?)`,
			id,
			parentHash,
			torrentName,
			resp.Filename,
			resp.Filesize,
			resp.Link,
			resp.Host,
		)
		if err != nil {
			log.Printf("failed to insert record: %v", err)
		}

		// Clear cache for parentHash
		db.Cache.Del([]byte(parentHash))
	}
}

func (db *Database) Get(parentHash, filename string) (*DavFile, error) {
	filenameV2, linkFragment := extractIDFromFilename(filename)
	id := GenerateID(parentHash, filenameV2, linkFragment)
	data, ok := db.Cache.Get([]byte(id))
	if !ok {
		resp, err := fetchFromDatabaseByID(db.Connection, id, linkFragment)
		if err != nil {
			return nil, err
		}

		buffer := &bytes.Buffer{}
		encoder := gob.NewEncoder(buffer)
		if err := encoder.Encode(resp); err != nil {
			return nil, err
		}

		db.Cache.Set([]byte(id), buffer.Bytes())
		return resp, nil
	}

	buffer := bytes.NewBuffer(data)
	decoder := gob.NewDecoder(buffer)
	var resp DavFile
	if err := decoder.Decode(&resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func extractIDFromFilename(filename string) (string, string) {
	filenameV2, err := url.PathUnescape(filename)
	if err != nil {
		filenameV2 = filename
	}
	ext := filepath.Ext(filenameV2)
	name := strings.TrimSuffix(filenameV2, ext)

	r := regexp.MustCompile(`\sDMM(\w+)`)
	matches := r.FindStringSubmatch(name)
	if len(matches) < 2 {
		// No ID found
		return filenameV2, ""
	}

	// Remove ID from filename
	originalName := strings.Replace(name, matches[0], "", 1)
	return originalName + ext, matches[1]
}

func (db *Database) GetMultiple(parentHash string) (*DavFiles, error) {
	key := []byte(parentHash)
	data, ok := db.Cache.Get(key)
	if !ok {
		resps, err := fetchMultipleFromDatabase(db.Connection, parentHash)
		if err != nil {
			return nil, err
		}

		buffer := &bytes.Buffer{}
		encoder := gob.NewEncoder(buffer)
		if err := encoder.Encode(resps); err != nil {
			return nil, err
		}

		db.Cache.Set(key, buffer.Bytes())
		return resps, nil
	}

	buffer := bytes.NewBuffer(data)
	decoder := gob.NewDecoder(buffer)
	var resps DavFiles
	if err := decoder.Decode(&resps); err != nil {
		return nil, err
	}

	return &resps, nil
}

func fetchFromDatabaseByID(conn *sql.DB, id, linkFragment string) (*DavFile, error) {
	log.Printf("fetching from database: %s", id)
	var resp DavFile

	query := `
			SELECT Filename, Filesize, Link
			FROM Links WHERE ID = ? AND Link LIKE ?`
	row := conn.QueryRow(query, id, "https://real-debrid.com/d/"+linkFragment+"%")

	err := row.Scan(&resp.Filename, &resp.Filesize, &resp.Link)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		log.Printf("failed to fetch record: %v", err)
		return nil, err
	}

	return &resp, nil
}

func fetchMultipleFromDatabase(conn *sql.DB, parentHash string) (*DavFiles, error) {
	log.Printf("fetching multiple from database: %s", parentHash)
	rows, err := conn.Query(`
		SELECT Filename, Filesize, Link
		FROM Links WHERE ParentHash = ?`,
		parentHash,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch records: %v", err)
	}
	defer rows.Close()

	var responses []*DavFile

	for rows.Next() {
		resp := &DavFile{}
		if err := rows.Scan(&resp.Filename, &resp.Filesize, &resp.Link); err != nil {
			return nil, fmt.Errorf("failed to scan row: %v", err)
		}
		responses = append(responses, resp)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error while iterating over rows: %v", err)
	}

	result := &DavFiles{
		Files: responses,
	}

	return result, nil
}