zurg/pkg/repo/mysql.go

package repo

import (
	"bytes"
	"database/sql"
	"encoding/gob"
	"fmt"
	"log"
	"path"

	"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(directory, filename string) string {
	fullPath := path.Join(directory, filename)
	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, directory string, resp realdebrid.UnrestrictResponse) {
	// Generate the ID for the link
	var id string
	if resp.Filename == "" {
		id = GenerateID(directory, resp.Link)
	} else {
		id = GenerateID(directory, resp.Filename)
	}
	// 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,
			directory,
			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(directory, filename string) (*DavFile, error) {
	id := GenerateID(directory, filename)
	data, ok := db.Cache.Get([]byte(id))
	if !ok {
		resp, err := fetchFromDatabaseByID(db.Connection, id)
		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 (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 string) (*DavFile, error) {
	log.Printf("fetching from database: %s", id)
	var resp DavFile

	err := conn.QueryRow(`
		SELECT Filename, Filesize, Link
		FROM Links WHERE ID = ?`,
		id,
	).Scan(
		&resp.Filename,
		&resp.Filesize,
		&resp.Link,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return &resp, nil
		}
		log.Printf("failed to fetch record: %v", 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
}