<commit_before>// Copyright 2014 Brian J. Downs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package spinner is a simple package to add a spinner to your application.
package spinner

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"
)

// CharSets contains the available character sets
var CharSets = [][]string{
	{"←", "↖", "↑", "↗", "→", "↘", "↓", "↙"},
	{"▁", "▃", "▄", "▅", "▆", "▇", "█", "▇", "▆", "▅", "▄", "▃"},
	{"▖", "▘", "▝", "▗"},
	{"┤", "┘", "┴", "└", "├", "┌", "┬", "┐"},
	{"◢", "◣", "◤", "◥"},
	{"◰", "◳", "◲", "◱"},
	{"◴", "◷", "◶", "◵"},
	{"◐", "◓", "◑", "◒"},
	{".", "o", "O", "@", "*"},
	{"|", "/", "-", "\\"},
	{"◡◡", "⊙⊙", "◠◠"},
	{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"},
	{">))'>", " >))'>", "  >))'>", "   >))'>", "    >))'>", "   <'((<", "  <'((<", " <'((<"},
	{"⠁", "⠂", "⠄", "⡀", "⢀", "⠠", "⠐", "⠈"},
	{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"},
	{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"},
	{"▉", "▊", "▋", "▌", "▍", "▎", "▏", "▎", "▍", "▌", "▋", "▊", "▉"},
	{"■", "□", "▪", "▫"},
	{"←", "↑", "→", "↓"},
	{"╫", "╪"},
	{"⇐", "⇖", "⇑", "⇗", "⇒", "⇘", "⇓", "⇙"},
	{"⠁", "⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈", "⠈"},
	{"⠈", "⠉", "⠋", "⠓", "⠒", "⠐", "⠐", "⠒", "⠖", "⠦", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈"},
	{"⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠴", "⠲", "⠒", "⠂", "⠂", "⠒", "⠚", "⠙", "⠉", "⠁"},
	{"⠋", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋"},
	{"ｦ", "ｧ", "ｨ", "ｩ", "ｪ", "ｫ", "ｬ", "ｭ", "ｮ", "ｯ", "ｱ", "ｲ", "ｳ", "ｴ", "ｵ", "ｶ", "ｷ", "ｸ", "ｹ", "ｺ", "ｻ", "ｼ", "ｽ", "ｾ", "ｿ", "ﾀ", "ﾁ", "ﾂ", "ﾃ", "ﾄ", "ﾅ", "ﾆ", "ﾇ", "ﾈ", "ﾉ", "ﾊ", "ﾋ", "ﾌ", "ﾍ", "ﾎ", "ﾏ", "ﾐ", "ﾑ", "ﾒ", "ﾓ", "ﾔ", "ﾕ", "ﾖ", "ﾗ", "ﾘ", "ﾙ", "ﾚ", "ﾛ", "ﾜ", "ﾝ"},
	{".", "..", "..."},
    {"▁","▂","▃","▄","▅","▆","▇","█","▉","▊","▋","▌","▍","▎","▏","▏","▎","▍","▌","▋","▊","▉","█","▇","▆","▅","▄","▃","▂","▁"},
}

// Spinner struct to hold the provided options
type Spinner struct {
	chars    []string
	Delay    time.Duration
	Prefix   string
	Suffix   string
	stopChan chan bool
	st       state
	w        io.Writer // to make testing better
	sync.Mutex
}

type state uint8

const (
	stopped state = iota
	running
)

var runlock sync.Mutex

// New provides a pointer to an instance of Spinner with the supplied options
func New(c []string, t time.Duration) *Spinner {
	s := &Spinner{
		Delay:    t,
		stopChan: make(chan bool, 1),
		w:        os.Stdout,
	}
	s.UpdateCharSet(c)
	return s
}

// Start will start the spinner
func (s *Spinner) Start() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		return
	}
	s.st = running

	go func() {
		runlock.Lock()
		defer runlock.Unlock()
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					out := fmt.Sprintf("%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.w, out)
					time.Sleep(s.Delay)
					erase(s.w, out)
				}
			}
		}
	}()
}

// erase deletes written characters
func erase(w io.Writer, a string) {
	n := utf8.RuneCountInString(a)
	for i := 0; i < n; i++ {
		fmt.Fprintf(w, "\b")
	}
}

// Stop stops the spinner
func (s *Spinner) Stop() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		s.stopChan <- true
		s.st = stopped
	}
}

// Restart will stop and start the spinner
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to that spinner
func (s *Spinner) Reverse() {
	s.Lock()
	defer s.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// UpdateSpeed is a convenience function to not have to make you
//create a new instance of the Spinner
func (s *Spinner) UpdateSpeed(delay time.Duration) { s.Delay = delay }

// UpdateCharSet will change the previously select character set to
// the provided one
func (s *Spinner) UpdateCharSet(chars []string) {
	// so that changes to the slice outside of the spinner don't change it
	// unexpectedly, create an internal copy
	s.Lock()
	defer s.Unlock()
	n := make([]string, len(chars))
	copy(n, chars)
	s.chars = n
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, 0)
	for i := 0; i < length; i++ {
		numSeq = append(numSeq, strconv.Itoa(i))
	}
	return numSeq
}
<commit_msg>Fixed formatting on new spinner<commit_after>// Copyright 2014 Brian J. Downs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package spinner is a simple package to add a spinner to your application.
package spinner

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"
)

// CharSets contains the available character sets
var CharSets = [][]string{
	{"←", "↖", "↑", "↗", "→", "↘", "↓", "↙"},
	{"▁", "▃", "▄", "▅", "▆", "▇", "█", "▇", "▆", "▅", "▄", "▃"},
	{"▖", "▘", "▝", "▗"},
	{"┤", "┘", "┴", "└", "├", "┌", "┬", "┐"},
	{"◢", "◣", "◤", "◥"},
	{"◰", "◳", "◲", "◱"},
	{"◴", "◷", "◶", "◵"},
	{"◐", "◓", "◑", "◒"},
	{".", "o", "O", "@", "*"},
	{"|", "/", "-", "\\"},
	{"◡◡", "⊙⊙", "◠◠"},
	{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"},
	{">))'>", " >))'>", "  >))'>", "   >))'>", "    >))'>", "   <'((<", "  <'((<", " <'((<"},
	{"⠁", "⠂", "⠄", "⡀", "⢀", "⠠", "⠐", "⠈"},
	{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"},
	{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"},
	{"▉", "▊", "▋", "▌", "▍", "▎", "▏", "▎", "▍", "▌", "▋", "▊", "▉"},
	{"■", "□", "▪", "▫"},
	{"←", "↑", "→", "↓"},
	{"╫", "╪"},
	{"⇐", "⇖", "⇑", "⇗", "⇒", "⇘", "⇓", "⇙"},
	{"⠁", "⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈", "⠈"},
	{"⠈", "⠉", "⠋", "⠓", "⠒", "⠐", "⠐", "⠒", "⠖", "⠦", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈"},
	{"⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠴", "⠲", "⠒", "⠂", "⠂", "⠒", "⠚", "⠙", "⠉", "⠁"},
	{"⠋", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋"},
	{"ｦ", "ｧ", "ｨ", "ｩ", "ｪ", "ｫ", "ｬ", "ｭ", "ｮ", "ｯ", "ｱ", "ｲ", "ｳ", "ｴ", "ｵ", "ｶ", "ｷ", "ｸ", "ｹ", "ｺ", "ｻ", "ｼ", "ｽ", "ｾ", "ｿ", "ﾀ", "ﾁ", "ﾂ", "ﾃ", "ﾄ", "ﾅ", "ﾆ", "ﾇ", "ﾈ", "ﾉ", "ﾊ", "ﾋ", "ﾌ", "ﾍ", "ﾎ", "ﾏ", "ﾐ", "ﾑ", "ﾒ", "ﾓ", "ﾔ", "ﾕ", "ﾖ", "ﾗ", "ﾘ", "ﾙ", "ﾚ", "ﾛ", "ﾜ", "ﾝ"},
	{".", "..", "..."},
	{"▁", "▂", "▃", "▄", "▅", "▆", "▇", "█", "▉", "▊", "▋", "▌", "▍", "▎", "▏", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█", "▇", "▆", "▅", "▄", "▃", "▂", "▁"},
}

// Spinner struct to hold the provided options
type Spinner struct {
	chars    []string
	Delay    time.Duration
	Prefix   string
	Suffix   string
	stopChan chan bool
	st       state
	w        io.Writer // to make testing better
	sync.Mutex
}

type state uint8

const (
	stopped state = iota
	running
)

var runlock sync.Mutex

// New provides a pointer to an instance of Spinner with the supplied options
func New(c []string, t time.Duration) *Spinner {
	s := &Spinner{
		Delay:    t,
		stopChan: make(chan bool, 1),
		w:        os.Stdout,
	}
	s.UpdateCharSet(c)
	return s
}

// Start will start the spinner
func (s *Spinner) Start() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		return
	}
	s.st = running

	go func() {
		runlock.Lock()
		defer runlock.Unlock()
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					out := fmt.Sprintf("%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.w, out)
					time.Sleep(s.Delay)
					erase(s.w, out)
				}
			}
		}
	}()
}

// erase deletes written characters
func erase(w io.Writer, a string) {
	n := utf8.RuneCountInString(a)
	for i := 0; i < n; i++ {
		fmt.Fprintf(w, "\b")
	}
}

// Stop stops the spinner
func (s *Spinner) Stop() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		s.stopChan <- true
		s.st = stopped
	}
}

// Restart will stop and start the spinner
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to that spinner
func (s *Spinner) Reverse() {
	s.Lock()
	defer s.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// UpdateSpeed is a convenience function to not have to make you
//create a new instance of the Spinner
func (s *Spinner) UpdateSpeed(delay time.Duration) { s.Delay = delay }

// UpdateCharSet will change the previously select character set to
// the provided one
func (s *Spinner) UpdateCharSet(chars []string) {
	// so that changes to the slice outside of the spinner don't change it
	// unexpectedly, create an internal copy
	s.Lock()
	defer s.Unlock()
	n := make([]string, len(chars))
	copy(n, chars)
	s.chars = n
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, 0)
	for i := 0; i < length; i++ {
		numSeq = append(numSeq, strconv.Itoa(i))
	}
	return numSeq
}
<|endoftext|>


<commit_before>// Copyright 2014 Brian J. Downs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package spinner is a simple package to add a spinner to your application.
package spinner

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
)

// CharSets contains the available character sets
var CharSets = [][]string{
	{"←", "↖", "↑", "↗", "→", "↘", "↓", "↙"},
	{"▁", "▃", "▄", "▅", "▆", "▇", "█", "▇", "▆", "▅", "▄", "▃"},
	{"▖", "▘", "▝", "▗"},
	{"┤", "┘", "┴", "└", "├", "┌", "┬", "┐"},
	{"◢", "◣", "◤", "◥"},
	{"◰", "◳", "◲", "◱"},
	{"◴", "◷", "◶", "◵"},
	{"◐", "◓", "◑", "◒"},
	{".", "o", "O", "@", "*"},
	{"|", "/", "-", "\\"},
	{"◡◡", "⊙⊙", "◠◠"},
	{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"},
	{">))'>", " >))'>", "  >))'>", "   >))'>", "    >))'>", "   <'((<", "  <'((<", " <'((<"},
	{"⠁", "⠂", "⠄", "⡀", "⢀", "⠠", "⠐", "⠈"},
	{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"},
	{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"},
	{"▉", "▊", "▋", "▌", "▍", "▎", "▏", "▎", "▍", "▌", "▋", "▊", "▉"},
	{"■", "□", "▪", "▫"},
	{"←", "↑", "→", "↓"},
	{"╫", "╪"},
	{"⇐", "⇖", "⇑", "⇗", "⇒", "⇘", "⇓", "⇙"},
	{"⠁", "⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈", "⠈"},
	{"⠈", "⠉", "⠋", "⠓", "⠒", "⠐", "⠐", "⠒", "⠖", "⠦", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈"},
	{"⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠴", "⠲", "⠒", "⠂", "⠂", "⠒", "⠚", "⠙", "⠉", "⠁"},
	{"⠋", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋"},
	{"ｦ", "ｧ", "ｨ", "ｩ", "ｪ", "ｫ", "ｬ", "ｭ", "ｮ", "ｯ", "ｱ", "ｲ", "ｳ", "ｴ", "ｵ", "ｶ", "ｷ", "ｸ", "ｹ", "ｺ", "ｻ", "ｼ", "ｽ", "ｾ", "ｿ", "ﾀ", "ﾁ", "ﾂ", "ﾃ", "ﾄ", "ﾅ", "ﾆ", "ﾇ", "ﾈ", "ﾉ", "ﾊ", "ﾋ", "ﾌ", "ﾍ", "ﾎ", "ﾏ", "ﾐ", "ﾑ", "ﾒ", "ﾓ", "ﾔ", "ﾕ", "ﾖ", "ﾗ", "ﾘ", "ﾙ", "ﾚ", "ﾛ", "ﾜ", "ﾝ"},
	{".", "..", "..."},
	{"▁", "▂", "▃", "▄", "▅", "▆", "▇", "█", "▉", "▊", "▋", "▌", "▍", "▎", "▏", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█", "▇", "▆", "▅", "▄", "▃", "▂", "▁"},
	{".", "o", "O", "°", "O", "o", "."},
	{"+", "x"},
	{"v", "<", "^", ">"},
}

// state is a type for the spinner status
type state uint8

// Spinner struct to hold the provided options
type Spinner struct {
	chars    []string                      // chosen character set
	Delay    time.Duration                 // speed of the spinner
	Prefix   string                        // Text preppended to the spinner
	Suffix   string                        // Text appended to the spinner
	stopChan chan bool                     // channel used to stop the spinner
	st       state                         // spinner status
	w        io.Writer                     // to make testing better
	color    func(a ...interface{}) string // default color is white
	sync.Mutex
}

const (
	stopped state = iota
	running
)

var runlock sync.Mutex

// validColors holds an array of the only colors allowed
var validColors = []string{"red", "green", "yellow", "blue", "magenta", "cyan", "white"}

// validColor will make sure the given color is actually allowed
func validColor(c string) bool {
	for _, i := range validColors {
		if c == i {
			return true
		}
	}
	return false
}

// New provides a pointer to an instance of Spinner with the supplied options
func New(c []string, t time.Duration) *Spinner {
	s := &Spinner{
		Delay:    t,
		stopChan: make(chan bool, 1),
		color:    color.New(color.FgWhite).SprintFunc(),
		w:        os.Stdout,
	}
	s.UpdateCharSet(c)
	return s
}

// Start will start the spinner
func (s *Spinner) Start() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		return
	}
	s.st = running
	go func() {
		runlock.Lock()
		defer runlock.Unlock()
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					out := fmt.Sprintf("%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
					fmt.Fprint(s.w, out)
					time.Sleep(s.Delay)
					erase(s.w, out)
				}
			}
		}
	}()
}

// erase deletes written characters
func erase(w io.Writer, a string) {
	n := utf8.RuneCountInString(a)
	for i := 0; i < n; i++ {
		fmt.Fprintf(w, "\b")
	}
}

// Color will set the struct field for the given color to be used
func (s *Spinner) Color(c string) error {
	if validColor(c) {
		switch {
		case c == "red":
			s.color = color.New(color.FgRed).SprintFunc()
			s.Restart()
		case c == "yellow":
			s.color = color.New(color.FgYellow).SprintFunc()
			s.Restart()
		case c == "green":
			s.color = color.New(color.FgGreen).SprintFunc()
			s.Restart()
		case c == "magenta":
			s.color = color.New(color.FgMagenta).SprintFunc()
			s.Restart()
		case c == "blue":
			s.color = color.New(color.FgBlue).SprintFunc()
			s.Restart()
		case c == "cyan":
			s.color = color.New(color.FgCyan).SprintFunc()
			s.Restart()
		case c == "white":
			s.color = color.New(color.FgWhite).SprintFunc()
			s.Restart()
		default:
			return errors.New("invalid color")
		}
	}
	return nil
}

// Stop stops the spinner
func (s *Spinner) Stop() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		s.stopChan <- true
		s.st = stopped
	}
}

// Restart will stop and start the spinner
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to that spinner
func (s *Spinner) Reverse() {
	s.Lock()
	defer s.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// UpdateSpeed is a convenience function to not have to make you
//create a new instance of the Spinner
func (s *Spinner) UpdateSpeed(delay time.Duration) { s.Delay = delay }

// UpdateCharSet will change the previously select character set to
// the provided one
func (s *Spinner) UpdateCharSet(chars []string) {
	// so that changes to the slice outside of the spinner don't change it
	// unexpectedly, create an internal copy
	s.Lock()
	defer s.Unlock()
	n := make([]string, len(chars))
	copy(n, chars)
	s.chars = n
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, 0)
	for i := 0; i < length; i++ {
		numSeq = append(numSeq, strconv.Itoa(i))
	}
	return numSeq
}
<commit_msg>add another spinner<commit_after>// Copyright 2014 Brian J. Downs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package spinner is a simple package to add a spinner to your application.
package spinner

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
)

// CharSets contains the available character sets
var CharSets = [][]string{
	{"←", "↖", "↑", "↗", "→", "↘", "↓", "↙"},
	{"▁", "▃", "▄", "▅", "▆", "▇", "█", "▇", "▆", "▅", "▄", "▃"},
	{"▖", "▘", "▝", "▗"},
	{"┤", "┘", "┴", "└", "├", "┌", "┬", "┐"},
	{"◢", "◣", "◤", "◥"},
	{"◰", "◳", "◲", "◱"},
	{"◴", "◷", "◶", "◵"},
	{"◐", "◓", "◑", "◒"},
	{".", "o", "O", "@", "*"},
	{"|", "/", "-", "\\"},
	{"◡◡", "⊙⊙", "◠◠"},
	{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"},
	{">))'>", " >))'>", "  >))'>", "   >))'>", "    >))'>", "   <'((<", "  <'((<", " <'((<"},
	{"⠁", "⠂", "⠄", "⡀", "⢀", "⠠", "⠐", "⠈"},
	{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"},
	{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"},
	{"▉", "▊", "▋", "▌", "▍", "▎", "▏", "▎", "▍", "▌", "▋", "▊", "▉"},
	{"■", "□", "▪", "▫"},
	{"←", "↑", "→", "↓"},
	{"╫", "╪"},
	{"⇐", "⇖", "⇑", "⇗", "⇒", "⇘", "⇓", "⇙"},
	{"⠁", "⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈", "⠈"},
	{"⠈", "⠉", "⠋", "⠓", "⠒", "⠐", "⠐", "⠒", "⠖", "⠦", "⠤", "⠠", "⠠", "⠤", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋", "⠉", "⠈"},
	{"⠁", "⠉", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠤", "⠄", "⠄", "⠤", "⠴", "⠲", "⠒", "⠂", "⠂", "⠒", "⠚", "⠙", "⠉", "⠁"},
	{"⠋", "⠙", "⠚", "⠒", "⠂", "⠂", "⠒", "⠲", "⠴", "⠦", "⠖", "⠒", "⠐", "⠐", "⠒", "⠓", "⠋"},
	{"ｦ", "ｧ", "ｨ", "ｩ", "ｪ", "ｫ", "ｬ", "ｭ", "ｮ", "ｯ", "ｱ", "ｲ", "ｳ", "ｴ", "ｵ", "ｶ", "ｷ", "ｸ", "ｹ", "ｺ", "ｻ", "ｼ", "ｽ", "ｾ", "ｿ", "ﾀ", "ﾁ", "ﾂ", "ﾃ", "ﾄ", "ﾅ", "ﾆ", "ﾇ", "ﾈ", "ﾉ", "ﾊ", "ﾋ", "ﾌ", "ﾍ", "ﾎ", "ﾏ", "ﾐ", "ﾑ", "ﾒ", "ﾓ", "ﾔ", "ﾕ", "ﾖ", "ﾗ", "ﾘ", "ﾙ", "ﾚ", "ﾛ", "ﾜ", "ﾝ"},
	{".", "..", "..."},
	{"▁", "▂", "▃", "▄", "▅", "▆", "▇", "█", "▉", "▊", "▋", "▌", "▍", "▎", "▏", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█", "▇", "▆", "▅", "▄", "▃", "▂", "▁"},
	{".", "o", "O", "°", "O", "o", "."},
	{"+", "x"},
	{"v", "<", "^", ">"},
	{">>--->", " >>--->", "  >>--->", "   >>--->", "    >>--->", "    <---<<", "   <---<<", "  <---<<", " <---<<", "<---<<"},
}

// state is a type for the spinner status
type state uint8

// Spinner struct to hold the provided options
type Spinner struct {
	chars    []string                      // chosen character set
	Delay    time.Duration                 // speed of the spinner
	Prefix   string                        // Text preppended to the spinner
	Suffix   string                        // Text appended to the spinner
	stopChan chan bool                     // channel used to stop the spinner
	st       state                         // spinner status
	w        io.Writer                     // to make testing better
	color    func(a ...interface{}) string // default color is white
	sync.Mutex
}

const (
	stopped state = iota
	running
)

var runlock sync.Mutex

// validColors holds an array of the only colors allowed
var validColors = []string{"red", "green", "yellow", "blue", "magenta", "cyan", "white"}

// validColor will make sure the given color is actually allowed
func validColor(c string) bool {
	for _, i := range validColors {
		if c == i {
			return true
		}
	}
	return false
}

// New provides a pointer to an instance of Spinner with the supplied options
func New(c []string, t time.Duration) *Spinner {
	s := &Spinner{
		Delay:    t,
		stopChan: make(chan bool, 1),
		color:    color.New(color.FgWhite).SprintFunc(),
		w:        os.Stdout,
	}
	s.UpdateCharSet(c)
	return s
}

// Start will start the spinner
func (s *Spinner) Start() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		return
	}
	s.st = running
	go func() {
		runlock.Lock()
		defer runlock.Unlock()
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					out := fmt.Sprintf("%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
					fmt.Fprint(s.w, out)
					time.Sleep(s.Delay)
					erase(s.w, out)
				}
			}
		}
	}()
}

// erase deletes written characters
func erase(w io.Writer, a string) {
	n := utf8.RuneCountInString(a)
	for i := 0; i < n; i++ {
		fmt.Fprintf(w, "\b")
	}
}

// Color will set the struct field for the given color to be used
func (s *Spinner) Color(c string) error {
	if validColor(c) {
		switch {
		case c == "red":
			s.color = color.New(color.FgRed).SprintFunc()
			s.Restart()
		case c == "yellow":
			s.color = color.New(color.FgYellow).SprintFunc()
			s.Restart()
		case c == "green":
			s.color = color.New(color.FgGreen).SprintFunc()
			s.Restart()
		case c == "magenta":
			s.color = color.New(color.FgMagenta).SprintFunc()
			s.Restart()
		case c == "blue":
			s.color = color.New(color.FgBlue).SprintFunc()
			s.Restart()
		case c == "cyan":
			s.color = color.New(color.FgCyan).SprintFunc()
			s.Restart()
		case c == "white":
			s.color = color.New(color.FgWhite).SprintFunc()
			s.Restart()
		default:
			return errors.New("invalid color")
		}
	}
	return nil
}

// Stop stops the spinner
func (s *Spinner) Stop() {
	s.Lock()
	defer s.Unlock()
	if s.st == running {
		s.stopChan <- true
		s.st = stopped
	}
}

// Restart will stop and start the spinner
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to that spinner
func (s *Spinner) Reverse() {
	s.Lock()
	defer s.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// UpdateSpeed is a convenience function to not have to make you
//create a new instance of the Spinner
func (s *Spinner) UpdateSpeed(delay time.Duration) { s.Delay = delay }

// UpdateCharSet will change the previously select character set to
// the provided one
func (s *Spinner) UpdateCharSet(chars []string) {
	// so that changes to the slice outside of the spinner don't change it
	// unexpectedly, create an internal copy
	s.Lock()
	defer s.Unlock()
	n := make([]string, len(chars))
	copy(n, chars)
	s.chars = n
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, 0)
	for i := 0; i < length; i++ {
		numSeq = append(numSeq, strconv.Itoa(i))
	}
	return numSeq
}
<|endoftext|>


<commit_before>package martiancheck

import (
  "net"
)

// IPChecker is an interface to check if an IP address is within a martian address range or not.
// This was designed as an interface to allow for both IPv4 and IPv6 to follow the same pattern.
type MartianChecker interface {
  Contains(ip net.IP) bool
  AddMartianNet(net net.IPNet)
  RemoveMartianNet(net net.IPNet) bool
  RemoveAllMartianNet() bool
}

// IPv4Checker to test if an address is within a martian range or not.
type MartianIPv4Checker struct {
  Martians []net.IPNet
}

// NewIPv4MartianChecker returns a configured IPv4 Checker struct.
// Provides all default martian ranges
func NewMartianIPv4Checker(populate bool) MartianIPv4Checker {
  ipch := IPv4Checker{}
  //create unused blocks
  ipch.IPUnused = make([]net.IPNet,0)
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{0,0,0,0},Mask:ipch.Masks["Mask8Bit"]})          //This Net RFC1122
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{10,0,0,0},Mask:ipch.Masks["Mask8Bit"]})         //Private IP RFC1918
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{127,0,0,0},Mask:ipch.Masks["Mask8Bit"]})        //Loopback IP RFC1122
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{169,254,0,0},Mask:ipch.Masks["Mask16Bit"]})     //Local link address
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{127,16,0,0},Mask:ipch.Masks["Mask12Bit"]})      //Private IP RFC1918
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,0,0,0},Mask:ipch.Masks["Mask24Bit"]})       //IETF Protocol Assignments RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,0,2,0},Mask:ipch.Masks["Mask24Bit"]})       //Test-Net-1 RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,168,0,0},Mask:ipch.Masks["Mask16Bit"]})     //Private IP RFC1918
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,18,0,0},Mask:ipch.Masks["Mask15Bit"]})      //Network Testing
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{198,51,100,0},Mask:ipch.Masks["Mask24Bit"]})    //Test-Net-2 RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{203,0,113,0},Mask:ipch.Masks["Mask24Bit"]})     //Test-Net-3 RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{224,0,0,0},Mask:ipch.Masks["Mask4Bit"]})        //Multicast RFC3171
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{240,0,0,0},Mask:ipch.Masks["Mask4Bit"]})        //Future Use RFC1122
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{255,255,255,255},Mask:ipch.Masks["Mask32Bit"]}) //Broadcast
  return ipch
}

func (ipch *IPv4Checker) Contains(ip net.IP) bool {
  for item := range ipch.IPUnused {
    if ipch.IPUnused[item].Contains(ip) {
      return true
    }
  }
  return false
}

func (ipch *IPv4Checker) AddMartianNet(net net.IPNet) {
  ipch.IPUnused = append(ipch.IPUnused,net)
}

func (ipch *IPv4Checker) RemoveMartianNet(net net.IPNet) bool {
  for item := range ipch.IPUnused {
    if net.String() == ipch.IPUnused[item].String() {
      copy(ipch.IPUnused[item:], ipch.IPUnused[item+1:])
      ipch.IPUnused = ipch.IPUnused[:len(ipch.IPUnused)-1]
      return true
    }
  }
  return false
}

func (ipch *IPv4Checker) RemoveAllMartianNet() {
  ipch.IPUnused = make([]net.IPNet,0)
}

func main() {
  ipch := NewIPv4MartianChecker(true)
  log.Println(len(ipch.IPUnused))
  log.Println(ipch.Contains([]byte{11,0,0,1}))
  ipch.AddMartianNet(net.IPNet{IP:[]byte{11,0,0,0},Mask:ipch.Masks["Mask24Bit"]})
  log.Println(ipch.Contains([]byte{11,0,0,1}))
  ipch.RemoveMartianNet(net.IPNet{IP:[]byte{11,0,0,0},Mask:ipch.Masks["Mask24Bit"]})
  log.Println(ipch.Contains([]byte{11,0,0,1}))
  ipch.RemoveAllMartianNet()
  log.Println(len(ipch.IPUnused))
}
<commit_msg>added eazymask to martian check<commit_after>package martiancheck

import (
  "net"
  "github.com/robwc/easymask/ezmv4"
)

// IPChecker is an interface to check if an IP address is within a martian address range or not.
// This was designed as an interface to allow for both IPv4 and IPv6 to follow the same pattern.
type MartianChecker interface {
  Contains(ip net.IP) bool
  AddMartianNet(net net.IPNet)
  RemoveMartianNet(net net.IPNet) bool
  RemoveAllMartianNet() bool
}

// IPv4Checker to test if an address is within a martian range or not.
type MartianIPv4Checker struct {
  Martians []net.IPNet
}

// NewIPv4MartianChecker returns a configured IPv4 Checker struct.
// Provides all default martian ranges
func NewMartianIPv4Checker(populate bool) MartianIPv4Checker {
  ipch := MartianChecker{}
  maskch := ezmv4.NewIPv4Masks()
  //create unused blocks
  ipch.IPUnused = make([]net.IPNet,0)
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{0,0,0,0},Mask:maskch.Masks["8"]})          //This Net RFC1122
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{10,0,0,0},Mask:maskch.Masks["8"]})         //Private IP RFC1918
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{127,0,0,0},Mask:maskch.Masks["8"]})        //Loopback IP RFC1122
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{169,254,0,0},Mask:maskch.Masks["16"]})     //Local link address
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{127,16,0,0},Mask:maskch.Masks["12"]})      //Private IP RFC1918
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,0,0,0},Mask:maskch.Masks["24"]})       //IETF Protocol Assignments RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,0,2,0},Mask:maskch.Masks["24"]})       //Test-Net-1 RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,168,0,0},Mask:maskch.Masks["16"]})     //Private IP RFC1918
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{192,18,0,0},Mask:maskch.Masks["15"]})      //Network Testing
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{198,51,100,0},Mask:maskch.Masks["24"]})    //Test-Net-2 RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{203,0,113,0},Mask:maskch.Masks["24"]})     //Test-Net-3 RFC5736
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{224,0,0,0},Mask:maskch.Masks["4"]})        //Multicast RFC3171
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{240,0,0,0},Mask:maskch.Masks["4"]})        //Future Use RFC1122
  ipch.IPUnused = append(ipch.IPUnused,net.IPNet{IP:[]byte{255,255,255,255},Mask:maskch.Masks["32"]}) //Broadcast
  return ipch
}

func (ipch *IPv4Checker) Contains(ip net.IP) bool {
  for item := range ipch.IPUnused {
    if ipch.IPUnused[item].Contains(ip) {
      return true
    }
  }
  return false
}

func (ipch *IPv4Checker) AddMartianNet(net net.IPNet) {
  ipch.IPUnused = append(ipch.IPUnused,net)
}

func (ipch *IPv4Checker) RemoveMartianNet(net net.IPNet) bool {
  for item := range ipch.IPUnused {
    if net.String() == ipch.IPUnused[item].String() {
      copy(ipch.IPUnused[item:], ipch.IPUnused[item+1:])
      ipch.IPUnused = ipch.IPUnused[:len(ipch.IPUnused)-1]
      return true
    }
  }
  return false
}

func (ipch *IPv4Checker) RemoveAllMartianNet() {
  ipch.IPUnused = make([]net.IPNet,0)
}

func main() {
  ipch := NewIPv4MartianChecker(true)
  maskch := ezmv4.NewIPv4Masks()
  log.Println(len(ipch.IPUnused))
  log.Println(ipch.Contains([]byte{11,0,0,1}))
  ipch.AddMartianNet(net.IPNet{IP:[]byte{11,0,0,0},Mask:maskch.Masks["24"]})
  log.Println(ipch.Contains([]byte{11,0,0,1}))
  ipch.RemoveMartianNet(net.IPNet{IP:[]byte{11,0,0,0},Mask:maskch.Masks["24"]})
  log.Println(ipch.Contains([]byte{11,0,0,1}))
  ipch.RemoveAllMartianNet()
  log.Println(len(ipch.IPUnused))
}
<|endoftext|>


<commit_before>package protocol

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
)

// JSONInput is what's sent into the function
// All HTTP request headers should be set in env
type JSONInput struct {
	Body string `json:"body"`
}

// JSONOutput function must return this format
// StatusCode value must be a HTTP status code
type JSONOutput struct {
	StatusCode int    `json:"status"`
	Body       string `json:"body"`
}

// JSONProtocol converts stdin/stdout streams from HTTP into JSON format.
type JSONProtocol struct {
	in  io.Writer
	out io.Reader
}

func (p *JSONProtocol) IsStreamable() bool {
	return true
}

func (h *JSONProtocol) Dispatch(w io.Writer, req *http.Request) error {
	var body bytes.Buffer
	if req.Body != nil {
		var dest io.Writer = &body

		// TODO copy w/ ctx
		nBytes, _ := strconv.ParseInt(
			req.Header.Get("Content-Length"), 10, 64)
		_, err := io.Copy(dest, io.LimitReader(req.Body, nBytes))
		if err != nil {
			// TODO: maybe mask this error if favour of something different
			return err
		}
	}

	// convert to JSON func format
	jin := &JSONInput{
		Body: body.String(),
	}
	b, err := json.Marshal(jin)
	if err != nil {
		// this shouldn't happen
		return fmt.Errorf("error marshalling JSONInput: %v", err)
	}
	h.in.Write(b)

	// TODO: put max size on how big the response can be so we don't blow up
	jout := &JSONOutput{}
	dec := json.NewDecoder(h.out)
	if err := dec.Decode(jout); err != nil {
		// TODO: how do we get an error back to the client??
		return fmt.Errorf("error unmarshalling JSONOutput: %v", err)
	}

	// res := &http.Response{}
	// res.Body = strings.NewReader(jout.Body)
	// TODO: shouldn't we pass back the full response object or something so we can set some things on it here?
	// For instance, user could set response content type or what have you.
	//io.Copy(cfg.Stdout, strings.NewReader(jout.Body))

	if rw, ok := w.(http.ResponseWriter); ok {
		b, err = json.Marshal(jout.Body)
		if err != nil {
			return fmt.Errorf("error unmarshalling JSONOutput.Body: %v", err)
		}
		rw.WriteHeader(jout.StatusCode)
		rw.Write(b) // TODO timeout
	} else {
		// logs can just copy the full thing in there, headers and all.
		b, err = json.Marshal(jout)
		if err != nil {
			return fmt.Errorf("error unmarshalling JSONOutput: %v", err)
		}

		w.Write(b) // TODO timeout

	}
	return nil

}
<commit_msg>Respond with any error that happends during JSON dispatching<commit_after>package protocol

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
)

// JSONInput is what's sent into the function
// All HTTP request headers should be set in env
type JSONInput struct {
	Body string `json:"body"`
}

// JSONOutput function must return this format
// StatusCode value must be a HTTP status code
type JSONOutput struct {
	StatusCode int    `json:"status"`
	Body       string `json:"body"`
}

// JSONProtocol converts stdin/stdout streams from HTTP into JSON format.
type JSONProtocol struct {
	in  io.Writer
	out io.Reader
}

func (p *JSONProtocol) IsStreamable() bool {
	return true
}

type Error struct {
	Message string `json:"message"`
}

type ErrMsg struct {
	Err Error `json:"error"`
}

func (h *JSONProtocol) Dispatch(w io.Writer, req *http.Request) error {
	var body bytes.Buffer
	if req.Body != nil {
		var dest io.Writer = &body

		// TODO copy w/ ctx
		nBytes, _ := strconv.ParseInt(
			req.Header.Get("Content-Length"), 10, 64)
		_, err := io.Copy(dest, io.LimitReader(req.Body, nBytes))
		if err != nil {
			respondWithError(w, err)
			return err
		}
	}

	// convert to JSON func format
	jin := &JSONInput{
		Body: body.String(),
	}
	b, err := json.Marshal(jin)
	if err != nil {
		// this shouldn't happen
		err = fmt.Errorf("error marshalling JSONInput: %v", err)
		respondWithError(w, err)
		return err
	}
	h.in.Write(b)

	maxContentSize := int64(1 * 1024 * 1024) // 1Mb should be enough
	jout := &JSONOutput{}
	dec := json.NewDecoder(io.LimitReader(h.out, maxContentSize))
	if err := dec.Decode(jout); err != nil {
		err = fmt.Errorf("Unable to decode JSON response object: %s", err.Error())
		respondWithError(w, err)
		return err
	}

	if rw, ok := w.(http.ResponseWriter); ok {
		b, err = json.Marshal(jout.Body)
		if err != nil {
			err = fmt.Errorf("error unmarshalling JSON body: %s", err.Error())
			respondWithError(w, err)
			return err
		}
		rw.WriteHeader(jout.StatusCode)
		rw.Write(b) // TODO timeout
	} else {
		// logs can just copy the full thing in there, headers and all.
		b, err = json.Marshal(jout)
		if err != nil {
			err = fmt.Errorf("error unmarshalling JSON response: %s", err.Error())
			respondWithError(w, err)
			return err
		}
		w.Write(b) // TODO timeout
	}
	return nil
}

func respondWithError(w io.Writer, err error) {
	errMsg := ErrMsg{
		Err: Error{
			Message: err.Error(),
		},
	}
	b, _ := json.Marshal(errMsg)
	statusCode := 500
	writeResponse(w, b, statusCode)
}

func writeResponse(w io.Writer, b []byte, statusCode int) {
	if rw, ok := w.(http.ResponseWriter); ok {
		rw.WriteHeader(statusCode)
		rw.Write(b) // TODO timeout
	} else {
		// logs can just copy the full thing in there, headers and all.
		w.Write(b) // TODO timeout
	}
}
<|endoftext|>


<commit_before>package remote

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/Azure/azure-sdk-for-go/arm/storage"
	mainStorage "github.com/Azure/azure-sdk-for-go/storage"
	"github.com/Azure/go-autorest/autorest/azure"
	riviera "github.com/jen20/riviera/azure"
)

func azureFactory(conf map[string]string) (Client, error) {
	storageAccountName, ok := conf["storage_account_name"]
	if !ok {
		return nil, fmt.Errorf("missing 'storage_account_name' configuration")
	}
	containerName, ok := conf["container_name"]
	if !ok {
		return nil, fmt.Errorf("missing 'container_name' configuration")
	}
	keyName, ok := conf["key"]
	if !ok {
		return nil, fmt.Errorf("missing 'key' configuration")
	}

	env, err := getAzureEnvironmentFromConf(conf)
	if err != nil {
		return nil, err
	}

	accessKey, ok := confOrEnv(conf, "access_key", "ARM_ACCESS_KEY")
	if !ok {
		resourceGroupName, ok := conf["resource_group_name"]
		if !ok {
			return nil, fmt.Errorf("missing 'resource_group' configuration")
		}

		var err error
		accessKey, err = getStorageAccountAccessKey(conf, resourceGroupName, storageAccountName, env)
		if err != nil {
			return nil, fmt.Errorf("Couldn't read access key from storage account: %s.", err)
		}
	}

	storageClient, err := mainStorage.NewClient(storageAccountName, accessKey, env.StorageEndpointSuffix,
		mainStorage.DefaultAPIVersion, true)
	if err != nil {
		return nil, fmt.Errorf("Error creating storage client for storage account %q: %s", storageAccountName, err)
	}

	blobClient := storageClient.GetBlobService()
	leaseID, _ := confOrEnv(conf, "lease_id", "ARM_LEASE_ID")

	return &AzureClient{
		blobClient:    &blobClient,
		containerName: containerName,
		keyName:       keyName,
		leaseID:       leaseID,
	}, nil
}

func getStorageAccountAccessKey(conf map[string]string, resourceGroupName, storageAccountName string, env azure.Environment) (string, error) {
	creds, err := getCredentialsFromConf(conf, env)
	if err != nil {
		return "", err
	}

	oauthConfig, err := env.OAuthConfigForTenant(creds.TenantID)
	if err != nil {
		return "", err
	}
	if oauthConfig == nil {
		return "", fmt.Errorf("Unable to configure OAuthConfig for tenant %s", creds.TenantID)
	}

	spt, err := azure.NewServicePrincipalToken(*oauthConfig, creds.ClientID, creds.ClientSecret, env.ResourceManagerEndpoint)
	if err != nil {
		return "", err
	}

	accountsClient := storage.NewAccountsClientWithBaseURI(env.ResourceManagerEndpoint, creds.SubscriptionID)
	accountsClient.Authorizer = spt

	keys, err := accountsClient.ListKeys(resourceGroupName, storageAccountName)
	if err != nil {
		return "", fmt.Errorf("Error retrieving keys for storage account %q: %s", storageAccountName, err)
	}

	if keys.Keys == nil {
		return "", fmt.Errorf("Nil key returned for storage account %q", storageAccountName)
	}

	accessKeys := *keys.Keys
	return *accessKeys[0].Value, nil
}

func getCredentialsFromConf(conf map[string]string, env azure.Environment) (*riviera.AzureResourceManagerCredentials, error) {
	subscriptionID, ok := confOrEnv(conf, "arm_subscription_id", "ARM_SUBSCRIPTION_ID")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_subscription_id' configuration")
	}
	clientID, ok := confOrEnv(conf, "arm_client_id", "ARM_CLIENT_ID")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_client_id' configuration")
	}
	clientSecret, ok := confOrEnv(conf, "arm_client_secret", "ARM_CLIENT_SECRET")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_client_secret' configuration")
	}
	tenantID, ok := confOrEnv(conf, "arm_tenant_id", "ARM_TENANT_ID")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_tenant_id' configuration")
	}

	return &riviera.AzureResourceManagerCredentials{
		SubscriptionID:          subscriptionID,
		ClientID:                clientID,
		ClientSecret:            clientSecret,
		TenantID:                tenantID,
		ActiveDirectoryEndpoint: env.ActiveDirectoryEndpoint,
		ResourceManagerEndpoint: env.ResourceManagerEndpoint,
	}, nil
}

func getAzureEnvironmentFromConf(conf map[string]string) (azure.Environment, error) {
	envName, ok := confOrEnv(conf, "environment", "ARM_ENVIRONMENT")
	if !ok {
		return azure.PublicCloud, nil
	}

	env, err := azure.EnvironmentFromName(envName)
	if err != nil {
		// try again with wrapped value to support readable values like german instead of AZUREGERMANCLOUD
		var innerErr error
		env, innerErr = azure.EnvironmentFromName(fmt.Sprintf("AZURE%sCLOUD", envName))
		if innerErr != nil {
			return env, fmt.Errorf("invalid 'environment' configuration: %s", err)
		}
	}

	return env, nil
}

func confOrEnv(conf map[string]string, confKey, envVar string) (string, bool) {
	value, ok := conf[confKey]
	if ok {
		return value, true
	}

	value = os.Getenv(envVar)

	return value, value != ""
}

type AzureClient struct {
	blobClient    *mainStorage.BlobStorageClient
	containerName string
	keyName       string
	leaseID       string
}

func (c *AzureClient) Get() (*Payload, error) {
	blob, err := c.blobClient.GetBlob(c.containerName, c.keyName)
	if err != nil {
		if storErr, ok := err.(mainStorage.AzureStorageServiceError); ok {
			if storErr.Code == "BlobNotFound" {
				return nil, nil
			}
		}
		return nil, err
	}

	defer blob.Close()

	data, err := ioutil.ReadAll(blob)
	if err != nil {
		return nil, err
	}

	payload := &Payload{
		Data: data,
	}

	// If there was no data, then return nil
	if len(payload.Data) == 0 {
		return nil, nil
	}

	return payload, nil
}

func (c *AzureClient) Put(data []byte) error {
	headers := map[string]string{
		"Content-Type": "application/json",
	}

	if c.leaseID != "" {
		headers["x-ms-lease-id"] = c.leaseID
	}

	return c.blobClient.CreateBlockBlobFromReader(
		c.containerName,
		c.keyName,
		uint64(len(data)),
		bytes.NewReader(data),
		headers,
	)
}

func (c *AzureClient) Delete() error {
	headers := map[string]string{}
	if c.leaseID != "" {
		headers["x-ms-lease-id"] = c.leaseID
	}

	return c.blobClient.DeleteBlob(c.containerName, c.keyName, headers)
}
<commit_msg>Fix error message in Azure state backend (#12424)<commit_after>package remote

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/Azure/azure-sdk-for-go/arm/storage"
	mainStorage "github.com/Azure/azure-sdk-for-go/storage"
	"github.com/Azure/go-autorest/autorest/azure"
	riviera "github.com/jen20/riviera/azure"
)

func azureFactory(conf map[string]string) (Client, error) {
	storageAccountName, ok := conf["storage_account_name"]
	if !ok {
		return nil, fmt.Errorf("missing 'storage_account_name' configuration")
	}
	containerName, ok := conf["container_name"]
	if !ok {
		return nil, fmt.Errorf("missing 'container_name' configuration")
	}
	keyName, ok := conf["key"]
	if !ok {
		return nil, fmt.Errorf("missing 'key' configuration")
	}

	env, err := getAzureEnvironmentFromConf(conf)
	if err != nil {
		return nil, err
	}

	accessKey, ok := confOrEnv(conf, "access_key", "ARM_ACCESS_KEY")
	if !ok {
		resourceGroupName, ok := conf["resource_group_name"]
		if !ok {
			return nil, fmt.Errorf("missing 'resource_group_name' configuration")
		}

		var err error
		accessKey, err = getStorageAccountAccessKey(conf, resourceGroupName, storageAccountName, env)
		if err != nil {
			return nil, fmt.Errorf("Couldn't read access key from storage account: %s.", err)
		}
	}

	storageClient, err := mainStorage.NewClient(storageAccountName, accessKey, env.StorageEndpointSuffix,
		mainStorage.DefaultAPIVersion, true)
	if err != nil {
		return nil, fmt.Errorf("Error creating storage client for storage account %q: %s", storageAccountName, err)
	}

	blobClient := storageClient.GetBlobService()
	leaseID, _ := confOrEnv(conf, "lease_id", "ARM_LEASE_ID")

	return &AzureClient{
		blobClient:    &blobClient,
		containerName: containerName,
		keyName:       keyName,
		leaseID:       leaseID,
	}, nil
}

func getStorageAccountAccessKey(conf map[string]string, resourceGroupName, storageAccountName string, env azure.Environment) (string, error) {
	creds, err := getCredentialsFromConf(conf, env)
	if err != nil {
		return "", err
	}

	oauthConfig, err := env.OAuthConfigForTenant(creds.TenantID)
	if err != nil {
		return "", err
	}
	if oauthConfig == nil {
		return "", fmt.Errorf("Unable to configure OAuthConfig for tenant %s", creds.TenantID)
	}

	spt, err := azure.NewServicePrincipalToken(*oauthConfig, creds.ClientID, creds.ClientSecret, env.ResourceManagerEndpoint)
	if err != nil {
		return "", err
	}

	accountsClient := storage.NewAccountsClientWithBaseURI(env.ResourceManagerEndpoint, creds.SubscriptionID)
	accountsClient.Authorizer = spt

	keys, err := accountsClient.ListKeys(resourceGroupName, storageAccountName)
	if err != nil {
		return "", fmt.Errorf("Error retrieving keys for storage account %q: %s", storageAccountName, err)
	}

	if keys.Keys == nil {
		return "", fmt.Errorf("Nil key returned for storage account %q", storageAccountName)
	}

	accessKeys := *keys.Keys
	return *accessKeys[0].Value, nil
}

func getCredentialsFromConf(conf map[string]string, env azure.Environment) (*riviera.AzureResourceManagerCredentials, error) {
	subscriptionID, ok := confOrEnv(conf, "arm_subscription_id", "ARM_SUBSCRIPTION_ID")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_subscription_id' configuration")
	}
	clientID, ok := confOrEnv(conf, "arm_client_id", "ARM_CLIENT_ID")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_client_id' configuration")
	}
	clientSecret, ok := confOrEnv(conf, "arm_client_secret", "ARM_CLIENT_SECRET")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_client_secret' configuration")
	}
	tenantID, ok := confOrEnv(conf, "arm_tenant_id", "ARM_TENANT_ID")
	if !ok {
		return nil, fmt.Errorf("missing 'arm_tenant_id' configuration")
	}

	return &riviera.AzureResourceManagerCredentials{
		SubscriptionID:          subscriptionID,
		ClientID:                clientID,
		ClientSecret:            clientSecret,
		TenantID:                tenantID,
		ActiveDirectoryEndpoint: env.ActiveDirectoryEndpoint,
		ResourceManagerEndpoint: env.ResourceManagerEndpoint,
	}, nil
}

func getAzureEnvironmentFromConf(conf map[string]string) (azure.Environment, error) {
	envName, ok := confOrEnv(conf, "environment", "ARM_ENVIRONMENT")
	if !ok {
		return azure.PublicCloud, nil
	}

	env, err := azure.EnvironmentFromName(envName)
	if err != nil {
		// try again with wrapped value to support readable values like german instead of AZUREGERMANCLOUD
		var innerErr error
		env, innerErr = azure.EnvironmentFromName(fmt.Sprintf("AZURE%sCLOUD", envName))
		if innerErr != nil {
			return env, fmt.Errorf("invalid 'environment' configuration: %s", err)
		}
	}

	return env, nil
}

func confOrEnv(conf map[string]string, confKey, envVar string) (string, bool) {
	value, ok := conf[confKey]
	if ok {
		return value, true
	}

	value = os.Getenv(envVar)

	return value, value != ""
}

type AzureClient struct {
	blobClient    *mainStorage.BlobStorageClient
	containerName string
	keyName       string
	leaseID       string
}

func (c *AzureClient) Get() (*Payload, error) {
	blob, err := c.blobClient.GetBlob(c.containerName, c.keyName)
	if err != nil {
		if storErr, ok := err.(mainStorage.AzureStorageServiceError); ok {
			if storErr.Code == "BlobNotFound" {
				return nil, nil
			}
		}
		return nil, err
	}

	defer blob.Close()

	data, err := ioutil.ReadAll(blob)
	if err != nil {
		return nil, err
	}

	payload := &Payload{
		Data: data,
	}

	// If there was no data, then return nil
	if len(payload.Data) == 0 {
		return nil, nil
	}

	return payload, nil
}

func (c *AzureClient) Put(data []byte) error {
	headers := map[string]string{
		"Content-Type": "application/json",
	}

	if c.leaseID != "" {
		headers["x-ms-lease-id"] = c.leaseID
	}

	return c.blobClient.CreateBlockBlobFromReader(
		c.containerName,
		c.keyName,
		uint64(len(data)),
		bytes.NewReader(data),
		headers,
	)
}

func (c *AzureClient) Delete() error {
	headers := map[string]string{}
	if c.leaseID != "" {
		headers["x-ms-lease-id"] = c.leaseID
	}

	return c.blobClient.DeleteBlob(c.containerName, c.keyName, headers)
}
<|endoftext|>


<commit_before>package migu

import (
	"database/sql"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"github.com/naoina/migu/dialect"
)

// Sync synchronizes the schema between Go's struct and the database.
// Go's struct may be provided via the filename of the source file, or via
// the src parameter.
//
// If src != nil, Sync parses the source from src and filename is not used.
// The type of the argument for the src parameter must be string, []byte, or
// io.Reader. If src == nil, Sync parses the file specified by filename.
//
// All query for synchronization will be performed within the transaction if
// storage engine supports the transaction. (e.g. MySQL's MyISAM engine does
// NOT support the transaction)
func Sync(db *sql.DB, filename string, src interface{}) error {
	sqls, err := Diff(db, filename, src)
	if err != nil {
		return err
	}
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	for _, sql := range sqls {
		if _, err := tx.Exec(sql); err != nil {
			tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

// Diff returns SQLs for schema synchronous between database and Go's struct.
func Diff(db *sql.DB, filename string, src interface{}) ([]string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, filename, src, parser.ParseComments)
	if err != nil {
		return nil, err
	}
	ast.FileExports(f)
	structASTMap := map[string]*ast.StructType{}
	ast.Inspect(f, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.TypeSpec:
			if t, ok := x.Type.(*ast.StructType); ok {
				structASTMap[x.Name.Name] = t
			}
			return false
		default:
			return true
		}
	})
	structMap := map[string][]*field{}
	for name, structAST := range structASTMap {
		for _, fld := range structAST.Fields.List {
			var typeName string
			switch t := fld.Type.(type) {
			case *ast.Ident:
				typeName = t.Name
			case *ast.SelectorExpr:
				typeName = t.X.(*ast.Ident).Name + "." + t.Sel.Name
			case *ast.StarExpr:
				typeName = "*" + t.X.(*ast.Ident).Name
			default:
				return nil, fmt.Errorf("migu: BUG: unknown type %T", t)
			}
			f := field{
				Type: typeName,
			}
			if fld.Tag != nil {
				s, err := strconv.Unquote(fld.Tag.Value)
				if err != nil {
					return nil, err
				}
				tag := reflect.StructTag(s)
				if def := tag.Get("default"); def != "" {
					f.Default = formatDefault(f.Type, def)
				}
			}
			if fld.Comment != nil {
				f.Comment = strings.TrimSpace(fld.Comment.Text())
			}
			for _, ident := range fld.Names {
				field := f
				field.Name = ident.Name
				structMap[name] = append(structMap[name], &field)
			}
		}
	}
	tableMap, err := getTableMap(db)
	if err != nil {
		return nil, err
	}
	names := make([]string, 0, len(structMap))
	for name := range structMap {
		names = append(names, name)
	}
	sort.Strings(names)
	d := &dialect.MySQL{}
	var migrations []string
	for _, name := range names {
		model := structMap[name]
		columns, ok := tableMap[name]
		tableName := toSnakeCase(name)
		if !ok {
			columns := make([]string, len(model))
			for i, f := range model {
				colType, null := d.ColumnType(f.Type)
				column := []string{d.Quote(toSnakeCase(f.Name)), colType}
				if !null {
					column = append(column, "NOT NULL")
				}
				if f.Default != "" {
					column = append(column, "DEFAULT", f.Default)
				}
				if f.PrimaryKey {
					column = append(column, "NOT NULL", "AUTO_INCREMENT", "PRIMARY KEY")
				} else if f.Unique {
					column = append(column, "UNIQUE")
				}
				if f.Comment != "" {
					column = append(column, "COMMENT", fmt.Sprintf("'%s'", f.Comment))
				}
				columns[i] = strings.Join(column, " ")
			}
			migrations = append(migrations, fmt.Sprintf(`CREATE TABLE %s (
  %s
)`, d.Quote(tableName), strings.Join(columns, ",\n  ")))
		} else {
			table := map[string]*columnSchema{}
			for _, column := range columns {
				table[toCamelCase(column.ColumnName)] = column
			}
			for _, f := range model {
				if column, ok := table[f.Name]; ok {
					types, err := column.GoFieldTypes()
					if err != nil {
						return nil, err
					}
					if !inStrings(types, f.Type) {
						colType, null := d.ColumnType(f.Type)
						m := fmt.Sprintf(`ALTER TABLE %s MODIFY %s %s`, d.Quote(tableName), d.Quote(toSnakeCase(f.Name)), colType)
						if !null {
							m += " NOT NULL"
						}
						migrations = append(migrations, m)
					}
				} else {
					colType, null := d.ColumnType(f.Type)
					m := fmt.Sprintf(`ALTER TABLE %s ADD %s %s`, d.Quote(tableName), d.Quote(toSnakeCase(f.Name)), colType)
					if !null {
						m += " NOT NULL"
					}
					migrations = append(migrations, m)
				}
				delete(table, f.Name)
			}
			for _, f := range table {
				migrations = append(migrations, fmt.Sprintf(`ALTER TABLE %s DROP %s`, d.Quote(tableName), d.Quote(toSnakeCase(f.ColumnName))))
			}
		}
		delete(structMap, name)
		delete(tableMap, name)
	}
	for name := range tableMap {
		migrations = append(migrations, fmt.Sprintf(`DROP TABLE %s`, d.Quote(toSnakeCase(name))))
	}
	return migrations, nil
}

type field struct {
	Name       string
	Type       string
	Comment    string
	Unique     bool
	PrimaryKey bool
	Default    string
	Size       int64
}

// Fprint generates Go's structs from database schema and writes to output.
func Fprint(output io.Writer, db *sql.DB) error {
	tableMap, err := getTableMap(db)
	if err != nil {
		return err
	}
	names := make([]string, 0, len(tableMap))
	for name := range tableMap {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		for _, schema := range tableMap[name] {
			if schema.DataType == "datetime" {
				if err := fprintln(output, importAST("time")); err != nil {
					return err
				}
				break
			}
		}
	}
	for _, name := range names {
		s, err := structAST(name, tableMap[name])
		if err != nil {
			return err
		}
		if err := fprintln(output, s); err != nil {
			return err
		}
	}
	return nil
}

func getTableMap(db *sql.DB) (map[string][]*columnSchema, error) {
	dbname, err := getCurrentDBName(db)
	if err != nil {
		return nil, err
	}
	query := `
SELECT
  TABLE_NAME,
  COLUMN_NAME,
  COLUMN_DEFAULT,
  IS_NULLABLE,
  DATA_TYPE,
  CHARACTER_MAXIMUM_LENGTH,
  CHARACTER_OCTET_LENGTH,
  NUMERIC_PRECISION,
  NUMERIC_SCALE,
  COLUMN_TYPE,
  COLUMN_KEY,
  COLUMN_COMMENT
FROM information_schema.COLUMNS
WHERE TABLE_SCHEMA = ?
ORDER BY TABLE_NAME, ORDINAL_POSITION`
	rows, err := db.Query(query, dbname)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	tableMap := map[string][]*columnSchema{}
	for rows.Next() {
		schema := &columnSchema{}
		if err := rows.Scan(
			&schema.TableName,
			&schema.ColumnName,
			&schema.ColumnDefault,
			&schema.IsNullable,
			&schema.DataType,
			&schema.CharacterMaximumLength,
			&schema.CharacterOctetLength,
			&schema.NumericPrecision,
			&schema.NumericScale,
			&schema.ColumnType,
			&schema.ColumnKey,
			&schema.ColumnComment,
		); err != nil {
			return nil, err
		}
		tableName := toCamelCase(schema.TableName)
		tableMap[tableName] = append(tableMap[tableName], schema)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return tableMap, nil
}

func getCurrentDBName(db *sql.DB) (string, error) {
	var dbname sql.NullString
	err := db.QueryRow(`SELECT DATABASE()`).Scan(&dbname)
	return dbname.String, err
}

func formatDefault(t, def string) string {
	switch t {
	case "string":
		return `'` + def + `'`
	default:
		return def
	}
}

func fprintln(output io.Writer, decl ast.Decl) error {
	if err := format.Node(output, token.NewFileSet(), decl); err != nil {
		return err
	}
	fmt.Fprintf(output, "\n\n")
	return nil
}

func importAST(pkg string) ast.Decl {
	return &ast.GenDecl{
		Tok: token.IMPORT,
		Specs: []ast.Spec{
			&ast.ImportSpec{
				Path: &ast.BasicLit{
					Kind:  token.STRING,
					Value: fmt.Sprintf(`"%s"`, pkg),
				},
			},
		},
	}
}

func structAST(name string, schemas []*columnSchema) (ast.Decl, error) {
	var fields []*ast.Field
	for _, schema := range schemas {
		f, err := schema.fieldAST()
		if err != nil {
			return nil, err
		}
		fields = append(fields, f)
	}
	return &ast.GenDecl{
		Tok: token.TYPE,
		Specs: []ast.Spec{
			&ast.TypeSpec{
				Name: ast.NewIdent(toCamelCase(name)),
				Type: &ast.StructType{
					Fields: &ast.FieldList{
						List: fields,
					},
				},
			},
		},
	}, nil
}

type columnSchema struct {
	TableName              string
	ColumnName             string
	OrdinalPosition        int64
	ColumnDefault          sql.NullString
	IsNullable             string
	DataType               string
	CharacterMaximumLength sql.NullInt64
	CharacterOctetLength   sql.NullInt64
	NumericPrecision       sql.NullInt64
	NumericScale           sql.NullInt64
	ColumnType             string
	ColumnKey              string
	ColumnComment          string
}

func (schema *columnSchema) fieldAST() (*ast.Field, error) {
	types, err := schema.GoFieldTypes()
	if err != nil {
		return nil, err
	}
	field := &ast.Field{
		Names: []*ast.Ident{
			ast.NewIdent(toCamelCase(schema.ColumnName)),
		},
		Type: ast.NewIdent(types[0]),
	}
	if schema.ColumnDefault.Valid {
		field.Tag = &ast.BasicLit{
			Kind:  token.STRING,
			Value: fmt.Sprintf(`"default:\"%s\""`, schema.ColumnDefault.String),
		}
	}
	return field, nil
}

func (schema *columnSchema) GoFieldTypes() ([]string, error) {
	switch schema.DataType {
	case "tinyint", "smallint", "mediumint", "int":
		if schema.isUnsigned() {
			if schema.isNullable() {
				return []string{"*uint", "sql.NullInt64"}, nil
			}
			return []string{"uint"}, nil
		}
		if schema.isNullable() {
			return []string{"*int", "sql.NullInt64"}, nil
		}
		return []string{"int"}, nil
	case "bigint":
		if schema.isUnsigned() {
			if schema.isNullable() {
				return []string{"*uint64", "sql.NullInt64"}, nil
			}
			return []string{"uint64"}, nil
		}
		if schema.isNullable() {
			return []string{"*int64", "sql.NullInt64"}, nil
		}
		return []string{"int64"}, nil
	case "varchar", "text", "mediumtext", "longtext":
		if schema.isNullable() {
			return []string{"*string", "sql.NullString"}, nil
		}
		return []string{"string"}, nil
	case "datetime":
		if schema.isNullable() {
			return []string{"*time.Time"}, nil
		}
		return []string{"time.Time"}, nil
	default:
		return nil, fmt.Errorf("BUG: unexpected data type: %s", schema.DataType)
	}
}

func (schema *columnSchema) isUnsigned() bool {
	return strings.Contains(schema.ColumnType, "unsigned")
}

func (schema *columnSchema) isNullable() bool {
	return strings.ToUpper(schema.IsNullable) == "YES"
}
<commit_msg>Refactoring<commit_after>package migu

import (
	"database/sql"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"github.com/naoina/migu/dialect"
)

// Sync synchronizes the schema between Go's struct and the database.
// Go's struct may be provided via the filename of the source file, or via
// the src parameter.
//
// If src != nil, Sync parses the source from src and filename is not used.
// The type of the argument for the src parameter must be string, []byte, or
// io.Reader. If src == nil, Sync parses the file specified by filename.
//
// All query for synchronization will be performed within the transaction if
// storage engine supports the transaction. (e.g. MySQL's MyISAM engine does
// NOT support the transaction)
func Sync(db *sql.DB, filename string, src interface{}) error {
	sqls, err := Diff(db, filename, src)
	if err != nil {
		return err
	}
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	for _, sql := range sqls {
		if _, err := tx.Exec(sql); err != nil {
			tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

// Diff returns SQLs for schema synchronous between database and Go's struct.
func Diff(db *sql.DB, filename string, src interface{}) ([]string, error) {
	structASTMap, err := makeStructASTMap(filename, src)
	if err != nil {
		return nil, err
	}
	structMap := map[string][]*field{}
	for name, structAST := range structASTMap {
		for _, fld := range structAST.Fields.List {
			typeName, err := detectTypeName(fld)
			if err != nil {
				return nil, err
			}
			f, err := newField(typeName, fld)
			if err != nil {
				return nil, err
			}
			for _, ident := range fld.Names {
				field := *f
				field.Name = ident.Name
				structMap[name] = append(structMap[name], &field)
			}
		}
	}
	tableMap, err := getTableMap(db)
	if err != nil {
		return nil, err
	}
	names := make([]string, 0, len(structMap))
	for name := range structMap {
		names = append(names, name)
	}
	sort.Strings(names)
	d := &dialect.MySQL{}
	var migrations []string
	for _, name := range names {
		model := structMap[name]
		columns, ok := tableMap[name]
		tableName := toSnakeCase(name)
		if !ok {
			columns := make([]string, len(model))
			for i, f := range model {
				columns[i] = columnSQL(d, f)
			}
			migrations = append(migrations, fmt.Sprintf(`CREATE TABLE %s (
  %s
)`, d.Quote(tableName), strings.Join(columns, ",\n  ")))
		} else {
			table := map[string]*columnSchema{}
			for _, column := range columns {
				table[toCamelCase(column.ColumnName)] = column
			}
			for _, f := range model {
				s, err := alterTableSQL(d, tableName, table, f)
				if err != nil {
					return nil, err
				}
				if s != "" {
					migrations = append(migrations, s)
				}
				delete(table, f.Name)
			}
			for _, f := range table {
				migrations = append(migrations, fmt.Sprintf(`ALTER TABLE %s DROP %s`, d.Quote(tableName), d.Quote(toSnakeCase(f.ColumnName))))
			}
		}
		delete(structMap, name)
		delete(tableMap, name)
	}
	for name := range tableMap {
		migrations = append(migrations, fmt.Sprintf(`DROP TABLE %s`, d.Quote(toSnakeCase(name))))
	}
	return migrations, nil
}

type field struct {
	Name       string
	Type       string
	Comment    string
	Unique     bool
	PrimaryKey bool
	Default    string
	Size       int64
}

func newField(name string, f *ast.Field) (*field, error) {
	ret := &field{
		Type: name,
	}
	if f.Tag != nil {
		s, err := strconv.Unquote(f.Tag.Value)
		if err != nil {
			return nil, err
		}
		tag := reflect.StructTag(s)
		if def := tag.Get("default"); def != "" {
			ret.Default = formatDefault(ret.Type, def)
		}
	}
	if f.Comment != nil {
		ret.Comment = strings.TrimSpace(f.Comment.Text())
	}
	return ret, nil
}

// Fprint generates Go's structs from database schema and writes to output.
func Fprint(output io.Writer, db *sql.DB) error {
	tableMap, err := getTableMap(db)
	if err != nil {
		return err
	}
	names := make([]string, 0, len(tableMap))
	for name := range tableMap {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		for _, schema := range tableMap[name] {
			if schema.DataType == "datetime" {
				if err := fprintln(output, importAST("time")); err != nil {
					return err
				}
				break
			}
		}
	}
	for _, name := range names {
		s, err := structAST(name, tableMap[name])
		if err != nil {
			return err
		}
		if err := fprintln(output, s); err != nil {
			return err
		}
	}
	return nil
}

func getTableMap(db *sql.DB) (map[string][]*columnSchema, error) {
	dbname, err := getCurrentDBName(db)
	if err != nil {
		return nil, err
	}
	query := `
SELECT
  TABLE_NAME,
  COLUMN_NAME,
  COLUMN_DEFAULT,
  IS_NULLABLE,
  DATA_TYPE,
  CHARACTER_MAXIMUM_LENGTH,
  CHARACTER_OCTET_LENGTH,
  NUMERIC_PRECISION,
  NUMERIC_SCALE,
  COLUMN_TYPE,
  COLUMN_KEY,
  COLUMN_COMMENT
FROM information_schema.COLUMNS
WHERE TABLE_SCHEMA = ?
ORDER BY TABLE_NAME, ORDINAL_POSITION`
	rows, err := db.Query(query, dbname)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	tableMap := map[string][]*columnSchema{}
	for rows.Next() {
		schema := &columnSchema{}
		if err := rows.Scan(
			&schema.TableName,
			&schema.ColumnName,
			&schema.ColumnDefault,
			&schema.IsNullable,
			&schema.DataType,
			&schema.CharacterMaximumLength,
			&schema.CharacterOctetLength,
			&schema.NumericPrecision,
			&schema.NumericScale,
			&schema.ColumnType,
			&schema.ColumnKey,
			&schema.ColumnComment,
		); err != nil {
			return nil, err
		}
		tableName := toCamelCase(schema.TableName)
		tableMap[tableName] = append(tableMap[tableName], schema)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return tableMap, nil
}

func getCurrentDBName(db *sql.DB) (string, error) {
	var dbname sql.NullString
	err := db.QueryRow(`SELECT DATABASE()`).Scan(&dbname)
	return dbname.String, err
}

func formatDefault(t, def string) string {
	switch t {
	case "string":
		return `'` + def + `'`
	default:
		return def
	}
}

func fprintln(output io.Writer, decl ast.Decl) error {
	if err := format.Node(output, token.NewFileSet(), decl); err != nil {
		return err
	}
	fmt.Fprintf(output, "\n\n")
	return nil
}

func makeStructASTMap(filename string, src interface{}) (map[string]*ast.StructType, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, filename, src, parser.ParseComments)
	if err != nil {
		return nil, err
	}
	ast.FileExports(f)
	structASTMap := map[string]*ast.StructType{}
	ast.Inspect(f, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.TypeSpec:
			if t, ok := x.Type.(*ast.StructType); ok {
				structASTMap[x.Name.Name] = t
			}
			return false
		default:
			return true
		}
	})
	return structASTMap, nil
}

func detectTypeName(f *ast.Field) (string, error) {
	switch t := f.Type.(type) {
	case *ast.Ident:
		return t.Name, nil
	case *ast.SelectorExpr:
		return t.X.(*ast.Ident).Name + "." + t.Sel.Name, nil
	case *ast.StarExpr:
		return "*" + t.X.(*ast.Ident).Name, nil
	default:
		return "", fmt.Errorf("migu: BUG: unknown type %T", t)
	}
}

func columnSQL(d dialect.Dialect, f *field) string {
	colType, null := d.ColumnType(f.Type)
	column := []string{d.Quote(toSnakeCase(f.Name)), colType}
	if !null {
		column = append(column, "NOT NULL")
	}
	if f.Default != "" {
		column = append(column, "DEFAULT", f.Default)
	}
	if f.PrimaryKey {
		column = append(column, "NOT NULL", "AUTO_INCREMENT", "PRIMARY KEY")
	} else if f.Unique {
		column = append(column, "UNIQUE")
	}
	if f.Comment != "" {
		column = append(column, "COMMENT", fmt.Sprintf("'%s'", f.Comment))
	}
	return strings.Join(column, " ")
}

func alterTableSQL(d dialect.Dialect, tableName string, table map[string]*columnSchema, f *field) (string, error) {
	column, exists := table[f.Name]
	if !exists {
		colType, null := d.ColumnType(f.Type)
		s := fmt.Sprintf(`ALTER TABLE %s ADD %s %s`, d.Quote(tableName), d.Quote(toSnakeCase(f.Name)), colType)
		if !null {
			s += " NOT NULL"
		}
		return s, nil
	}
	types, err := column.GoFieldTypes()
	if err != nil {
		return "", err
	}
	if !inStrings(types, f.Type) {
		colType, null := d.ColumnType(f.Type)
		s := fmt.Sprintf(`ALTER TABLE %s MODIFY %s %s`, d.Quote(tableName), d.Quote(toSnakeCase(f.Name)), colType)
		if !null {
			s += " NOT NULL"
		}
		return s, nil
	}
	return "", nil
}

func importAST(pkg string) ast.Decl {
	return &ast.GenDecl{
		Tok: token.IMPORT,
		Specs: []ast.Spec{
			&ast.ImportSpec{
				Path: &ast.BasicLit{
					Kind:  token.STRING,
					Value: fmt.Sprintf(`"%s"`, pkg),
				},
			},
		},
	}
}

func structAST(name string, schemas []*columnSchema) (ast.Decl, error) {
	var fields []*ast.Field
	for _, schema := range schemas {
		f, err := schema.fieldAST()
		if err != nil {
			return nil, err
		}
		fields = append(fields, f)
	}
	return &ast.GenDecl{
		Tok: token.TYPE,
		Specs: []ast.Spec{
			&ast.TypeSpec{
				Name: ast.NewIdent(toCamelCase(name)),
				Type: &ast.StructType{
					Fields: &ast.FieldList{
						List: fields,
					},
				},
			},
		},
	}, nil
}

type columnSchema struct {
	TableName              string
	ColumnName             string
	OrdinalPosition        int64
	ColumnDefault          sql.NullString
	IsNullable             string
	DataType               string
	CharacterMaximumLength sql.NullInt64
	CharacterOctetLength   sql.NullInt64
	NumericPrecision       sql.NullInt64
	NumericScale           sql.NullInt64
	ColumnType             string
	ColumnKey              string
	ColumnComment          string
}

func (schema *columnSchema) fieldAST() (*ast.Field, error) {
	types, err := schema.GoFieldTypes()
	if err != nil {
		return nil, err
	}
	field := &ast.Field{
		Names: []*ast.Ident{
			ast.NewIdent(toCamelCase(schema.ColumnName)),
		},
		Type: ast.NewIdent(types[0]),
	}
	if schema.ColumnDefault.Valid {
		field.Tag = &ast.BasicLit{
			Kind:  token.STRING,
			Value: fmt.Sprintf(`"default:\"%s\""`, schema.ColumnDefault.String),
		}
	}
	return field, nil
}

func (schema *columnSchema) GoFieldTypes() ([]string, error) {
	switch schema.DataType {
	case "tinyint", "smallint", "mediumint", "int":
		if schema.isUnsigned() {
			if schema.isNullable() {
				return []string{"*uint", "sql.NullInt64"}, nil
			}
			return []string{"uint"}, nil
		}
		if schema.isNullable() {
			return []string{"*int", "sql.NullInt64"}, nil
		}
		return []string{"int"}, nil
	case "bigint":
		if schema.isUnsigned() {
			if schema.isNullable() {
				return []string{"*uint64", "sql.NullInt64"}, nil
			}
			return []string{"uint64"}, nil
		}
		if schema.isNullable() {
			return []string{"*int64", "sql.NullInt64"}, nil
		}
		return []string{"int64"}, nil
	case "varchar", "text", "mediumtext", "longtext":
		if schema.isNullable() {
			return []string{"*string", "sql.NullString"}, nil
		}
		return []string{"string"}, nil
	case "datetime":
		if schema.isNullable() {
			return []string{"*time.Time"}, nil
		}
		return []string{"time.Time"}, nil
	default:
		return nil, fmt.Errorf("BUG: unexpected data type: %s", schema.DataType)
	}
}

func (schema *columnSchema) isUnsigned() bool {
	return strings.Contains(schema.ColumnType, "unsigned")
}

func (schema *columnSchema) isNullable() bool {
	return strings.ToUpper(schema.IsNullable) == "YES"
}
<|endoftext|>


<commit_before>package capzcrs

import (
	"context"
	"fmt"
	"reflect"
	"strconv"

	"github.com/giantswarm/apiextensions/v2/pkg/annotation"
	providerv1alpha1 "github.com/giantswarm/apiextensions/v2/pkg/apis/provider/v1alpha1"
	"github.com/giantswarm/apiextensions/v2/pkg/label"
	"github.com/giantswarm/microerror"
	"github.com/giantswarm/to"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	capzv1alpha3 "sigs.k8s.io/cluster-api-provider-azure/api/v1alpha3"
	capiv1alpha3 "sigs.k8s.io/cluster-api/api/v1alpha3"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/giantswarm/azure-operator/v5/service/controller/key"
)

type crmapping struct {
	obj runtime.Object

	needUpdateFunc func(orig, desired runtime.Object) (bool, error)
	mergeFunc      func(orig, desired runtime.Object) (runtime.Object, error)
}

func (r *Resource) EnsureCreated(ctx context.Context, obj interface{}) error {
	cr, err := key.ToCustomResource(obj)
	if err != nil {
		return microerror.Mask(err)
	}

	r.logger.LogCtx(ctx, "level", "debug", "message", "mapping AzureConfig CR to CAPI & CAPZ CRs")

	{
		objKey := client.ObjectKey{
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
		}
		cluster := new(capiv1alpha3.Cluster)
		err = r.ctrlClient.Get(ctx, objKey, cluster)
		if errors.IsNotFound(err) {
			// all good
		} else if err != nil {
			return microerror.Mask(err)
		} else if !cluster.GetDeletionTimestamp().IsZero() {
			r.logger.LogCtx(ctx, "level", "debug", "message", "Cluster is being deleted, skipping mapping AzureConfig CR to CAPI & CAPZ CRs")
			return nil
		}
	}

	var mappedCRs []crmapping
	{
		o, err := r.mapAzureConfigToAzureCluster(ctx, cr)
		if err != nil {
			return microerror.Mask(err)
		}
		mappedCRs = append(mappedCRs, crmapping{
			obj:            o,
			needUpdateFunc: detectAzureClusterUpdate,
			mergeFunc:      mergeAzureCluster,
		})

		infraRef := &corev1.ObjectReference{
			Kind:      "AzureCluster",
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
		}
		o, err = r.mapAzureConfigToCluster(ctx, cr, infraRef)
		if err != nil {
			return microerror.Mask(err)
		}
		mappedCRs = append(mappedCRs, crmapping{
			obj:            o,
			needUpdateFunc: nil,
			mergeFunc:      nil,
		})

		o, err = r.mapAzureConfigToAzureMachine(ctx, cr)
		if err != nil {
			return microerror.Mask(err)
		}
		mappedCRs = append(mappedCRs, crmapping{
			obj:            o,
			needUpdateFunc: nil,
			mergeFunc:      nil,
		})
	}

	err = r.updateCRs(ctx, mappedCRs)
	if err != nil {
		return microerror.Mask(err)
	}

	r.logger.LogCtx(ctx, "level", "debug", "message", "mapped AzureConfig CR to CAPI & CAPZ CRs")

	return nil
}

func (r *Resource) mapAzureConfigToCluster(ctx context.Context, cr providerv1alpha1.AzureConfig, infraRef *corev1.ObjectReference) (runtime.Object, error) {
	cluster := &capiv1alpha3.Cluster{
		TypeMeta: metav1.TypeMeta{
			APIVersion: capiv1alpha3.GroupVersion.String(),
			Kind:       "Cluster",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
			Labels: map[string]string{
				// XXX: azure-operator reconciles Cluster & MachinePool to set OwnerReferences (for now).
				label.AzureOperatorVersion:    key.OperatorVersion(&cr),
				label.ClusterOperatorVersion:  cr.Labels[label.ClusterOperatorVersion],
				label.Cluster:                 cr.Name,
				capiv1alpha3.ClusterLabelName: cr.Name,
				label.Organization:            key.OrganizationID(&cr),
				label.ReleaseVersion:          key.ReleaseVersion(&cr),
			},
		},
		Spec: capiv1alpha3.ClusterSpec{
			ClusterNetwork: &capiv1alpha3.ClusterNetwork{
				APIServerPort: toInt32P(int32(key.APISecurePort(cr))),
				Services: &capiv1alpha3.NetworkRanges{
					CIDRBlocks: []string{
						key.ClusterIPRange(cr),
					},
				},
				ServiceDomain: key.ClusterBaseDomain(cr),
			},
			ControlPlaneEndpoint: capiv1alpha3.APIEndpoint{
				Host: key.ClusterAPIEndpoint(cr),
				Port: 443,
			},
			InfrastructureRef: infraRef,
		},
	}
	return cluster, nil
}

func (r *Resource) mapAzureConfigToAzureCluster(ctx context.Context, cr providerv1alpha1.AzureConfig) (runtime.Object, error) {
	azureCluster := &capzv1alpha3.AzureCluster{
		TypeMeta: metav1.TypeMeta{
			APIVersion: capzv1alpha3.GroupVersion.String(),
			Kind:       "AzureCluster",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
			Labels: map[string]string{
				label.AzureOperatorVersion:    key.OperatorVersion(&cr),
				label.Cluster:                 cr.Name,
				capiv1alpha3.ClusterLabelName: cr.Name,
				label.Organization:            key.OrganizationID(&cr),
				label.ReleaseVersion:          key.ReleaseVersion(&cr),
			},
			Annotations: map[string]string{
				annotation.ClusterDescription: cr.Annotations[annotation.ClusterDescription],
			},
		},
		Spec: capzv1alpha3.AzureClusterSpec{
			Location: r.location,
			ControlPlaneEndpoint: capiv1alpha3.APIEndpoint{
				Host: key.ClusterAPIEndpoint(cr),
				Port: 443,
			},
			ResourceGroup: key.ClusterID(&cr),
			NetworkSpec: capzv1alpha3.NetworkSpec{
				Vnet: capzv1alpha3.VnetSpec{
					CidrBlock:     key.VnetCIDR(cr),
					Name:          key.VnetName(cr),
					ResourceGroup: key.ResourceGroupName(cr),
				},
			},
		},
	}

	return azureCluster, nil
}

func (r *Resource) mapAzureConfigToAzureMachine(ctx context.Context, cr providerv1alpha1.AzureConfig) (runtime.Object, error) {
	if len(cr.Spec.Azure.Masters) == 0 {
		return nil, microerror.Maskf(invalidConfigError, "no master nodes defined")
	}
	vmSize := cr.Spec.Azure.Masters[0].VMSize
	azureMachine := &capzv1alpha3.AzureMachine{
		TypeMeta: metav1.TypeMeta{
			APIVersion: capzv1alpha3.GroupVersion.String(),
			Kind:       "AzureMachine",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      fmt.Sprintf("%s-master-0", key.ClusterID(&cr)),
			Namespace: key.OrganizationNamespace(&cr),
			Labels: map[string]string{
				label.AzureOperatorVersion:                key.OperatorVersion(&cr),
				label.Cluster:                             cr.Name,
				capiv1alpha3.ClusterLabelName:             cr.Name,
				capiv1alpha3.MachineControlPlaneLabelName: "true",
				label.Organization:                        key.OrganizationID(&cr),
				label.ReleaseVersion:                      key.ReleaseVersion(&cr),
			},
		},
		Spec: capzv1alpha3.AzureMachineSpec{
			VMSize: vmSize,
			Image: &capzv1alpha3.Image{
				Marketplace: &capzv1alpha3.AzureMarketplaceImage{
					Publisher: "kinvolk",
					Offer:     "flatcar-container-linux-free",
					SKU:       "stable",
					Version:   "2345.3.1",
				},
			},
			OSDisk: capzv1alpha3.OSDisk{
				OSType:     "Linux",
				DiskSizeGB: int32(50),
				ManagedDisk: capzv1alpha3.ManagedDisk{
					StorageAccountType: "Premium_LRS",
				},
			},
			Location: r.location,
			// We use ignition for SSH keys deployment.
			SSHPublicKey: "",
		},
	}

	if len(key.AvailabilityZones(cr, r.location)) > 0 {
		azureMachine.Spec.FailureDomain = to.StringP(strconv.Itoa(key.AvailabilityZones(cr, r.location)[0]))
	}

	return azureMachine, nil
}

func (r *Resource) updateCRs(ctx context.Context, crmappings []crmapping) error {
	for _, m := range crmappings {
		// Construct new instance via reflection to ensure clean zero value object.
		readCR := reflect.New(reflect.TypeOf(m.obj).Elem()).Interface().(runtime.Object)

		// Acquire accessors for ObjectMeta and TypeMeta fields of CR.
		desiredMeta, err := meta.Accessor(m.obj)
		if err != nil {
			return microerror.Mask(err)
		}
		desiredType, err := meta.TypeAccessor(m.obj)
		if err != nil {
			return microerror.Mask(err)
		}

		nsName := types.NamespacedName{
			Name:      desiredMeta.GetName(),
			Namespace: desiredMeta.GetNamespace(),
		}

		r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("reading present %s %s", desiredType.GetKind(), nsName.String()))

		err = r.ctrlClient.Get(ctx, nsName, readCR)
		if errors.IsNotFound(err) {
			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("%s %s did not exist. creating", desiredType.GetKind(), nsName.String()))

			// It's ok. Let's create it.
			err = r.ctrlClient.Create(ctx, m.obj)
			if err != nil {
				return microerror.Mask(err)
			}

			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("created %s %s", desiredType.GetKind(), nsName.String()))
			continue
		} else if err != nil {
			return microerror.Mask(err)
		}

		// If needUpdateFunc or mergeFunc are nil, it means that given CR is
		// not intended to be updated. Only created when it doesn't already
		// exist.
		if m.needUpdateFunc == nil || m.mergeFunc == nil {
			return nil
		}

		updateNeeded, err := m.needUpdateFunc(readCR, m.obj)
		if err != nil {
			return microerror.Mask(err)
		}

		if updateNeeded {
			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("found that %s %s needs updating", desiredType.GetKind(), nsName.String()))

			merged, err := m.mergeFunc(readCR, m.obj)
			if err != nil {
				return microerror.Mask(err)
			}

			readMeta, err := meta.Accessor(readCR)
			if err != nil {
				return microerror.Mask(err)
			}
			mergedMeta, err := meta.Accessor(merged)
			if err != nil {
				return microerror.Mask(err)
			}

			// Copy read CR's resource version to update object for optimistic
			// locking.
			mergedMeta.SetResourceVersion(readMeta.GetResourceVersion())

			err = r.ctrlClient.Update(ctx, merged)
			if err != nil {
				return microerror.Mask(err)
			}

			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("updated %s %s", desiredType.GetKind(), nsName.String()))
		} else {
			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("no need to %s %s", desiredType.GetKind(), nsName.String()))
		}
	}

	return nil
}

func detectAzureClusterUpdate(orig, desired runtime.Object) (bool, error) {
	o := orig.(*capzv1alpha3.AzureCluster)
	d := desired.(*capzv1alpha3.AzureCluster)

	return (o.Spec.NetworkSpec.Vnet.CidrBlock != d.Spec.NetworkSpec.Vnet.CidrBlock ||
		o.Spec.NetworkSpec.Vnet.Name != d.Spec.NetworkSpec.Vnet.Name ||
		o.Spec.NetworkSpec.Vnet.ResourceGroup != d.Spec.NetworkSpec.Vnet.ResourceGroup ||
		o.Spec.ResourceGroup != d.Spec.ResourceGroup ||
		o.Spec.Location != d.Spec.Location ||
		o.Spec.ControlPlaneEndpoint != d.Spec.ControlPlaneEndpoint), nil
}

func mergeAzureCluster(orig, desired runtime.Object) (runtime.Object, error) {
	o := orig.(*capzv1alpha3.AzureCluster)
	d := desired.(*capzv1alpha3.AzureCluster)

	// Only copy specific parts of desired opbject
	o.Spec.NetworkSpec.Vnet.CidrBlock = d.Spec.NetworkSpec.Vnet.CidrBlock
	o.Spec.NetworkSpec.Vnet.Name = d.Spec.NetworkSpec.Vnet.Name
	o.Spec.NetworkSpec.Vnet.ResourceGroup = d.Spec.NetworkSpec.Vnet.ResourceGroup
	o.Spec.ResourceGroup = d.Spec.ResourceGroup
	o.Spec.Location = d.Spec.Location
	o.Spec.ControlPlaneEndpoint = d.Spec.ControlPlaneEndpoint

	return o, nil
}

func toInt32P(v int32) *int32 {
	return &v
}
<commit_msg>Rewrite object instantiation with comments<commit_after>package capzcrs

import (
	"context"
	"fmt"
	"reflect"
	"strconv"

	"github.com/giantswarm/apiextensions/v2/pkg/annotation"
	providerv1alpha1 "github.com/giantswarm/apiextensions/v2/pkg/apis/provider/v1alpha1"
	"github.com/giantswarm/apiextensions/v2/pkg/label"
	"github.com/giantswarm/microerror"
	"github.com/giantswarm/to"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	capzv1alpha3 "sigs.k8s.io/cluster-api-provider-azure/api/v1alpha3"
	capiv1alpha3 "sigs.k8s.io/cluster-api/api/v1alpha3"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/giantswarm/azure-operator/v5/service/controller/key"
)

type crmapping struct {
	obj runtime.Object

	needUpdateFunc func(orig, desired runtime.Object) (bool, error)
	mergeFunc      func(orig, desired runtime.Object) (runtime.Object, error)
}

func (r *Resource) EnsureCreated(ctx context.Context, obj interface{}) error {
	cr, err := key.ToCustomResource(obj)
	if err != nil {
		return microerror.Mask(err)
	}

	r.logger.LogCtx(ctx, "level", "debug", "message", "mapping AzureConfig CR to CAPI & CAPZ CRs")

	{
		objKey := client.ObjectKey{
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
		}
		cluster := new(capiv1alpha3.Cluster)
		err = r.ctrlClient.Get(ctx, objKey, cluster)
		if errors.IsNotFound(err) {
			// all good
		} else if err != nil {
			return microerror.Mask(err)
		} else if !cluster.GetDeletionTimestamp().IsZero() {
			r.logger.LogCtx(ctx, "level", "debug", "message", "Cluster is being deleted, skipping mapping AzureConfig CR to CAPI & CAPZ CRs")
			return nil
		}
	}

	var mappedCRs []crmapping
	{
		o, err := r.mapAzureConfigToAzureCluster(ctx, cr)
		if err != nil {
			return microerror.Mask(err)
		}
		mappedCRs = append(mappedCRs, crmapping{
			obj:            o,
			needUpdateFunc: detectAzureClusterUpdate,
			mergeFunc:      mergeAzureCluster,
		})

		infraRef := &corev1.ObjectReference{
			Kind:      "AzureCluster",
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
		}
		o, err = r.mapAzureConfigToCluster(ctx, cr, infraRef)
		if err != nil {
			return microerror.Mask(err)
		}
		mappedCRs = append(mappedCRs, crmapping{
			obj:            o,
			needUpdateFunc: nil,
			mergeFunc:      nil,
		})

		o, err = r.mapAzureConfigToAzureMachine(ctx, cr)
		if err != nil {
			return microerror.Mask(err)
		}
		mappedCRs = append(mappedCRs, crmapping{
			obj:            o,
			needUpdateFunc: nil,
			mergeFunc:      nil,
		})
	}

	err = r.updateCRs(ctx, mappedCRs)
	if err != nil {
		return microerror.Mask(err)
	}

	r.logger.LogCtx(ctx, "level", "debug", "message", "mapped AzureConfig CR to CAPI & CAPZ CRs")

	return nil
}

func (r *Resource) mapAzureConfigToCluster(ctx context.Context, cr providerv1alpha1.AzureConfig, infraRef *corev1.ObjectReference) (runtime.Object, error) {
	cluster := &capiv1alpha3.Cluster{
		TypeMeta: metav1.TypeMeta{
			APIVersion: capiv1alpha3.GroupVersion.String(),
			Kind:       "Cluster",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
			Labels: map[string]string{
				// XXX: azure-operator reconciles Cluster & MachinePool to set OwnerReferences (for now).
				label.AzureOperatorVersion:    key.OperatorVersion(&cr),
				label.ClusterOperatorVersion:  cr.Labels[label.ClusterOperatorVersion],
				label.Cluster:                 cr.Name,
				capiv1alpha3.ClusterLabelName: cr.Name,
				label.Organization:            key.OrganizationID(&cr),
				label.ReleaseVersion:          key.ReleaseVersion(&cr),
			},
		},
		Spec: capiv1alpha3.ClusterSpec{
			ClusterNetwork: &capiv1alpha3.ClusterNetwork{
				APIServerPort: toInt32P(int32(key.APISecurePort(cr))),
				Services: &capiv1alpha3.NetworkRanges{
					CIDRBlocks: []string{
						key.ClusterIPRange(cr),
					},
				},
				ServiceDomain: key.ClusterBaseDomain(cr),
			},
			ControlPlaneEndpoint: capiv1alpha3.APIEndpoint{
				Host: key.ClusterAPIEndpoint(cr),
				Port: 443,
			},
			InfrastructureRef: infraRef,
		},
	}
	return cluster, nil
}

func (r *Resource) mapAzureConfigToAzureCluster(ctx context.Context, cr providerv1alpha1.AzureConfig) (runtime.Object, error) {
	azureCluster := &capzv1alpha3.AzureCluster{
		TypeMeta: metav1.TypeMeta{
			APIVersion: capzv1alpha3.GroupVersion.String(),
			Kind:       "AzureCluster",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      cr.Name,
			Namespace: key.OrganizationNamespace(&cr),
			Labels: map[string]string{
				label.AzureOperatorVersion:    key.OperatorVersion(&cr),
				label.Cluster:                 cr.Name,
				capiv1alpha3.ClusterLabelName: cr.Name,
				label.Organization:            key.OrganizationID(&cr),
				label.ReleaseVersion:          key.ReleaseVersion(&cr),
			},
			Annotations: map[string]string{
				annotation.ClusterDescription: cr.Annotations[annotation.ClusterDescription],
			},
		},
		Spec: capzv1alpha3.AzureClusterSpec{
			Location: r.location,
			ControlPlaneEndpoint: capiv1alpha3.APIEndpoint{
				Host: key.ClusterAPIEndpoint(cr),
				Port: 443,
			},
			ResourceGroup: key.ClusterID(&cr),
			NetworkSpec: capzv1alpha3.NetworkSpec{
				Vnet: capzv1alpha3.VnetSpec{
					CidrBlock:     key.VnetCIDR(cr),
					Name:          key.VnetName(cr),
					ResourceGroup: key.ResourceGroupName(cr),
				},
			},
		},
	}

	return azureCluster, nil
}

func (r *Resource) mapAzureConfigToAzureMachine(ctx context.Context, cr providerv1alpha1.AzureConfig) (runtime.Object, error) {
	if len(cr.Spec.Azure.Masters) == 0 {
		return nil, microerror.Maskf(invalidConfigError, "no master nodes defined")
	}
	vmSize := cr.Spec.Azure.Masters[0].VMSize
	azureMachine := &capzv1alpha3.AzureMachine{
		TypeMeta: metav1.TypeMeta{
			APIVersion: capzv1alpha3.GroupVersion.String(),
			Kind:       "AzureMachine",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      fmt.Sprintf("%s-master-0", key.ClusterID(&cr)),
			Namespace: key.OrganizationNamespace(&cr),
			Labels: map[string]string{
				label.AzureOperatorVersion:                key.OperatorVersion(&cr),
				label.Cluster:                             cr.Name,
				capiv1alpha3.ClusterLabelName:             cr.Name,
				capiv1alpha3.MachineControlPlaneLabelName: "true",
				label.Organization:                        key.OrganizationID(&cr),
				label.ReleaseVersion:                      key.ReleaseVersion(&cr),
			},
		},
		Spec: capzv1alpha3.AzureMachineSpec{
			VMSize: vmSize,
			Image: &capzv1alpha3.Image{
				Marketplace: &capzv1alpha3.AzureMarketplaceImage{
					Publisher: "kinvolk",
					Offer:     "flatcar-container-linux-free",
					SKU:       "stable",
					Version:   "2345.3.1",
				},
			},
			OSDisk: capzv1alpha3.OSDisk{
				OSType:     "Linux",
				DiskSizeGB: int32(50),
				ManagedDisk: capzv1alpha3.ManagedDisk{
					StorageAccountType: "Premium_LRS",
				},
			},
			Location: r.location,
			// We use ignition for SSH keys deployment.
			SSHPublicKey: "",
		},
	}

	if len(key.AvailabilityZones(cr, r.location)) > 0 {
		azureMachine.Spec.FailureDomain = to.StringP(strconv.Itoa(key.AvailabilityZones(cr, r.location)[0]))
	}

	return azureMachine, nil
}

func (r *Resource) updateCRs(ctx context.Context, crmappings []crmapping) error {
	for _, m := range crmappings {
		// Construct new instance via reflection to ensure clean zero value object.
		var readCR runtime.Object
		{
			// Get the underlying type of desired runtime.Object.
			t := reflect.TypeOf(m.obj)

			// We know that underlying type is a pointer so let's dereference
			// it before cloning so that we get the actual object instance
			// instead of just an instance of a pointer.
			e := t.Elem()

			// Construct a new instance of that type and receive
			// `reflect.Value` object containing pointer to that instance.
			v := reflect.New(e)

			// Finally, extract the encapsulated `interface{}` from the
			// `reflect.Value` and cast it to instance of `runtime.Object`
			// interface.
			readCR = v.Interface().(runtime.Object)
		}

		// Acquire accessors for ObjectMeta and TypeMeta fields of CR.
		desiredMeta, err := meta.Accessor(m.obj)
		if err != nil {
			return microerror.Mask(err)
		}
		desiredType, err := meta.TypeAccessor(m.obj)
		if err != nil {
			return microerror.Mask(err)
		}

		nsName := types.NamespacedName{
			Name:      desiredMeta.GetName(),
			Namespace: desiredMeta.GetNamespace(),
		}

		r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("reading present %s %s", desiredType.GetKind(), nsName.String()))

		err = r.ctrlClient.Get(ctx, nsName, readCR)
		if errors.IsNotFound(err) {
			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("%s %s did not exist. creating", desiredType.GetKind(), nsName.String()))

			// It's ok. Let's create it.
			err = r.ctrlClient.Create(ctx, m.obj)
			if err != nil {
				return microerror.Mask(err)
			}

			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("created %s %s", desiredType.GetKind(), nsName.String()))
			continue
		} else if err != nil {
			return microerror.Mask(err)
		}

		// If needUpdateFunc or mergeFunc are nil, it means that given CR is
		// not intended to be updated. Only created when it doesn't already
		// exist.
		if m.needUpdateFunc == nil || m.mergeFunc == nil {
			return nil
		}

		updateNeeded, err := m.needUpdateFunc(readCR, m.obj)
		if err != nil {
			return microerror.Mask(err)
		}

		if updateNeeded {
			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("found that %s %s needs updating", desiredType.GetKind(), nsName.String()))

			merged, err := m.mergeFunc(readCR, m.obj)
			if err != nil {
				return microerror.Mask(err)
			}

			readMeta, err := meta.Accessor(readCR)
			if err != nil {
				return microerror.Mask(err)
			}
			mergedMeta, err := meta.Accessor(merged)
			if err != nil {
				return microerror.Mask(err)
			}

			// Copy read CR's resource version to update object for optimistic
			// locking.
			mergedMeta.SetResourceVersion(readMeta.GetResourceVersion())

			err = r.ctrlClient.Update(ctx, merged)
			if err != nil {
				return microerror.Mask(err)
			}

			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("updated %s %s", desiredType.GetKind(), nsName.String()))
		} else {
			r.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("no need to %s %s", desiredType.GetKind(), nsName.String()))
		}
	}

	return nil
}

func detectAzureClusterUpdate(orig, desired runtime.Object) (bool, error) {
	o := orig.(*capzv1alpha3.AzureCluster)
	d := desired.(*capzv1alpha3.AzureCluster)

	return (o.Spec.NetworkSpec.Vnet.CidrBlock != d.Spec.NetworkSpec.Vnet.CidrBlock ||
		o.Spec.NetworkSpec.Vnet.Name != d.Spec.NetworkSpec.Vnet.Name ||
		o.Spec.NetworkSpec.Vnet.ResourceGroup != d.Spec.NetworkSpec.Vnet.ResourceGroup ||
		o.Spec.ResourceGroup != d.Spec.ResourceGroup ||
		o.Spec.Location != d.Spec.Location ||
		o.Spec.ControlPlaneEndpoint != d.Spec.ControlPlaneEndpoint), nil
}

func mergeAzureCluster(orig, desired runtime.Object) (runtime.Object, error) {
	o := orig.(*capzv1alpha3.AzureCluster)
	d := desired.(*capzv1alpha3.AzureCluster)

	// Only copy specific parts of desired opbject
	o.Spec.NetworkSpec.Vnet.CidrBlock = d.Spec.NetworkSpec.Vnet.CidrBlock
	o.Spec.NetworkSpec.Vnet.Name = d.Spec.NetworkSpec.Vnet.Name
	o.Spec.NetworkSpec.Vnet.ResourceGroup = d.Spec.NetworkSpec.Vnet.ResourceGroup
	o.Spec.ResourceGroup = d.Spec.ResourceGroup
	o.Spec.Location = d.Spec.Location
	o.Spec.ControlPlaneEndpoint = d.Spec.ControlPlaneEndpoint

	return o, nil
}

func toInt32P(v int32) *int32 {
	return &v
}
<|endoftext|>