timetable-mgr/pis/update.go

package pis

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

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

// Downloads the release tarball, extracts then applies to database
func runUpdate(tarballUrl string) error {
	err := downloadFile(tarballUrl, destPath)
	if err != nil {
		return err
	}

	// Extract to disk
	file, err := os.Open(destPath)
	if err != nil {
		return err
	}
	defer file.Close()

	if err := extractFiles(file, extractPath); err != nil {
		return err
	}

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

	pisSlice, err := processYaml(pisData)
	if err != nil {
		return err
	}

	fmt.Println(&pisSlice)
	/*
		err = dbAccess.DropCollection(dbAccess.PisCollection)
		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
	}

	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()

	tarReader := tar.NewReader(uncompressedStream)

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

		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:
			fmt.Printf("Unable to handle filetype %c in %s\n", header.Typeflag, 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")) {
			fmt.Printf("Processing YAML file in 'pis' directory: %s\n", 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 {
			fmt.Printf("Error removing %s: %v\n", path, err)
		}
	}
}