timetable-mgr/pis/update.go

package pis

import (
	"archive/tar"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"git.fjla.uk/owlboard/timetable-mgr/dbAccess"
	"git.fjla.uk/owlboard/timetable-mgr/log"
	"go.uber.org/zap"
)

const (
	destPath    = "/tmp/pis.tar.gz"
	extractPath = "/tmp/extract"
)

// Downloads the release tarball, extracts then applies to database
func runUpdate(tarballUrl string) error {
	log.Info("PIS Update Process started")
	err := downloadFile(tarballUrl, destPath)
	if err != nil {
		return err
	}

	// Extract to disk
	file, err := os.Open(destPath)
	if err != nil {
		return err
	}
	log.Info("Saved PIS Release to disk")
	defer file.Close()

	if err := extractFiles(file, extractPath); err != nil {
		return err
	}
	log.Info("Extracted PIS Release")

	// Load YAML to string
	pisData, err := extractYamlData(extractPath)
	if err != nil {
		return err
	}

	pisSlice, err := processYaml(pisData)
	if err != nil {
		return err
	}
	log.Info("Loaded PIS Files to Slice")

	// Temporarily use "pis_testing" collection to ensure all works as expected
	err = dbAccess.DropCollection("pis_testing")
	if err != nil {
		return err
	}

	count, err := dbAccess.PutPisData(pisSlice)
	if err != nil {
		return err
	}

	log.Info("Insterted new PIS Data", zap.Int64("PIS Codes", count))

	err = dbAccess.CreatePisIndeces()
	if err != nil {
		log.Error("Failed to create PIS Indeces, poor performance expected", zap.Error(err))
	}

	// Cleanup files
	cleanupFiles(destPath, extractPath)
	return nil
}

// Download the tarball to disk
func downloadFile(url, filepath string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	log.Info("PIS Release downloaded")

	defer resp.Body.Close()

	out, err := os.Create(filepath)
	if err != nil {
		return err
	}

	defer out.Close()

	_, err = io.Copy(out, resp.Body)
	return err
}

// Extract tarball to disk
func extractFiles(gzipStream io.Reader, dest string) error {
	uncompressedStream, err := gzip.NewReader(gzipStream)
	if err != nil {
		return err
	}
	defer uncompressedStream.Close()
	log.Info("Extracting PIS File")

	tarReader := tar.NewReader(uncompressedStream)

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		// Handle pax_global_header or other unsupported types
		if header.Typeflag == tar.TypeXGlobalHeader || header.Name == "pax_global_header" {
			// Skip this special header file
			continue
		}

		filePath := filepath.Join(dest, header.Name)

		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(filepath.Join(dest, header.Name), 0755); err != nil {
				return err
			}
		case tar.TypeReg:
			outFile, err := os.Create(filePath)
			if err != nil {
				return err
			}
			if _, err := io.Copy(outFile, tarReader); err != nil {
				return err
			}
			outFile.Close()
		default:
			log.Warn("Unable to handle filetype", zap.String("Typeflag", string(header.Typeflag)), zap.String("Filename", header.Name))
		}
	}
	return nil
}

// Return YAML PIS files as a string
func extractYamlData(dir string) (string, error) {
	var allContent strings.Builder // Using a string builder to accumulate content

	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err // Only returning error since Walk callback doesn't accept string
		}

		// Check if the path contains 'pis' and has a .yaml or .yml extension
		if strings.Contains(path, "/pis/") && !info.IsDir() &&
			(strings.HasSuffix(info.Name(), ".yaml") || strings.HasSuffix(info.Name(), ".yml")) {
			log.Debug("Processing YAML", zap.String("directory", path))

			file, err := os.Open(path)
			if err != nil {
				return fmt.Errorf("failed to open YAML file %s: %v", path, err)
			}
			defer file.Close()

			content, err := io.ReadAll(file)
			if err != nil {
				return fmt.Errorf("failed to read YAML file %s: %v", path, err)
			}

			// Accumulate content from each YAML file in 'pis' directory
			allContent.Write(content)
			allContent.WriteString("\n") // Add a newline between file contents
		}

		return nil
	})

	if err != nil {
		return "", err
	}

	// Return the accumulated content as a single string
	return allContent.String(), nil
}

// Cleans up downloaded and extracted files
func cleanupFiles(paths ...string) {
	for _, path := range paths {
		err := os.RemoveAll(path)
		if err != nil {
			log.Warn("Error removing file", zap.String("path", path), zap.Error(err))
		}
	}
	log.Info("Removed PIS Release files")
}

// Saves pisSlice to file
func DEBUG_saveSliceToFile(pisSlice interface{}, filename string) error {
	data, err := json.MarshalIndent(pisSlice, "", "  ")
	if err != nil {
		return fmt.Errorf("error marshalling slice to JSON: %v", err)
	}

	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("error creating file: %v", err)
	}
	defer file.Close()

	_, err = file.Write(data)
	if err != nil {
		return fmt.Errorf("error writing JSON to file: %v", err)
	}

	return nil
}