<commit_before>package glmenu

import (
	"github.com/4ydx/gltext"
	"github.com/go-gl/gl/v3.3-core/gl"
	"github.com/go-gl/glfw/v3.1/glfw"
	"github.com/go-gl/mathgl/mgl32"
	"golang.org/x/image/math/fixed"
	"os"
)

type Point struct {
	X, Y float32
}

var vertexShaderSource string = `
#version 330

uniform mat4 matrix;

in vec4 position;

void main() {
  gl_Position = matrix * position;
}
` + "\x00"

var fragmentShaderSource string = `
#version 330

uniform vec4 background;
out vec4 fragment_color;

void main() {
  fragment_color = background;
}
` + "\x00"

type MouseClick int

const (
	MouseUnclicked MouseClick = iota
	MouseLeft
	MouseRight
	MouseCenter
)

type Menu struct {
	//trigger
	OnShow         func()
	OnEnterRelease func()

	// options
	Visible      bool
	ShowOnKey    glfw.Key
	Height       float32
	Width        float32
	IsAutoCenter bool
	lowerLeft    Point

	backgroundUniform int32
	Background        mgl32.Vec4

	// interactive objects
	Font          *gltext.Font
	Labels        []*Label
	TextBoxes     []*TextBox
	TextScaleRate float32 // increment during a scale operation

	// opengl oriented
	WindowWidth   float32
	WindowHeight  float32
	program       uint32 // shader program
	glMatrix      int32  // ortho matrix
	position      uint32 // index location
	vao           uint32
	vbo           uint32
	ebo           uint32
	vboData       []float32
	vboIndexCount int
	eboData       []int32
	eboIndexCount int
}

func (menu *Menu) NewLabel(str string) *Label {
	label := &Label{
		Menu: menu,
		Text: gltext.NewText(menu.Font, 1.0, 1.1),
	}
	menu.Labels = append(menu.Labels, label)
	label.SetString(str)
	label.Text.SetScale(1)
	return label
}

func (menu *Menu) AddTextBox(textbox *TextBox, str string, width int32, height int32, borderWidth int32) {
	textbox.Load(menu, width, height, borderWidth)
	textbox.SetString(str)
	textbox.Text.SetScale(1)
	textbox.Text.SetPosition(0, 0)
	textbox.Text.SetColor(0, 0, 0)
	menu.TextBoxes = append(menu.TextBoxes, textbox)
}

func (menu *Menu) Show() {
	for i := range menu.Labels {
		menu.Labels[i].Reset()
	}
	menu.Visible = true
	if menu.OnShow != nil {
		menu.OnShow()
	}
}

func (menu *Menu) Hide() {
	for i := range menu.Labels {
		menu.Labels[i].Reset()
	}
	menu.Visible = false
}

func (menu *Menu) Toggle() {
	for i := range menu.Labels {
		menu.Labels[i].Reset()
	}
	menu.Visible = !menu.Visible
}

// Load will draw a background centered on the screen or positioned based on offsetBy values
func NewMenu(width float32, height float32, scale fixed.Int26_6, offsetBy mgl32.Vec2) (*Menu, error) {
	glfloat_size := 4
	glint_size := 4

	menu := &Menu{}
	menu.Visible = false
	menu.ShowOnKey = glfw.KeyM
	menu.Width = width
	menu.Height = height

	// load font
	fd, err := os.Open("font/luximr.ttf")
	if err != nil {
		return nil, err
	}
	defer fd.Close()

	runesPerRow := fixed.Int26_6(16)
	runeRanges := make(gltext.RuneRanges, 0)
	runeRange := gltext.RuneRange{Low: 32, High: 127}
	runeRanges = append(runeRanges, runeRange)

	menu.Font, err = gltext.NewTruetype(fd, scale, runeRanges, runesPerRow)
	if err != nil {
		return nil, err
	}

	// 2DO: make this time dependent rather than fps dependent
	menu.TextScaleRate = 0.01

	// create shader program and define attributes and uniforms
	menu.program, err = gltext.NewProgram(vertexShaderSource, fragmentShaderSource)
	if err != nil {
		return nil, err
	}
	menu.glMatrix = gl.GetUniformLocation(menu.program, gl.Str("matrix\x00"))
	menu.backgroundUniform = gl.GetUniformLocation(menu.program, gl.Str("background\x00"))
	menu.position = uint32(gl.GetAttribLocation(menu.program, gl.Str("position\x00")))

	gl.GenVertexArrays(1, &menu.vao)
	gl.GenBuffers(1, &menu.vbo)
	gl.GenBuffers(1, &menu.ebo)

	// vao
	gl.BindVertexArray(menu.vao)

	// 2DO: Change text depth to get it to render? For now this works.
	gl.Enable(gl.DEPTH_TEST)
	gl.DepthFunc(gl.LEQUAL)

	// vbo
	// specify the buffer for which the VertexAttribPointer calls apply
	gl.BindBuffer(gl.ARRAY_BUFFER, menu.vbo)

	gl.EnableVertexAttribArray(menu.position)
	gl.VertexAttribPointer(
		menu.position,
		2,
		gl.FLOAT,
		false,
		0, // no stride... yet
		gl.PtrOffset(0),
	)

	// ebo
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, menu.ebo)

	// i am guessing that order is important here
	gl.BindVertexArray(0)
	gl.BindBuffer(gl.ARRAY_BUFFER, 0)
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, 0)

	// ebo, vbo data
	menu.vboIndexCount = 4 * 2 // four indices (2 points per index)
	menu.eboIndexCount = 6     // 6 triangle indices for a quad
	menu.vboData = make([]float32, menu.vboIndexCount, menu.vboIndexCount)
	menu.eboData = make([]int32, menu.eboIndexCount, menu.eboIndexCount)
	menu.lowerLeft = menu.findCenter(offsetBy)
	menu.makeBufferData()

	// setup context
	gl.BindVertexArray(menu.vao)
	gl.BindBuffer(gl.ARRAY_BUFFER, menu.vbo)
	gl.BufferData(
		gl.ARRAY_BUFFER, glfloat_size*menu.vboIndexCount, gl.Ptr(menu.vboData), gl.DYNAMIC_DRAW)
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, menu.ebo)
	gl.BufferData(
		gl.ELEMENT_ARRAY_BUFFER, glint_size*menu.eboIndexCount, gl.Ptr(menu.eboData), gl.DYNAMIC_DRAW)
	gl.BindVertexArray(0)

	// not necesssary, but i just want to better understand using vertex arrays
	gl.BindBuffer(gl.ARRAY_BUFFER, 0)
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, 0)
	return menu, nil
}

func (menu *Menu) ResizeWindow(width float32, height float32) {
	menu.WindowWidth = width
	menu.WindowHeight = height
	menu.Font.ResizeWindow(width, height)
}

func (menu *Menu) makeBufferData() {
	// index (0,0)
	menu.vboData[0] = menu.lowerLeft.X // position
	menu.vboData[1] = menu.lowerLeft.Y

	// index (1,0)
	menu.vboData[2] = menu.lowerLeft.X + menu.Width
	menu.vboData[3] = menu.lowerLeft.Y

	// index (1,1)
	menu.vboData[4] = menu.lowerLeft.X + menu.Width
	menu.vboData[5] = menu.lowerLeft.Y + menu.Height

	// index (0,1)
	menu.vboData[6] = menu.lowerLeft.X
	menu.vboData[7] = menu.lowerLeft.Y + menu.Height

	menu.eboData[0] = 0
	menu.eboData[1] = 1
	menu.eboData[2] = 2
	menu.eboData[3] = 0
	menu.eboData[4] = 2
	menu.eboData[5] = 3
}

func (menu *Menu) Release() {
	gl.DeleteBuffers(1, &menu.vbo)
	gl.DeleteBuffers(1, &menu.ebo)
	gl.DeleteBuffers(1, &menu.vao)
	for i := range menu.Labels {
		menu.Labels[i].Text.Release()
		if menu.Labels[i].Shadow != nil && menu.Labels[i].Shadow.Text != nil {
			menu.Labels[i].Shadow.Text.Release()
		}
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].Text.Release()
	}
}

func (menu *Menu) Draw() bool {
	if !menu.Visible {
		return menu.Visible
	}
	gl.UseProgram(menu.program)

	gl.UniformMatrix4fv(menu.glMatrix, 1, false, &menu.Font.OrthographicMatrix[0])
	gl.Uniform4fv(menu.backgroundUniform, 1, &menu.Background[0])

	gl.Enable(gl.BLEND)
	gl.BlendEquation(gl.FUNC_ADD)
	gl.BlendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA)

	gl.BindVertexArray(menu.vao)
	gl.DrawElements(gl.TRIANGLES, int32(menu.eboIndexCount), gl.UNSIGNED_INT, nil)
	gl.BindVertexArray(0)
	gl.Disable(gl.BLEND)

	for i := range menu.Labels {
		if !menu.Labels[i].IsHover {
			if menu.Labels[i].OnNotHover != nil {
				menu.Labels[i].OnNotHover(menu.Labels[i])
				if menu.Labels[i].Shadow != nil {
					menu.Labels[i].OnNotHover(&menu.Labels[i].Shadow.Label)
				}
			}
		}
		menu.Labels[i].Draw()
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].Draw()
	}
	return menu.Visible
}

func (menu *Menu) MouseClick(xPos, yPos float64, button MouseClick) {
	if !menu.Visible {
		return
	}
	yPos = float64(menu.WindowHeight) - yPos
	for i := range menu.Labels {
		menu.Labels[i].IsClicked(xPos, yPos, button)
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].IsClicked(xPos, yPos, button)
	}
}

func (menu *Menu) MouseRelease(xPos, yPos float64, button MouseClick) {
	if !menu.Visible {
		return
	}
	yPos = float64(menu.WindowHeight) - yPos
	for i := range menu.Labels {
		menu.Labels[i].IsReleased(xPos, yPos, button)
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].IsReleased(xPos, yPos, button)
	}
}

func (menu *Menu) MouseHover(xPos, yPos float64) {
	if !menu.Visible {
		return
	}
	yPos = float64(menu.WindowHeight) - yPos
	for i := range menu.Labels {
		if menu.Labels[i].OnHover != nil {
			menu.Labels[i].IsHovered(xPos, yPos)
		}
	}
}

func (menu *Menu) findCenter(offsetBy mgl32.Vec2) (lowerLeft Point) {
	menuWidthHalf := menu.Width / 2
	menuHeightHalf := menu.Height / 2

	lowerLeft.X = -menuWidthHalf + offsetBy.X()
	lowerLeft.Y = -menuHeightHalf + offsetBy.Y()
	return
}

func (menu *Menu) KeyRelease(key glfw.Key, withShift bool) {
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].KeyRelease(key, withShift)
	}
	if menu.OnEnterRelease != nil && key == glfw.KeyEnter {
		menu.OnEnterRelease()
	}
}
<commit_msg>Cleanup.<commit_after>package glmenu

import (
	"github.com/4ydx/gltext"
	"github.com/go-gl/gl/v3.3-core/gl"
	"github.com/go-gl/glfw/v3.1/glfw"
	"github.com/go-gl/mathgl/mgl32"
)

type Point struct {
	X, Y float32
}

var vertexShaderSource string = `
#version 330

uniform mat4 matrix;

in vec4 position;

void main() {
  gl_Position = matrix * position;
}
` + "\x00"

var fragmentShaderSource string = `
#version 330

uniform vec4 background;
out vec4 fragment_color;

void main() {
  fragment_color = background;
}
` + "\x00"

type MouseClick int

const (
	MouseUnclicked MouseClick = iota
	MouseLeft
	MouseRight
	MouseCenter
)

type Menu struct {
	//trigger
	OnShow         func()
	OnEnterRelease func()

	// options
	Visible      bool
	ShowOnKey    glfw.Key
	Height       float32
	Width        float32
	IsAutoCenter bool
	lowerLeft    Point

	backgroundUniform int32
	Background        mgl32.Vec4

	// interactive objects
	Font          *gltext.Font
	Labels        []*Label
	TextBoxes     []*TextBox
	TextScaleRate float32 // increment during a scale operation

	// opengl oriented
	WindowWidth   float32
	WindowHeight  float32
	program       uint32 // shader program
	glMatrix      int32  // ortho matrix
	position      uint32 // index location
	vao           uint32
	vbo           uint32
	ebo           uint32
	vboData       []float32
	vboIndexCount int
	eboData       []int32
	eboIndexCount int
}

func (menu *Menu) NewLabel(str string) *Label {
	label := &Label{
		Menu: menu,
		Text: gltext.NewText(menu.Font, 1.0, 1.1),
	}
	menu.Labels = append(menu.Labels, label)
	label.SetString(str)
	label.Text.SetScale(1)
	return label
}

func (menu *Menu) AddTextBox(textbox *TextBox, str string, width int32, height int32, borderWidth int32) {
	textbox.Load(menu, width, height, borderWidth)
	textbox.SetString(str)
	textbox.Text.SetScale(1)
	textbox.Text.SetPosition(0, 0)
	textbox.Text.SetColor(0, 0, 0)
	menu.TextBoxes = append(menu.TextBoxes, textbox)
}

func (menu *Menu) Show() {
	for i := range menu.Labels {
		menu.Labels[i].Reset()
	}
	menu.Visible = true
	if menu.OnShow != nil {
		menu.OnShow()
	}
}

func (menu *Menu) Hide() {
	for i := range menu.Labels {
		menu.Labels[i].Reset()
	}
	menu.Visible = false
}

func (menu *Menu) Toggle() {
	for i := range menu.Labels {
		menu.Labels[i].Reset()
	}
	menu.Visible = !menu.Visible
}

// NewMenu creates a new menu object with a background centered on the screen or positioned using offsetBy
func NewMenu(font *gltext.Font, width float32, height float32, offsetBy mgl32.Vec2) (*Menu, error) {
	glfloat_size := 4
	glint_size := 4

	menu := &Menu{
		Font:      font,
		Visible:   false,
		ShowOnKey: glfw.KeyM,
		Width:     width,
		Height:    height,
	}

	// 2DO: make this time dependent rather than fps dependent
	menu.TextScaleRate = 0.01

	// create shader program and define attributes and uniforms
	var err error
	menu.program, err = gltext.NewProgram(vertexShaderSource, fragmentShaderSource)
	if err != nil {
		return nil, err
	}
	menu.glMatrix = gl.GetUniformLocation(menu.program, gl.Str("matrix\x00"))
	menu.backgroundUniform = gl.GetUniformLocation(menu.program, gl.Str("background\x00"))
	menu.position = uint32(gl.GetAttribLocation(menu.program, gl.Str("position\x00")))

	gl.GenVertexArrays(1, &menu.vao)
	gl.GenBuffers(1, &menu.vbo)
	gl.GenBuffers(1, &menu.ebo)

	// vao
	gl.BindVertexArray(menu.vao)

	// 2DO: Change text depth to get it to render? For now this works.
	gl.Enable(gl.DEPTH_TEST)
	gl.DepthFunc(gl.LEQUAL)

	// vbo
	// specify the buffer for which the VertexAttribPointer calls apply
	gl.BindBuffer(gl.ARRAY_BUFFER, menu.vbo)

	gl.EnableVertexAttribArray(menu.position)
	gl.VertexAttribPointer(
		menu.position,
		2,
		gl.FLOAT,
		false,
		0, // no stride... yet
		gl.PtrOffset(0),
	)

	// ebo
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, menu.ebo)

	// i am guessing that order is important here
	gl.BindVertexArray(0)
	gl.BindBuffer(gl.ARRAY_BUFFER, 0)
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, 0)

	// ebo, vbo data
	menu.vboIndexCount = 4 * 2 // four indices (2 points per index)
	menu.eboIndexCount = 6     // 6 triangle indices for a quad
	menu.vboData = make([]float32, menu.vboIndexCount, menu.vboIndexCount)
	menu.eboData = make([]int32, menu.eboIndexCount, menu.eboIndexCount)
	menu.lowerLeft = menu.findCenter(offsetBy)
	menu.makeBufferData()

	// setup context
	gl.BindVertexArray(menu.vao)
	gl.BindBuffer(gl.ARRAY_BUFFER, menu.vbo)
	gl.BufferData(
		gl.ARRAY_BUFFER, glfloat_size*menu.vboIndexCount, gl.Ptr(menu.vboData), gl.DYNAMIC_DRAW)
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, menu.ebo)
	gl.BufferData(
		gl.ELEMENT_ARRAY_BUFFER, glint_size*menu.eboIndexCount, gl.Ptr(menu.eboData), gl.DYNAMIC_DRAW)
	gl.BindVertexArray(0)

	// not necesssary, but i just want to better understand using vertex arrays
	gl.BindBuffer(gl.ARRAY_BUFFER, 0)
	gl.BindBuffer(gl.ELEMENT_ARRAY_BUFFER, 0)
	return menu, nil
}

func (menu *Menu) ResizeWindow(width float32, height float32) {
	menu.WindowWidth = width
	menu.WindowHeight = height
	menu.Font.ResizeWindow(width, height)
}

func (menu *Menu) makeBufferData() {
	// index (0,0)
	menu.vboData[0] = menu.lowerLeft.X // position
	menu.vboData[1] = menu.lowerLeft.Y

	// index (1,0)
	menu.vboData[2] = menu.lowerLeft.X + menu.Width
	menu.vboData[3] = menu.lowerLeft.Y

	// index (1,1)
	menu.vboData[4] = menu.lowerLeft.X + menu.Width
	menu.vboData[5] = menu.lowerLeft.Y + menu.Height

	// index (0,1)
	menu.vboData[6] = menu.lowerLeft.X
	menu.vboData[7] = menu.lowerLeft.Y + menu.Height

	menu.eboData[0] = 0
	menu.eboData[1] = 1
	menu.eboData[2] = 2
	menu.eboData[3] = 0
	menu.eboData[4] = 2
	menu.eboData[5] = 3
}

func (menu *Menu) Release() {
	gl.DeleteBuffers(1, &menu.vbo)
	gl.DeleteBuffers(1, &menu.ebo)
	gl.DeleteBuffers(1, &menu.vao)
	for i := range menu.Labels {
		menu.Labels[i].Text.Release()
		if menu.Labels[i].Shadow != nil && menu.Labels[i].Shadow.Text != nil {
			menu.Labels[i].Shadow.Text.Release()
		}
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].Text.Release()
	}
}

func (menu *Menu) Draw() bool {
	if !menu.Visible {
		return menu.Visible
	}
	gl.UseProgram(menu.program)

	gl.UniformMatrix4fv(menu.glMatrix, 1, false, &menu.Font.OrthographicMatrix[0])
	gl.Uniform4fv(menu.backgroundUniform, 1, &menu.Background[0])

	gl.Enable(gl.BLEND)
	gl.BlendEquation(gl.FUNC_ADD)
	gl.BlendFunc(gl.SRC_ALPHA, gl.ONE_MINUS_SRC_ALPHA)

	gl.BindVertexArray(menu.vao)
	gl.DrawElements(gl.TRIANGLES, int32(menu.eboIndexCount), gl.UNSIGNED_INT, nil)
	gl.BindVertexArray(0)
	gl.Disable(gl.BLEND)

	for i := range menu.Labels {
		if !menu.Labels[i].IsHover {
			if menu.Labels[i].OnNotHover != nil {
				menu.Labels[i].OnNotHover(menu.Labels[i])
				if menu.Labels[i].Shadow != nil {
					menu.Labels[i].OnNotHover(&menu.Labels[i].Shadow.Label)
				}
			}
		}
		menu.Labels[i].Draw()
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].Draw()
	}
	return menu.Visible
}

func (menu *Menu) MouseClick(xPos, yPos float64, button MouseClick) {
	if !menu.Visible {
		return
	}
	yPos = float64(menu.WindowHeight) - yPos
	for i := range menu.Labels {
		menu.Labels[i].IsClicked(xPos, yPos, button)
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].IsClicked(xPos, yPos, button)
	}
}

func (menu *Menu) MouseRelease(xPos, yPos float64, button MouseClick) {
	if !menu.Visible {
		return
	}
	yPos = float64(menu.WindowHeight) - yPos
	for i := range menu.Labels {
		menu.Labels[i].IsReleased(xPos, yPos, button)
	}
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].IsReleased(xPos, yPos, button)
	}
}

func (menu *Menu) MouseHover(xPos, yPos float64) {
	if !menu.Visible {
		return
	}
	yPos = float64(menu.WindowHeight) - yPos
	for i := range menu.Labels {
		if menu.Labels[i].OnHover != nil {
			menu.Labels[i].IsHovered(xPos, yPos)
		}
	}
}

func (menu *Menu) findCenter(offsetBy mgl32.Vec2) (lowerLeft Point) {
	menuWidthHalf := menu.Width / 2
	menuHeightHalf := menu.Height / 2

	lowerLeft.X = -menuWidthHalf + offsetBy.X()
	lowerLeft.Y = -menuHeightHalf + offsetBy.Y()
	return
}

func (menu *Menu) KeyRelease(key glfw.Key, withShift bool) {
	for i := range menu.TextBoxes {
		menu.TextBoxes[i].KeyRelease(key, withShift)
	}
	if menu.OnEnterRelease != nil && key == glfw.KeyEnter {
		menu.OnEnterRelease()
	}
}
<|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) {
	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
	}
	if hasDatetimeColumn(tableMap) {
		if err := fprintln(output, importAST("time")); 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 {
		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 hasDatetimeColumn(t map[string][]*columnSchema) bool {
	for _, schemas := range t {
		for _, schema := range schemas {
			if schema.DataType == "datetime" {
				return true
			}
		}
	}
	return false
}

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>Add `migu` struct tag and move `default' to it<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
		}
		if err := parseStructTag(ret, reflect.StructTag(s)); err != nil {
			return nil, err
		}
	}
	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
	}
	if hasDatetimeColumn(tableMap) {
		if err := fprintln(output, importAST("time")); 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 {
		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 hasDatetimeColumn(t map[string][]*columnSchema) bool {
	for _, schemas := range t {
		for _, schema := range schemas {
			if schema.DataType == "datetime" {
				return true
			}
		}
	}
	return false
}

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
}

func parseStructTag(f *field, tag reflect.StructTag) error {
	migu := tag.Get("migu")
	if migu == "" {
		return nil
	}
	for _, opt := range strings.Split(migu, ",") {
		optval := strings.SplitN(opt, ":", 2)
		switch optval[0] {
		case "default":
			var val string
			if len(optval) > 1 {
				val = optval[1]
			}
			f.Default = formatDefault(f.Type, val)
		default:
			return fmt.Errorf("unknown option: `%s'", opt)
		}
	}
	return 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]),
	}
	var tags []string
	if schema.ColumnDefault.Valid {
		tags = append(tags, "default:"+schema.ColumnDefault.String)
	}
	if len(tags) > 0 {
		field.Tag = &ast.BasicLit{
			Kind:  token.STRING,
			Value: fmt.Sprintf("`migu:\"%s\"`", strings.Join(tags, ",")),
		}
	}
	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 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
	}
	return &ast.Field{
		Names: []*ast.Ident{
			ast.NewIdent(toCamelCase(schema.ColumnName)),
		},
		Type: ast.NewIdent(types[0]),
		// Tag: &ast.BasicLit{
		// Kind:  token.STRING,
		// Value: "",
		// },
	}, 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>Fprint outputs the `default' struct tag if default is specified to column<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) {
	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"
}
<|endoftext|>


<commit_before>package k

import (
	"log"
	"net/url"
	"os"
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

// MustURL calls net/url.Parse() and panics if it returns a non-nil
// error. Useful for URL constants.
func MustURL(rawurl string) *url.URL {
	u, err := url.Parse(rawurl)
	if err != nil {
		panic(err)
	}
	return u
}

// DefaultEnv returns the value of an environment variable, provided
// it has been set. If it is unset (i.e. empty), the specified default
// value is returned.
func DefaultEnv(key, def string) string {
	if val := os.Getenv(key); val != "" {
		return val
	}
	return def
}
<commit_msg>Add MustBytes<commit_after>package k

import (
	"log"
	"net/url"
	"os"
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

// FIXME: Make Must* functions consistent

// MustURL calls net/url.Parse() and panics if it returns a non-nil
// error. Useful for URL constants.
func MustURL(rawurl string) *url.URL {
	u, err := url.Parse(rawurl)
	if err != nil {
		panic(err)
	}
	return u
}

// MustBytes is a function which usually wraps functions which return
// a byte slice and an error. It panics, if the given error is not nil.
func MustBytes(data []byte, err error) []byte {
	if err != nil {
		panic(err)
	}
	return data
}

// DefaultEnv returns the value of an environment variable, provided
// it has been set. If it is unset (i.e. empty), the specified default
// value is returned.
func DefaultEnv(key, def string) string {
	if val := os.Getenv(key); val != "" {
		return val
	}
	return def
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package apidoc

import (
	"math/rand"
	"net/http"
	"time"

	"github.com/issue9/rands"

	"github.com/caixw/apidoc/v7/core"
	"github.com/caixw/apidoc/v7/internal/ast"
	"github.com/caixw/apidoc/v7/internal/locale"
	"github.com/caixw/apidoc/v7/internal/mock"
)

// Range 表示数值的范围
type Range struct {
	Min, Max int
}

func (r *Range) sanitize() *core.SyntaxError {
	if r.Max <= r.Min {
		return core.NewSyntaxError(core.Location{}, "Min", locale.ErrInvalidValue)
	}
	return nil
}

// MockOptions mock 的一些随机设置项
type MockOptions struct {
	Indent    string            // 缩进字符串
	Servers   map[string]string // 为文档中所有 server 以及对应的路由前缀。
	SliceSize Range             // 指定用于生成数组大小范围的数值

	NumberSize  Range // 指定用于生成数值数据的范围
	EnableFloat bool  // 是否允许生成浮点数

	StringSize  Range  // 指定生成随机字符串的长度范围
	StringAlpha []byte // 指定生成字符串可用的字符

	URLDomains        []string // 指定生成 url 类型数据时可用的域名，默认为 example.com
	EmailDomains      []string // 指定生成 email 类型数据时可用的域名，默认为 example.com
	EmailUsernameSize Range    // 指定生成 email 类型数据的用户名长度范围，默认 [3,8]

	ImageBasePrefix string // 图片的基地址

	DateStart time.Time
	DateEnd   time.Time
	dateSize  int64 // 由 gen 方法中根据 DateStart 和 DateEnd 生成
}

var defaultMockOptions = &MockOptions{
	Indent:    "\t",
	SliceSize: Range{Min: 5, Max: 50},

	NumberSize:  Range{Min: 100, Max: 10000},
	EnableFloat: false,

	StringSize:  Range{Min: 50, Max: 1024},
	StringAlpha: rands.AlphaNumber,

	URLDomains:        []string{"https://example.com/"},
	EmailDomains:      []string{"example.com"},
	EmailUsernameSize: Range{Min: 3, Max: 8},

	ImageBasePrefix: "/__images__",

	DateStart: time.Now().Add(-time.Hour * 24 * 365),
	DateEnd:   time.Now().Add(time.Hour * 24 * 3650),
}

func (o *MockOptions) sanitize() *core.SyntaxError {
	if err := o.SliceSize.sanitize(); err != nil {
		err.Field = "SliceSize." + err.Field
		return err
	}

	if err := o.NumberSize.sanitize(); err != nil {
		err.Field = "NumberSize." + err.Field
		return err
	}

	if err := o.StringSize.sanitize(); err != nil {
		err.Field = "StringSize." + err.Field
		return err
	}

	if err := o.EmailUsernameSize.sanitize(); err != nil {
		err.Field = "EmailUsernameSize." + err.Field
		return err
	}

	if len(o.StringAlpha) == 0 {
		return core.NewSyntaxError(core.Location{}, "StringAlpha", locale.ErrRequired)
	}

	if len(o.URLDomains) == 0 {
		return core.NewSyntaxError(core.Location{}, "URLDomains", locale.ErrRequired)
	}

	if len(o.EmailDomains) == 0 {
		return core.NewSyntaxError(core.Location{}, "EmailDomains", locale.ErrRequired)
	}

	o.dateSize = o.DateEnd.Unix() - o.DateStart.Unix() - 86400
	if o.dateSize <= 0 {
		return core.NewSyntaxError(core.Location{}, "DateStart", locale.ErrInvalidValue)
	}

	return nil
}

func (o *MockOptions) gen() (*mock.GenOptions, error) {
	if o == nil {
		o = defaultMockOptions
	} else if err := o.sanitize(); err != nil {
		err.Field = "MockOptions." + err.Field
		return nil, err
	}

	return &mock.GenOptions{
		Number: func(p *ast.Param) interface{} {
			switch p.Type.V() {
			case ast.TypeFloat:
				return o.float()
			case ast.TypeInt:
				return o.integer()
			}

			if !o.EnableFloat {
				return o.integer()
			}

			if rand.Int()%2 == 0 {
				return o.integer()
			}
			return o.float()
		},

		String: func(p *ast.Param) string {
			switch p.Type.V() {
			case ast.TypeEmail:
				return o.email()
			case ast.TypeURL:
				return o.url()
			case ast.TypeImage:
				return o.image()
			case ast.TypeDate:
				return o.date()
			case ast.TypeTime:
				return o.time()
			case ast.TypeDateTime:
				return o.dateTime()
			}
			return rands.String(o.StringSize.Min, o.StringSize.Max, o.StringAlpha)
		},

		Bool: func() bool {
			return rand.Int()%2 == 0
		},

		SliceSize: func() int {
			return rand.Intn(o.SliceSize.Max-o.SliceSize.Min) + o.SliceSize.Min
		},

		Index: func(max int) int {
			return rand.Intn(max)
		},
	}, nil
}

func (o *MockOptions) integer() int {
	return rand.Intn(o.NumberSize.Max-o.NumberSize.Min) + o.NumberSize.Min
}

func (o *MockOptions) float() float32 {
	return float32(o.NumberSize.Min) + rand.Float32()*float32(o.NumberSize.Max-o.NumberSize.Min)
}

func (o *MockOptions) url() string {
	url := o.URLDomains[rand.Intn(len(o.URLDomains))]
	if url[len(url)-1] != '/' {
		url += "/"
	}

	size := rand.Intn(4)
	for i := 0; i < size; i++ {
		url += rands.String(1, 5, rands.AlphaNumber) + "/"
	}
	return url
}

func (o *MockOptions) email() string {
	domain := o.EmailDomains[rand.Intn(len(o.EmailDomains))]
	username := rands.String(o.EmailUsernameSize.Min, o.EmailUsernameSize.Max, rands.AlphaNumber)
	return username + "@" + domain
}

func (o *MockOptions) image() string {
	path := o.ImageBasePrefix
	if path[len(path)-1] != '/' {
		path += "/"
	}
	return path + rands.String(1, 5, rands.AlphaNumber)
}

func (o *MockOptions) date() string {
	s := rand.Int63n(o.dateSize)
	return o.DateStart.Add(time.Duration(s) * time.Second).Format(ast.DateFormat)
}

func (o *MockOptions) time() string {
	d := rand.Int63n(86400)
	return o.DateStart.Add(time.Duration(d) * time.Second).Format(ast.TimeFormat)
}

func (o *MockOptions) dateTime() string {
	return o.date() + "T" + o.time()
}

// Mock 生成 Mock 中间件
//
// data 为文档内容；
// options 用于生成 Mock 数据的随机项，如果为 nil，则会使用一些默认值；
func Mock(h *core.MessageHandler, data []byte, o *MockOptions) (http.Handler, error) {
	g, err := o.gen()
	if err != nil {
		return nil, err
	}

	d := &ast.APIDoc{}
	d.Parse(h, core.Block{Data: data})
	return mock.New(h, d, o.Indent, o.ImageBasePrefix, o.Servers, g)
}

// MockFile 生成 Mock 中间件
//
// path 为文档路径，可以是本地路径也可以是 URL，根据是否为 http 或是 https 开头做判断；
// options 用于生成 Mock 数据的随机项，如果为 nil，则会使用一些默认值；
func MockFile(h *core.MessageHandler, path core.URI, o *MockOptions) (http.Handler, error) {
	g, err := o.gen()
	if err != nil {
		return nil, err
	}

	return mock.Load(h, path, o.Indent, o.ImageBasePrefix, o.Servers, g)
}
<commit_msg>docs: 更新函数的文档说明<commit_after>// SPDX-License-Identifier: MIT

package apidoc

import (
	"math/rand"
	"net/http"
	"time"

	"github.com/issue9/rands"

	"github.com/caixw/apidoc/v7/core"
	"github.com/caixw/apidoc/v7/internal/ast"
	"github.com/caixw/apidoc/v7/internal/locale"
	"github.com/caixw/apidoc/v7/internal/mock"
)

// Range 表示数值的范围
type Range struct {
	Min, Max int
}

func (r *Range) sanitize() *core.SyntaxError {
	if r.Max <= r.Min {
		return core.NewSyntaxError(core.Location{}, "Min", locale.ErrInvalidValue)
	}
	return nil
}

// MockOptions mock 的一些随机设置项
type MockOptions struct {
	Indent    string            // 缩进字符串
	Servers   map[string]string // 为文档中所有 server 以及对应的路由前缀。
	SliceSize Range             // 指定用于生成数组大小范围的数值

	NumberSize  Range // 指定用于生成数值数据的范围
	EnableFloat bool  // 是否允许生成浮点数

	StringSize  Range  // 指定生成随机字符串的长度范围
	StringAlpha []byte // 指定生成字符串可用的字符

	URLDomains        []string // 指定生成 url 类型数据时可用的域名，默认为 example.com
	EmailDomains      []string // 指定生成 email 类型数据时可用的域名，默认为 example.com
	EmailUsernameSize Range    // 指定生成 email 类型数据的用户名长度范围，默认 [3,8]

	ImageBasePrefix string // 图片的基地址

	DateStart time.Time
	DateEnd   time.Time
	dateSize  int64 // 由 gen 方法中根据 DateStart 和 DateEnd 生成
}

var defaultMockOptions = &MockOptions{
	Indent:    "\t",
	SliceSize: Range{Min: 5, Max: 50},

	NumberSize:  Range{Min: 100, Max: 10000},
	EnableFloat: false,

	StringSize:  Range{Min: 50, Max: 1024},
	StringAlpha: rands.AlphaNumber,

	URLDomains:        []string{"https://example.com/"},
	EmailDomains:      []string{"example.com"},
	EmailUsernameSize: Range{Min: 3, Max: 8},

	ImageBasePrefix: "/__images__",

	DateStart: time.Now().Add(-time.Hour * 24 * 365),
	DateEnd:   time.Now().Add(time.Hour * 24 * 3650),
}

func (o *MockOptions) sanitize() *core.SyntaxError {
	if err := o.SliceSize.sanitize(); err != nil {
		err.Field = "SliceSize." + err.Field
		return err
	}

	if err := o.NumberSize.sanitize(); err != nil {
		err.Field = "NumberSize." + err.Field
		return err
	}

	if err := o.StringSize.sanitize(); err != nil {
		err.Field = "StringSize." + err.Field
		return err
	}

	if err := o.EmailUsernameSize.sanitize(); err != nil {
		err.Field = "EmailUsernameSize." + err.Field
		return err
	}

	if len(o.StringAlpha) == 0 {
		return core.NewSyntaxError(core.Location{}, "StringAlpha", locale.ErrRequired)
	}

	if len(o.URLDomains) == 0 {
		return core.NewSyntaxError(core.Location{}, "URLDomains", locale.ErrRequired)
	}

	if len(o.EmailDomains) == 0 {
		return core.NewSyntaxError(core.Location{}, "EmailDomains", locale.ErrRequired)
	}

	o.dateSize = o.DateEnd.Unix() - o.DateStart.Unix() - 86400
	if o.dateSize <= 0 {
		return core.NewSyntaxError(core.Location{}, "DateStart", locale.ErrInvalidValue)
	}

	return nil
}

func (o *MockOptions) gen() (*mock.GenOptions, error) {
	if o == nil {
		o = defaultMockOptions
	} else if err := o.sanitize(); err != nil {
		err.Field = "MockOptions." + err.Field
		return nil, err
	}

	return &mock.GenOptions{
		Number: func(p *ast.Param) interface{} {
			switch p.Type.V() {
			case ast.TypeFloat:
				return o.float()
			case ast.TypeInt:
				return o.integer()
			}

			if !o.EnableFloat {
				return o.integer()
			}

			if rand.Int()%2 == 0 {
				return o.integer()
			}
			return o.float()
		},

		String: func(p *ast.Param) string {
			switch p.Type.V() {
			case ast.TypeEmail:
				return o.email()
			case ast.TypeURL:
				return o.url()
			case ast.TypeImage:
				return o.image()
			case ast.TypeDate:
				return o.date()
			case ast.TypeTime:
				return o.time()
			case ast.TypeDateTime:
				return o.dateTime()
			}
			return rands.String(o.StringSize.Min, o.StringSize.Max, o.StringAlpha)
		},

		Bool: func() bool {
			return rand.Int()%2 == 0
		},

		SliceSize: func() int {
			return rand.Intn(o.SliceSize.Max-o.SliceSize.Min) + o.SliceSize.Min
		},

		Index: func(max int) int {
			return rand.Intn(max)
		},
	}, nil
}

func (o *MockOptions) integer() int {
	return rand.Intn(o.NumberSize.Max-o.NumberSize.Min) + o.NumberSize.Min
}

func (o *MockOptions) float() float32 {
	return float32(o.NumberSize.Min) + rand.Float32()*float32(o.NumberSize.Max-o.NumberSize.Min)
}

func (o *MockOptions) url() string {
	url := o.URLDomains[rand.Intn(len(o.URLDomains))]
	if url[len(url)-1] != '/' {
		url += "/"
	}

	size := rand.Intn(4)
	for i := 0; i < size; i++ {
		url += rands.String(1, 5, rands.AlphaNumber) + "/"
	}
	return url
}

func (o *MockOptions) email() string {
	domain := o.EmailDomains[rand.Intn(len(o.EmailDomains))]
	username := rands.String(o.EmailUsernameSize.Min, o.EmailUsernameSize.Max, rands.AlphaNumber)
	return username + "@" + domain
}

func (o *MockOptions) image() string {
	path := o.ImageBasePrefix
	if path[len(path)-1] != '/' {
		path += "/"
	}
	return path + rands.String(1, 5, rands.AlphaNumber)
}

func (o *MockOptions) date() string {
	s := rand.Int63n(o.dateSize)
	return o.DateStart.Add(time.Duration(s) * time.Second).Format(ast.DateFormat)
}

func (o *MockOptions) time() string {
	d := rand.Int63n(86400)
	return o.DateStart.Add(time.Duration(d) * time.Second).Format(ast.TimeFormat)
}

func (o *MockOptions) dateTime() string {
	return o.date() + "T" + o.time()
}

// Mock 生成 Mock 中间件
//
// data 为文档内容；
// o 用于生成 Mock 数据的随机项，如果为 nil，则会采用默认配置项；
func Mock(h *core.MessageHandler, data []byte, o *MockOptions) (http.Handler, error) {
	g, err := o.gen()
	if err != nil {
		return nil, err
	}

	d := &ast.APIDoc{}
	d.Parse(h, core.Block{Data: data})
	return mock.New(h, d, o.Indent, o.ImageBasePrefix, o.Servers, g)
}

// MockFile 生成 Mock 中间件
//
// path 为文档路径；
// o 用于生成 Mock 数据的随机项，如果为 nil，则会采用默认配置项；
func MockFile(h *core.MessageHandler, path core.URI, o *MockOptions) (http.Handler, error) {
	g, err := o.gen()
	if err != nil {
		return nil, err
	}

	return mock.Load(h, path, o.Indent, o.ImageBasePrefix, o.Servers, g)
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "net"
)

// Connection struct
type Connection struct {
    server string
    socket net.Conn

    nick, user string
}

// Create a new IRC connection object
func IRC(nick, user string) *Connection {
    irc := &Connection {
        nick: nick,
        user: user,
    }
    return irc
}

// Method to connect to IRC server
func (irc *Connection) Connect(server string) error {
    var err error
    irc.socket, err = net.Dial("tcp", server)
    if err != nil {
        return err
    }
    return nil
}

// Main loop
func main() {
    fmt.Println("Moss is starting...")

    const server = "irc.freenode.org:6667"
    const channel = "#coredump"

    irc := IRC("mossbot", "mossbot")
    err := irc.Connect(server)
    if err != nil {
        panic(err)
    }
}
<commit_msg>Write loop<commit_after>package main

import (
    "bufio"
    "fmt"
    "net"
)

// Connection struct
type Connection struct {
    // Server connection
    server string
    socket net.Conn

    // User information
    nick, user string

    // Read/write channels
    chanRead, chanWrite chan string

    // Exit channels
    chanExit chan bool
}

// Create a new IRC connection object
func IRC(nick, user string) *Connection {
    irc := &Connection {
        nick: nick,
        user: user,
        chanExit: make(chan bool),
    }
    return irc
}

// Method to connect to IRC server
func (irc *Connection) Connect(server string) error {
    var err error
    irc.server = server
    irc.socket, err = net.Dial("tcp", server)
    if err != nil {
        return err
    }
    fmt.Printf("Connected to %s (%s)\n", irc.server, irc.socket.RemoteAddr())

    irc.chanRead = make(chan string, 16)
    irc.chanWrite = make(chan string, 16)

    go irc.readLoop()
    go irc.writeLoop()

    return nil
}

// Write loop
func (irc *Connection) writeLoop() {
    for {
        data := <-irc.chanWrite

        fmt.Printf("--> %s\n", data)
        _, err := irc.socket.Write([]byte(data))
        if err != nil {
            // TODO log errors
            break
        }
    }
}

// Read loop
func (irc *Connection) readLoop() {
}

// Main loop
func main() {
    fmt.Println("Moss is starting...")

    const server = "irc.freenode.org:6667"
    const channel = "#coredump"

    irc := IRC("mossbot", "mossbot")
    err := irc.Connect(server)
    if err != nil {
        panic(err)
    }

    irc.chanWrite <- fmt.Sprintf("JOIN %s\r\n", channel)

    <-irc.chanExit
}
<|endoftext|>


<commit_before>// Package mqtt implements MQTT clients and servers
package mqtt

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"

	proto "github.com/huin/mqtt"
)

type subscriptions struct {
	workers int
	posts   chan (post)

	mu        sync.Mutex // guards access to fields below
	subs      map[string][]*connection
	wildcards []wild
}

// The length of the queue that subscription processing
// workers are taking from.
const postQueue = 100

func newSubscriptions(workers int) *subscriptions {
	s := &subscriptions{
		subs:    make(map[string][]*connection),
		posts:   make(chan post, postQueue),
		workers: workers,
	}
	for i := 0; i < s.workers; i++ {
		go s.run(i)
	}
	return s
}

func (s *subscriptions) add(topic string, c *connection) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if isWildcard(topic) {
		w := newWild(topic, c)
		if w.valid() {
			s.wildcards = append(s.wildcards, w)
		}
	} else {
		s.subs[topic] = append(s.subs[topic], c)
	}
}

type wild struct {
	wild []string
	c    *connection
}

func newWild(topic string, c *connection) wild {
	return wild{wild: strings.Split(topic, "/"), c: c}
}

func (w wild) matches(parts []string) bool {
	i := 0
	for i < len(parts) {
		// topic is longer, no match
		if i >= len(w.wild) {
			return false
		}
		// matched up to here, and now the wildcard says "all others will match"
		if w.wild[i] == "#" {
			return true
		}
		// text does not match, and there wasn't a + to excuse it
		if parts[i] != w.wild[i] && w.wild[i] != "+" {
			return false
		}
		i++
	}

	// make finance/stock/ibm/# match finance/stock/ibm
	if i == len(w.wild)-1 && w.wild[len(w.wild)-1] == "#" {
		return true
	}

	if i == len(w.wild) {
		return true
	}
	return false
}

// Find all connections that are subscribed to this topic.
func (s *subscriptions) subscribers(topic string) []*connection {
	s.mu.Lock()
	defer s.mu.Unlock()

	// non-wildcard subscribers
	res := s.subs[topic]

	// process wildcards
	parts := strings.Split(topic, "/")
	for _, w := range s.wildcards {
		if w.matches(parts) {
			res = append(res, w.c)
		}
	}

	return res
}

// Remove all subscriptions that refer to a connection.
func (s *subscriptions) unsubAll(c *connection) {
	s.mu.Lock()
	for _, v := range s.subs {
		for i := range v {
			if v[i] == c {
				v[i] = nil
			}
		}
	}

	// remove any associated entries in the wildcard list
	var wildNew []wild
	for i := 0; i < len(s.wildcards); i++ {
		if s.wildcards[i].c != c {
			wildNew = append(wildNew, s.wildcards[i])
		}
	}
	s.wildcards = wildNew

	s.mu.Unlock()
}

// Remove the subscription to topic for a given connection.
func (s *subscriptions) unsub(topic string, c *connection) {
	s.mu.Lock()
	if subs, ok := s.subs[topic]; ok {
		nils := 0

		// Search the list, removing references to our connection.
		// At the same time, count the nils to see if this list is now empty.
		for i := 0; i < len(subs); i++ {
			if subs[i] == c {
				subs[i] = nil
			}
			if subs[i] == nil {
				nils++
			}
		}

		if nils == len(subs) {
			delete(s.subs, topic)
		}
	}
	s.mu.Unlock()
}

// The subscription processing worker.
func (s *subscriptions) run(id int) {
	tag := fmt.Sprintf("worker %d ", id)
	log.Print(tag, "started")
	for post := range s.posts {
		conns := s.subscribers(post.m.TopicName)
		for _, c := range conns {
			if c != nil {
				c.submit(post.m)
			}
		}
	}
}

func (s *subscriptions) submit(c *connection, m *proto.Publish) {
	s.posts <- post{c: c, m: m}
}

// A post is a unit of work for the subscription processing workers.
type post struct {
	c *connection
	m *proto.Publish
}

type Server struct {
	l    net.Listener
	Done chan bool
	subs *subscriptions
}

func NewServer(l net.Listener) *Server {
	s := &Server{
		l:    l,
		Done: make(chan bool),
		subs: newSubscriptions(2), // 2 workers for now, to see it working
	}

	go s.run()
	return s
}

func (s *Server) run() {
	for {
		conn, err := s.l.Accept()
		if err != nil {
			log.Print("Accept", err)
		}

		cli := newConnection(s, conn)
		go cli.reader()
		go cli.writer()
	}
}

type connection struct {
	svr      *Server
	conn     net.Conn
	jobs     chan job
	clientid string
}

var clients map[string]*connection = make(map[string]*connection)
var clientsMu sync.Mutex

const sendingQueueLength = 100

func newConnection(svr *Server, conn net.Conn) *connection {
	return &connection{
		svr:  svr,
		conn: conn,
		jobs: make(chan job, sendingQueueLength),
	}
}

type signal struct{}
type receipt chan signal

// Wait for the receipt to indicate that the job is done.
func (r receipt) wait() {
	// TODO: timeout
	<-r
}

type job struct {
	m proto.Message
	r receipt
}

// Add this	connection to the map, or find out that an existing connection
// already exists for the same client-id.
func (c *connection) add() *connection {
	clientsMu.Lock()
	defer clientsMu.Unlock()

	existing, ok := clients[c.clientid]
	if !ok {
		// this client id already exists, return it
		return existing
	}

	clients[c.clientid] = c
	return nil
}

// Delete a connection; the conection must be closed by the caller first.
func (c *connection) del() {
	clientsMu.Lock()
	defer clientsMu.Unlock()
	delete(clients, c.clientid)
	return
}

// Replace any existing connection with this one. The one to be replaced,
// if any, must be closed first by the caller.
func (c *connection) replace() {
	clientsMu.Lock()
	defer clientsMu.Unlock()

	// Check that any existing connection is already closed.
	existing, ok := clients[c.clientid]
	if ok {
		die := false
		select {
		case _, ok := <-existing.jobs:
			// what? we are expecting that this channel is closed!
			if ok {
				die = true
			}
		default:
			die = true
		}
		if die {
			panic("attempting to replace a connection that is not closed")
		}

		delete(clients, c.clientid)
	}

	clients[c.clientid] = c
	return
}

// Queue a message; no notification of sending is done.
func (c *connection) submit(m proto.Message) {
	j := job{m: m}
	c.jobs <- j
	return
}

// Queue a message, returns a channel that will be readable
// when the message is sent.
func (c *connection) submitSync(m proto.Message) receipt {
	j := job{m: m, r: make(receipt)}
	c.jobs <- j
	return j.r
}

func (c *connection) reader() {
	// On exit, close the connection and arrange for the writer to exit
	// by closing the output channel.
	defer func() {
		c.conn.Close()
		close(c.jobs)
		log.Print("reader: done")
	}()

	for {
		// TODO: timeout (first message and/or keepalives)
		m, err := proto.DecodeOneMessage(c.conn, nil)
		if err != nil {
			log.Print("reader: ", err)
			return
		}

		switch m := m.(type) {
		case *proto.Connect:
			rc := proto.RetCodeAccepted

			if m.ProtocolName != "MQIsdp" ||
				m.ProtocolVersion != 3 {
				log.Print("reader: reject connection from ", m.ProtocolName, " version ", m.ProtocolVersion)
				rc = proto.RetCodeUnacceptableProtocolVersion
			}

			if len(m.ClientId) < 1 || len(m.ClientId) > 23 {
				rc = proto.RetCodeIdentifierRejected
			}
			c.clientid = m.ClientId
			if existing := c.add(); existing != nil {
				disconnect := &proto.Disconnect{}
				r := existing.submitSync(disconnect)
				r.wait()
				existing.del()
			}
			c.add()

			// TODO: Last will

			connack := &proto.ConnAck{
				ReturnCode: rc,
			}
			c.submit(connack)

			// close connection if it was a bad connect
			if rc != proto.RetCodeAccepted {
				return
			}

		case *proto.Publish:
			// TODO: Proper QoS support
			if m.Header.QosLevel != proto.QosAtMostOnce {
				log.Printf("reader: no support for QoS %v yet", m.Header.QosLevel)
				return
			}
			if isWildcard(m.TopicName) {
				log.Print("reader: ignoring PUBLISH with wildcard topic ", m.TopicName)
			} else {
				c.svr.subs.submit(c, m)
			}
			c.submit(&proto.PubAck{MessageId: m.MessageId})

		case *proto.PingReq:
			c.submit(&proto.PingResp{})

		case *proto.Subscribe:
			suback := &proto.SubAck{
				MessageId: m.MessageId,
				TopicsQos: make([]proto.QosLevel, len(m.Topics)),
			}
			for i, tq := range m.Topics {
				// TODO: Handle varying QoS correctly
				c.svr.subs.add(tq.Topic, c)
				suback.TopicsQos[i] = proto.QosAtLeastOnce
			}
			c.submit(suback)

		case *proto.Unsubscribe:
			for _, t := range m.Topics {
				c.svr.subs.unsub(t, c)
			}
			ack := &proto.UnsubAck{MessageId: m.MessageId}
			c.submit(ack)

		case *proto.Disconnect:
			return

		default:
			log.Printf("reader: unknown msg type %T", m)
			return
		}
	}
}

func (c *connection) writer() {

	// Close connection on exit in order to cause reader to exit.
	defer func() {
		c.conn.Close()
		c.del()
		c.svr.subs.unsubAll(c)
		log.Print("writer: done")
	}()

	for job := range c.jobs {
		if job.m == nil {
			log.Print("writer: no message?")
			return
		}

		// TODO: write timeout
		err := job.m.Encode(c.conn)
		if job.r != nil {
			// notifiy the sender that this message is sent
			close(job.r)
		}
		if err != nil {
			log.Print("writer: ", err)
			return
		}

		if _, ok := job.m.(*proto.Disconnect); ok {
			log.Print("writer: sent disconnect message")
			return
		}
	}
}

// header is used to initialize a proto.Header when the zero value
// is not correct. The zero value of proto.Header is
// the equivalent of header(dupFalse, proto.QosAtMostOnce, retainFalse)
// and is correct for most messages.
func header(d dupFlag, q proto.QosLevel, r retainFlag) proto.Header {
	return proto.Header{
		DupFlag: bool(d), QosLevel: q, Retain: bool(r),
	}
}

type retainFlag bool
type dupFlag bool

const (
	retainFalse retainFlag = false
	retainTrue             = true
	dupFalse    dupFlag    = false
	dupTrue                = true
)

func isWildcard(topic string) bool {
	if strings.Contains(topic, "#") || strings.Contains(topic, "+") {
		return true
	}
	return false
}

func (w wild) valid() bool {
	for i, part := range w.wild {
		// catch things like finance#
		if isWildcard(part) && len(part) != 1 {
			return false
		}
		// # can only occur as the last part
		if part == "#" && i != len(w.wild)-1 {
			return false
		}
	}
	return true
}
<commit_msg>Added retain.<commit_after>// Package mqtt implements MQTT clients and servers
package mqtt

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"

	proto "github.com/huin/mqtt"
)

type subscriptions struct {
	workers int
	posts   chan (post)

	mu        sync.Mutex // guards access to fields below
	subs      map[string][]*connection
	wildcards []wild
	retain		map[string]*proto.Publish
}

// The length of the queue that subscription processing
// workers are taking from.
const postQueue = 100

func newSubscriptions(workers int) *subscriptions {
	s := &subscriptions{
		subs:    make(map[string][]*connection),
		retain:  make(map[string]*proto.Publish),
		posts:   make(chan post, postQueue),
		workers: workers,
	}
	for i := 0; i < s.workers; i++ {
		go s.run(i)
	}
	return s
}

func (s *subscriptions)sendRetain(topic string, c *connection) {
	s.mu.Lock()
				if message, ok := s.retain[topic]; ok {
					c.submit(message)
				}
	s.mu.Unlock()
}

func (s *subscriptions) add(topic string, c *connection) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if isWildcard(topic) {
		w := newWild(topic, c)
		if w.valid() {
			s.wildcards = append(s.wildcards, w)
		}
	} else {
		s.subs[topic] = append(s.subs[topic], c)
	}
}

type wild struct {
	wild []string
	c    *connection
}

func newWild(topic string, c *connection) wild {
	return wild{wild: strings.Split(topic, "/"), c: c}
}

func (w wild) matches(parts []string) bool {
	i := 0
	for i < len(parts) {
		// topic is longer, no match
		if i >= len(w.wild) {
			return false
		}
		// matched up to here, and now the wildcard says "all others will match"
		if w.wild[i] == "#" {
			return true
		}
		// text does not match, and there wasn't a + to excuse it
		if parts[i] != w.wild[i] && w.wild[i] != "+" {
			return false
		}
		i++
	}

	// make finance/stock/ibm/# match finance/stock/ibm
	if i == len(w.wild)-1 && w.wild[len(w.wild)-1] == "#" {
		return true
	}

	if i == len(w.wild) {
		return true
	}
	return false
}

// Find all connections that are subscribed to this topic.
func (s *subscriptions) subscribers(topic string) []*connection {
	s.mu.Lock()
	defer s.mu.Unlock()

	// non-wildcard subscribers
	res := s.subs[topic]

	// process wildcards
	parts := strings.Split(topic, "/")
	for _, w := range s.wildcards {
		if w.matches(parts) {
			res = append(res, w.c)
		}
	}

	return res
}

// Remove all subscriptions that refer to a connection.
func (s *subscriptions) unsubAll(c *connection) {
	s.mu.Lock()
	for _, v := range s.subs {
		for i := range v {
			if v[i] == c {
				v[i] = nil
			}
		}
	}

	// remove any associated entries in the wildcard list
	var wildNew []wild
	for i := 0; i < len(s.wildcards); i++ {
		if s.wildcards[i].c != c {
			wildNew = append(wildNew, s.wildcards[i])
		}
	}
	s.wildcards = wildNew

	s.mu.Unlock()
}

// Remove the subscription to topic for a given connection.
func (s *subscriptions) unsub(topic string, c *connection) {
	s.mu.Lock()
	if subs, ok := s.subs[topic]; ok {
		nils := 0

		// Search the list, removing references to our connection.
		// At the same time, count the nils to see if this list is now empty.
		for i := 0; i < len(subs); i++ {
			if subs[i] == c {
				subs[i] = nil
			}
			if subs[i] == nil {
				nils++
			}
		}

		if nils == len(subs) {
			delete(s.subs, topic)
		}
	}
	s.mu.Unlock()
}

// The subscription processing worker.
func (s *subscriptions) run(id int) {
	tag := fmt.Sprintf("worker %d ", id)
	log.Print(tag, "started")
	for post := range s.posts {

		// Remember the original retain setting, but send out immediate
		// copies without retain: "When a server sends a PUBLISH to a client
		// as a result of a subscription that already existed when the
		// original PUBLISH arrived, the Retain flag should not be set,
		// regardless of the Retain flag of the original PUBLISH.
		retain := post.m.Header.Retain
		post.m.Header.Retain = false

		if retain && post.m.Payload.Size() == 0 {
			s.mu.Lock()
			delete(s.retain, post.m.TopicName)
			s.mu.Unlock()
			// Once the "delete retain" is done, return.
			return
		}

		// Find all the connections that should be notified of this message.
		conns := s.subscribers(post.m.TopicName)
		for _, c := range conns {
			if c != nil {
				c.submit(post.m)
			}
		}

		if retain {
			post.m.Header.Retain = true
			s.mu.Lock()
			s.retain[post.m.TopicName] = post.m
			s.mu.Unlock()
		}
	}
}

func (s *subscriptions) submit(c *connection, m *proto.Publish) {
	s.posts <- post{c: c, m: m}
}

// A post is a unit of work for the subscription processing workers.
type post struct {
	c *connection
	m *proto.Publish
}

type Server struct {
	l    net.Listener
	Done chan bool
	subs *subscriptions
}

func NewServer(l net.Listener) *Server {
	s := &Server{
		l:    l,
		Done: make(chan bool),
		subs: newSubscriptions(2), // 2 workers for now, to see it working
	}

	go s.run()
	return s
}

func (s *Server) run() {
	for {
		conn, err := s.l.Accept()
		if err != nil {
			log.Print("Accept", err)
		}

		cli := newConnection(s, conn)
		go cli.reader()
		go cli.writer()
	}
}

type connection struct {
	svr      *Server
	conn     net.Conn
	jobs     chan job
	clientid string
}

var clients map[string]*connection = make(map[string]*connection)
var clientsMu sync.Mutex

const sendingQueueLength = 100

func newConnection(svr *Server, conn net.Conn) *connection {
	return &connection{
		svr:  svr,
		conn: conn,
		jobs: make(chan job, sendingQueueLength),
	}
}

type signal struct{}
type receipt chan signal

// Wait for the receipt to indicate that the job is done.
func (r receipt) wait() {
	// TODO: timeout
	<-r
}

type job struct {
	m proto.Message
	r receipt
}

// Add this	connection to the map, or find out that an existing connection
// already exists for the same client-id.
func (c *connection) add() *connection {
	clientsMu.Lock()
	defer clientsMu.Unlock()

	existing, ok := clients[c.clientid]
	if !ok {
		// this client id already exists, return it
		return existing
	}

	clients[c.clientid] = c
	return nil
}

// Delete a connection; the conection must be closed by the caller first.
func (c *connection) del() {
	clientsMu.Lock()
	defer clientsMu.Unlock()
	delete(clients, c.clientid)
	return
}

// Replace any existing connection with this one. The one to be replaced,
// if any, must be closed first by the caller.
func (c *connection) replace() {
	clientsMu.Lock()
	defer clientsMu.Unlock()

	// Check that any existing connection is already closed.
	existing, ok := clients[c.clientid]
	if ok {
		die := false
		select {
		case _, ok := <-existing.jobs:
			// what? we are expecting that this channel is closed!
			if ok {
				die = true
			}
		default:
			die = true
		}
		if die {
			panic("attempting to replace a connection that is not closed")
		}

		delete(clients, c.clientid)
	}

	clients[c.clientid] = c
	return
}

// Queue a message; no notification of sending is done.
func (c *connection) submit(m proto.Message) {
	j := job{m: m}
	c.jobs <- j
	return
}

// Queue a message, returns a channel that will be readable
// when the message is sent.
func (c *connection) submitSync(m proto.Message) receipt {
	j := job{m: m, r: make(receipt)}
	c.jobs <- j
	return j.r
}

func (c *connection) reader() {
	// On exit, close the connection and arrange for the writer to exit
	// by closing the output channel.
	defer func() {
		c.conn.Close()
		close(c.jobs)
		log.Print("reader: done")
	}()

	for {
		// TODO: timeout (first message and/or keepalives)
		m, err := proto.DecodeOneMessage(c.conn, nil)
		if err != nil {
			log.Print("reader: ", err)
			return
		}

		switch m := m.(type) {
		case *proto.Connect:
			rc := proto.RetCodeAccepted

			if m.ProtocolName != "MQIsdp" ||
				m.ProtocolVersion != 3 {
				log.Print("reader: reject connection from ", m.ProtocolName, " version ", m.ProtocolVersion)
				rc = proto.RetCodeUnacceptableProtocolVersion
			}

			if len(m.ClientId) < 1 || len(m.ClientId) > 23 {
				rc = proto.RetCodeIdentifierRejected
			}
			c.clientid = m.ClientId
			if existing := c.add(); existing != nil {
				disconnect := &proto.Disconnect{}
				r := existing.submitSync(disconnect)
				r.wait()
				existing.del()
			}
			c.add()

			// TODO: Last will

			connack := &proto.ConnAck{
				ReturnCode: rc,
			}
			c.submit(connack)

			// close connection if it was a bad connect
			if rc != proto.RetCodeAccepted {
				return
			}

		case *proto.Publish:
			// TODO: Proper QoS support
			if m.Header.QosLevel != proto.QosAtMostOnce {
				log.Printf("reader: no support for QoS %v yet", m.Header.QosLevel)
				return
			}
			if isWildcard(m.TopicName) {
				log.Print("reader: ignoring PUBLISH with wildcard topic ", m.TopicName)
			} else {
				c.svr.subs.submit(c, m)
			}
			c.submit(&proto.PubAck{MessageId: m.MessageId})

		case *proto.PingReq:
			c.submit(&proto.PingResp{})

		case *proto.Subscribe:
			suback := &proto.SubAck{
				MessageId: m.MessageId,
				TopicsQos: make([]proto.QosLevel, len(m.Topics)),
			}
			for i, tq := range m.Topics {
				// TODO: Handle varying QoS correctly
				c.svr.subs.add(tq.Topic, c)
				suback.TopicsQos[i] = proto.QosAtLeastOnce
			}
			c.submit(suback)

			// Process retained messages.
			for _, tq := range m.Topics {
			if isWildcard(tq.Topic) {
				// TODO: wildcard retains
			} else {
				c.svr.subs.sendRetain(tq.Topic, c)
			}
			}

		case *proto.Unsubscribe:
			for _, t := range m.Topics {
				c.svr.subs.unsub(t, c)
			}
			ack := &proto.UnsubAck{MessageId: m.MessageId}
			c.submit(ack)

		case *proto.Disconnect:
			return

		default:
			log.Printf("reader: unknown msg type %T", m)
			return
		}
	}
}

func (c *connection) writer() {

	// Close connection on exit in order to cause reader to exit.
	defer func() {
		c.conn.Close()
		c.del()
		c.svr.subs.unsubAll(c)
		log.Print("writer: done")
	}()

	for job := range c.jobs {
		if job.m == nil {
			log.Print("writer: no message?")
			return
		}

		// TODO: write timeout
		err := job.m.Encode(c.conn)
		if job.r != nil {
			// notifiy the sender that this message is sent
			close(job.r)
		}
		if err != nil {
			log.Print("writer: ", err)
			return
		}

		if _, ok := job.m.(*proto.Disconnect); ok {
			log.Print("writer: sent disconnect message")
			return
		}
	}
}

// header is used to initialize a proto.Header when the zero value
// is not correct. The zero value of proto.Header is
// the equivalent of header(dupFalse, proto.QosAtMostOnce, retainFalse)
// and is correct for most messages.
func header(d dupFlag, q proto.QosLevel, r retainFlag) proto.Header {
	return proto.Header{
		DupFlag: bool(d), QosLevel: q, Retain: bool(r),
	}
}

type retainFlag bool
type dupFlag bool

const (
	retainFalse retainFlag = false
	retainTrue             = true
	dupFalse    dupFlag    = false
	dupTrue                = true
)

func isWildcard(topic string) bool {
	if strings.Contains(topic, "#") || strings.Contains(topic, "+") {
		return true
	}
	return false
}

func (w wild) valid() bool {
	for i, part := range w.wild {
		// catch things like finance#
		if isWildcard(part) && len(part) != 1 {
			return false
		}
		// # can only occur as the last part
		if part == "#" && i != len(w.wild)-1 {
			return false
		}
	}
	return true
}
<|endoftext|>