e621-sdk-go/pkg/e621/endpoints/dbexport.go

package endpoints

import (
	"fmt"
	"git.dragse.it/anthrove/e621-sdk-go/pkg/e621/model"
	"git.dragse.it/anthrove/e621-sdk-go/pkg/e621/utils"
	"golang.org/x/net/html"
	"net/http"
	"strings"
)

// GetDBExportList retrieves a list of files available in the e621 database export.
//
// This function performs an HTTP GET request to the e621 database export endpoint and parses
// the HTML content to extract the links to export files with the ".csv.gz" extension.
//
// Parameters:
//   - requestContext: The context for the API request, including the host, user agent, username, and API key.
//
// Returns:
//   - []string: A slice of file names with the ".csv.gz" extension.
//   - error: An error, if any, encountered during the API request, response handling, or HTML parsing.
func GetDBExportList(requestContext model.RequestContext) ([]string, error) {
	// Create a new HTTP GET request.
	r, err := http.NewRequest("GET", fmt.Sprintf("https://e621.net/db_export/"), nil)
	if err != nil {
		return nil, err
	}

	r.Header.Set("User-Agent", requestContext.UserAgent)
	r.Header.Add("Accept", "application/json")
	r.SetBasicAuth(requestContext.Username, requestContext.APIKey)

	// Send the request using the provided http.Client.
	resp, err := requestContext.Client.Do(r)
	if err != nil {
		return nil, err
	}

	// Check if the HTTP response status code indicates success (2xx range).
	if resp.StatusCode < 200 || resp.StatusCode > 300 {
		// If the status code is outside the 2xx range, return an error based on the status code.
		return nil, utils.StatusCodesToError(resp.StatusCode)
	}

	// Parse the HTML content
	tokenizer := html.NewTokenizer(resp.Body)
	var files []string

	// Iterate through the HTML tokens
	for {
		tokenType := tokenizer.Next()

		switch tokenType {
		case html.ErrorToken:
			// End of the HTML document
			return files, nil
		case html.StartTagToken, html.SelfClosingTagToken:
			token := tokenizer.Token()
			if token.Data == "a" {
				// Check if the anchor tag has an href attribute
				for _, attr := range token.Attr {
					if attr.Key == "href" {
						// Filter out the parent directory link and only add links with ".csv.gz" extension
						if !strings.HasPrefix(attr.Val, "../") && strings.HasSuffix(attr.Val, ".csv.gz") {
							files = append(files, attr.Val)
						}
						break
					}
				}
			}
		}
	}
}

// GetDBExportFile retrieves a specific file from the e621 database export.
//
// This function performs an HTTP GET request to the e621 database export endpoint to fetch a
// particular file identified by its name.
//
// Parameters:
//   - requestContext: The context for the API request, including the host, user agent, username, and API key.
//   - file: The name of the file to be fetched from the database export.
//
// Returns:
//   - *http.Response: The HTTP response containing the requested file (probably a csv.gz).
//   - error: An error, if any, encountered during the API request or response handling.
func GetDBExportFile(requestContext model.RequestContext, file string) (*http.Response, error) {
	if file == "" {
		return nil, fmt.Errorf("no file specified")
	}

	// Create a new HTTP GET request.
	r, err := http.NewRequest("GET", fmt.Sprintf("https://e621.net/db_export/%s", file), nil)
	if err != nil {
		return nil, err
	}

	r.Header.Set("User-Agent", requestContext.UserAgent)
	r.Header.Add("Accept", "application/json")
	r.SetBasicAuth(requestContext.Username, requestContext.APIKey)

	// Send the request using the provided http.Client.
	resp, err := requestContext.Client.Do(r)
	if err != nil {
		return nil, err
	}

	// Check if the HTTP response status code indicates success (2xx range).
	if resp.StatusCode < 200 || resp.StatusCode > 300 {
		// If the status code is outside the 2xx range, return an error based on the status code.
		return nil, utils.StatusCodesToError(resp.StatusCode)
	}

	return resp, nil

}