<commit_before>// Create Json Streams in Go.
// Chanson makes easy to fetch data from a channel and encode it.
// It is not an encoder it self, by default it relies on json.Encoder but its flexible enough to let you use whatever you want.
package chanson

import (
	"encoding/json"
	"io"
	"strconv"
)

type Chanson struct {
	w   io.Writer
	enc *json.Encoder
}

// Value is the types that functions like Array.Push() or Object.Set() can accepts as values.
// Custom Value types are:
//	- func(Array)
//	- func(Object)
//	- func(io.Writer)
// If Value is none of the above, it will be encoded using json.Encoder
type Value interface{}

// New returns a new json stream.
// The stream will use w for write the output
func New(w io.Writer) Chanson {
	cs := Chanson{
		w:   w,
		enc: json.NewEncoder(w),
	}
	return cs
}

// Object will execute the callback inside an object context
// this is: "{" f() "}"
func (cs Chanson) Object(f func(Object)) {
	cs.w.Write([]byte("{"))
	if f != nil {
		f(Object{cs: &cs, empty: true})
	}
	cs.w.Write([]byte("}"))
}

// Object will execute the callback inside an array context
// this is: "[" f() "]"
func (cs Chanson) Array(f func(Array)) {
	cs.w.Write([]byte("["))
	if f != nil {
		f(newArray(&cs))
	}
	cs.w.Write([]byte("]"))
}

type Object struct {
	cs    *Chanson
	empty bool
}

// Sets an element into the object
func (obj *Object) Set(key string, val Value) {
	if !obj.empty {
		obj.cs.w.Write([]byte(","))
	} else {
		obj.empty = false
	}

	obj.cs.w.Write([]byte(strconv.Quote(key)))
	obj.cs.w.Write([]byte(":"))
	handleValue(obj.cs, val)
}

type Array struct {
	cs    *Chanson
	empty bool
}

func newArray(cs *Chanson) Array {
	return Array{cs: cs, empty: true}
}

// Pushes an item into the array
func (a *Array) Push(val Value) {
	if !a.empty {
		a.cs.w.Write([]byte(","))
	} else {
		a.empty = false
	}

	handleValue(a.cs, val)
}

func handleValue(cs *Chanson, val Value) {
	switch t := val.(type) {
	case func(Array):
		cs.Array(t)
	case func(Object):
		cs.Object(t)
	case func(io.Writer):
		t(cs.w)
	default:
		err := cs.enc.Encode(val)
		if err != nil {
			//TODO: should panic?!
			cs.w.Write([]byte("null"))
		}
	}
}
<commit_msg>Pass cs by value rather than by reference<commit_after>// Create Json Streams in Go.
// Chanson makes easy to fetch data from a channel and encode it.
// It is not an encoder it self, by default it relies on json.Encoder but its flexible enough to let you use whatever you want.
package chanson

import (
	"encoding/json"
	"io"
	"strconv"
)

type Chanson struct {
	w   io.Writer
	enc *json.Encoder
}

// Value is the types that functions like Array.Push() or Object.Set() can accepts as values.
// Custom Value types are:
//	- func(Array)
//	- func(Object)
//	- func(io.Writer)
// If Value is none of the above, it will be encoded using json.Encoder
type Value interface{}

// New returns a new json stream.
// The stream will use w for write the output
func New(w io.Writer) Chanson {
	cs := Chanson{
		w:   w,
		enc: json.NewEncoder(w),
	}
	return cs
}

// Object will execute the callback inside an object context
// this is: "{" f() "}"
func (cs Chanson) Object(f func(Object)) {
	cs.w.Write([]byte("{"))
	if f != nil {
		f(Object{cs: &cs, empty: true})
	}
	cs.w.Write([]byte("}"))
}

// Object will execute the callback inside an array context
// this is: "[" f() "]"
func (cs Chanson) Array(f func(Array)) {
	cs.w.Write([]byte("["))
	if f != nil {
		f(newArray(&cs))
	}
	cs.w.Write([]byte("]"))
}

type Object struct {
	cs    *Chanson
	empty bool
}

// Sets an element into the object
func (obj *Object) Set(key string, val Value) {
	if !obj.empty {
		obj.cs.w.Write([]byte(","))
	} else {
		obj.empty = false
	}

	obj.cs.w.Write([]byte(strconv.Quote(key)))
	obj.cs.w.Write([]byte(":"))
	handleValue(*obj.cs, val)
}

type Array struct {
	cs    *Chanson
	empty bool
}

func newArray(cs *Chanson) Array {
	return Array{cs: cs, empty: true}
}

// Pushes an item into the array
func (a *Array) Push(val Value) {
	if !a.empty {
		a.cs.w.Write([]byte(","))
	} else {
		a.empty = false
	}

	handleValue(*a.cs, val)
}

func handleValue(cs Chanson, val Value) {
	switch t := val.(type) {
	case func(Array):
		cs.Array(t)
	case func(Object):
		cs.Object(t)
	case func(io.Writer):
		t(cs.w)
	default:
		err := cs.enc.Encode(val)
		if err != nil {
			//TODO: should panic?!
			cs.w.Write([]byte("null"))
		}
	}
}
<|endoftext|>


<commit_before>package qb

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestConstraints(t *testing.T) {

	assert.Equal(t, NotNull(), Constraint{"NOT NULL"})
	assert.Equal(t, Default(5), Constraint{"DEFAULT '5'"})
	assert.Equal(t, Default("-"), Constraint{"DEFAULT '-'"})
	assert.Equal(t, Unique(), Constraint{"UNIQUE"})
	assert.Equal(t, Unique("email", "name"), Constraint{"UNIQUE(email, name)"})
	//assert.Equal(t, Primary(), Constraint{"PRIMARY KEY"})
	//assert.Equal(t, Primary("email", "password"), Constraint{"PRIMARY KEY(email, password)"})
	//assert.Equal(t, Foreign("user_id", "profile", "user_id"), Constraint{"FOREIGN KEY (user_id) REFERENCES profile(user_id)"})
}
<commit_msg>cleanup constraint tests<commit_after>package qb

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestConstraints(t *testing.T) {
	assert.Equal(t, NotNull(), Constraint{"NOT NULL"})
	assert.Equal(t, Default(5), Constraint{"DEFAULT '5'"})
	assert.Equal(t, Default("-"), Constraint{"DEFAULT '-'"})
	assert.Equal(t, Unique(), Constraint{"UNIQUE"})
	assert.Equal(t, Unique("email", "name"), Constraint{"UNIQUE(email, name)"})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/joshsoftware/curem/config"
)

func TestGetContactsHandler(t *testing.T) {
	ts := httptest.NewServer(r)
	defer ts.Close()
	c, err := NewContact(
		"Encom Inc.",
		"Flynn",
		"flynn@encom.com",
		"",
		"",
		"USA",
	)
	if err != nil {
		t.Errorf("%s", err)
	}
	x := []contact{*c}
	y, err := json.Marshal(x)
	resp, err := http.Get(ts.URL + "/contacts")
	if err != nil {
		t.Errorf("%s", err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Errorf("%s", err)
	}
	bodystring := strings.TrimSpace(string(body))
	if bodystring != string(y) {
		t.Errorf("expected %s, but got %s", string(y), bodystring)
	}
	err = config.ContactsCollection.DropCollection()
	if err != nil {
		t.Errorf("%s", err)
	}
}
<commit_msg>Add test for getContactHandler<commit_after>package main

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/joshsoftware/curem/config"
)

func TestGetContactsHandler(t *testing.T) {
	ts := httptest.NewServer(r)
	defer ts.Close()
	c, err := NewContact(
		"Encom Inc.",
		"Flynn",
		"flynn@encom.com",
		"",
		"",
		"USA",
	)
	if err != nil {
		t.Errorf("%s", err)
	}
	x := []contact{*c}
	y, err := json.Marshal(x)
	resp, err := http.Get(ts.URL + "/contacts")
	if err != nil {
		t.Errorf("%s", err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Errorf("%s", err)
	}
	bodystring := strings.TrimSpace(string(body))
	if bodystring != string(y) {
		t.Errorf("expected %s, but got %s", string(y), bodystring)
	}
	err = config.ContactsCollection.DropCollection()
	if err != nil {
		t.Errorf("%s", err)
	}
}

func TestGetContactHandler(t *testing.T) {
	ts := httptest.NewServer(r)
	defer ts.Close()
	c, err := NewContact(
		"Encom Inc.",
		"Flynn",
		"flynn@encom.com",
		"",
		"",
		"USA",
	)
	if err != nil {
		t.Errorf("%s", err)
	}
	y, err := json.Marshal(c)
	resp, err := http.Get(ts.URL + "/contacts/" + c.Slug)
	if err != nil {
		t.Errorf("%s", err)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Errorf("%s", err)
	}
	bodystring := strings.TrimSpace(string(body))
	if bodystring != string(y) {
		t.Errorf("expected %s, but got %s", string(y), bodystring)
	}
	err = config.ContactsCollection.DropCollection()
	if err != nil {
		t.Errorf("%s", err)
	}
}
<|endoftext|>


<commit_before>package container

import (
	"bytes"
	"compress/flate"
	"io"
)

type Codec string

const (
	Null    Codec = "null"
	Deflate Codec = "deflate"
	Snappy  Codec = "snappy"
)

type CloseableResettableWriter interface {
	Close() error
	Reset(io.Writer)
}

type Writer struct {
	writer           io.Writer
	syncMarker       [16]byte
	codec            Codec
	recordsPerBlock  int64
	blockBuffer      *bytes.Buffer
	compressedWriter io.Writer
	nextBlockRecords int64
	headerWritten    bool
}

func NewWriter(writer io.Writer, codec Codec, recordsPerBlock int64) (*Writer, error) {
	blockBytes := make([]byte, 0)
	blockBuffer := bytes.NewBuffer(blockBytes)

	avroWriter := &Writer{
		writer:          writer,
		syncMarker:      [16]byte{'g', 'o', 'g', 'e', 'n', 'a', 'v', 'r', 'o', 'm', 'a', 'g', 'i', 'c', '1', '0'},
		codec:           codec,
		recordsPerBlock: recordsPerBlock,
		blockBuffer:     blockBuffer,
		headerWritten:   false,
	}
	var err error
	if codec == Deflate {
		avroWriter.compressedWriter, err = flate.NewWriter(avroWriter.blockBuffer, flate.DefaultCompression)
		if err != nil {
			return nil, err
		}
	} else if codec == Snappy {
		avroWriter.compressedWriter = newSnappyWriter(avroWriter.blockBuffer)
	} else if codec == Null {
		avroWriter.compressedWriter = avroWriter.blockBuffer
	}

	return avroWriter, nil
}

func (avroWriter *Writer) writeHeader(schema string) error {
	header := &AvroContainerHeader{
		Magic: [4]byte{'O', 'b', 'j', 1},
		Meta: map[string][]byte{
			"avro.schema": []byte(schema),
			"avro.codec":  []byte(avroWriter.codec),
		},
		Sync: avroWriter.syncMarker,
	}
	return header.Serialize(avroWriter.writer)
}

func (avroWriter *Writer) WriteRecord(record AvroRecord) error {
	var err error
	// Lazily write the header when the first record is written
	if !avroWriter.headerWritten {
		avroWriter.headerWritten = true
		err = avroWriter.writeHeader(record.Schema())
		if err != nil {
			return err
		}
	}
	// Serialize the new record into the compressed writer
	err = record.Serialize(avroWriter.compressedWriter)
	if err != nil {
		return err
	}
	avroWriter.nextBlockRecords += 1

	// If the block if full, flush and reset the compressed writer,
	// write the header and the block contents
	if avroWriter.nextBlockRecords >= avroWriter.recordsPerBlock {
		return avroWriter.Flush()
	}

	return nil
}

func (avroWriter *Writer) Flush() error {
	// Write out all of the buffered records as a new block
	// Must be called before closing to ensure the last block is written
	if fwWriter, ok := avroWriter.compressedWriter.(CloseableResettableWriter); ok {
		fwWriter.Close()
		fwWriter.Reset(avroWriter.blockBuffer)
	}

	block := &AvroContainerBlock{
		NumRecords:  avroWriter.nextBlockRecords,
		RecordBytes: avroWriter.blockBuffer.Bytes(),
		Sync:        avroWriter.syncMarker,
	}
	err := block.Serialize(avroWriter.writer)
	if err != nil {
		return err
	}

	avroWriter.blockBuffer.Reset()
	avroWriter.nextBlockRecords = 0

	return nil
}
<commit_msg>Add godocs for container package<commit_after>// Container provides a Writer which is capable of serializing gogen-avro structs and writing them in the Avro Object Container File (OCF) format
package container

import (
	"bytes"
	"compress/flate"
	"io"
)

/*
  A Codec specifies how the blocks within a container file should be compressed.
*/
type Codec string

const (
	// No compression
	Null    Codec = "null"
	// Deflate compression
	Deflate Codec = "deflate"
	// Snappy compression
	Snappy  Codec = "snappy"
)

type CloseableResettableWriter interface {
	Close() error
	Reset(io.Writer)
}

/*
  Writer wraps an io.Writer and writes the file and block-level framing required for an OCF file
*/
type Writer struct {
	writer           io.Writer
	syncMarker       [16]byte
	codec            Codec
	recordsPerBlock  int64
	blockBuffer      *bytes.Buffer
	compressedWriter io.Writer
	nextBlockRecords int64
	headerWritten    bool
}

/*
  Create a new Writer wrapping the provided io.Writer with the given Codec and number of records per block.
  The Writer will lazily write the container file header when WriteRecord is called the first time.
  You must call Flush on the Writer before closing the underlying io.Writer, to ensure the final block is written.
*/
func NewWriter(writer io.Writer, codec Codec, recordsPerBlock int64) (*Writer, error) {
	blockBytes := make([]byte, 0)
	blockBuffer := bytes.NewBuffer(blockBytes)

	avroWriter := &Writer{
		writer:          writer,
		syncMarker:      [16]byte{'g', 'o', 'g', 'e', 'n', 'a', 'v', 'r', 'o', 'm', 'a', 'g', 'i', 'c', '1', '0'},
		codec:           codec,
		recordsPerBlock: recordsPerBlock,
		blockBuffer:     blockBuffer,
		headerWritten:   false,
	}
	var err error
	if codec == Deflate {
		avroWriter.compressedWriter, err = flate.NewWriter(avroWriter.blockBuffer, flate.DefaultCompression)
		if err != nil {
			return nil, err
		}
	} else if codec == Snappy {
		avroWriter.compressedWriter = newSnappyWriter(avroWriter.blockBuffer)
	} else if codec == Null {
		avroWriter.compressedWriter = avroWriter.blockBuffer
	}

	return avroWriter, nil
}

func (avroWriter *Writer) writeHeader(schema string) error {
	header := &AvroContainerHeader{
		Magic: [4]byte{'O', 'b', 'j', 1},
		Meta: map[string][]byte{
			"avro.schema": []byte(schema),
			"avro.codec":  []byte(avroWriter.codec),
		},
		Sync: avroWriter.syncMarker,
	}
	return header.Serialize(avroWriter.writer)
}

/*
  Write an AvroRecord to the container file. All gogen-avro generated structs
  fulfill the AvroRecord interface. Note that all records in a given container file
  must be of the same Avro type.
*/
func (avroWriter *Writer) WriteRecord(record AvroRecord) error {
	var err error
	// Lazily write the header when the first record is written
	if !avroWriter.headerWritten {
		avroWriter.headerWritten = true
		err = avroWriter.writeHeader(record.Schema())
		if err != nil {
			return err
		}
	}
	// Serialize the new record into the compressed writer
	err = record.Serialize(avroWriter.compressedWriter)
	if err != nil {
		return err
	}
	avroWriter.nextBlockRecords += 1

	// If the block if full, flush and reset the compressed writer,
	// write the header and the block contents
	if avroWriter.nextBlockRecords >= avroWriter.recordsPerBlock {
		return avroWriter.Flush()
	}

	return nil
}

/*
  Write the current block to the file, even if it hasn't been filled.
  This must be called before the underlying io.Writer is closed.
*/
func (avroWriter *Writer) Flush() error {
	// Write out all of the buffered records as a new block
	// Must be called before closing to ensure the last block is written
	if fwWriter, ok := avroWriter.compressedWriter.(CloseableResettableWriter); ok {
		fwWriter.Close()
		fwWriter.Reset(avroWriter.blockBuffer)
	}

	block := &AvroContainerBlock{
		NumRecords:  avroWriter.nextBlockRecords,
		RecordBytes: avroWriter.blockBuffer.Bytes(),
		Sync:        avroWriter.syncMarker,
	}
	err := block.Serialize(avroWriter.writer)
	if err != nil {
		return err
	}

	avroWriter.blockBuffer.Reset()
	avroWriter.nextBlockRecords = 0

	return nil
}
<|endoftext|>


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

package mem

import (
	"encoding/binary"
	"fmt"
	"math/rand"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/dgraph-io/badger/value"
	"github.com/dgraph-io/badger/y"
)

func extract(m *Table) ([]string, []string) {
	var keys, vals []string
	it := m.NewIterator()
	for it.SeekToFirst(); it.Valid(); it.Next() {
		keys = append(keys, string(it.Key()))
		v := it.Value()
		if (value.BitDelete & v[0]) != 0 {
			vals = append(vals, "DEL")
		} else {
			vals = append(vals, string(v[1:]))
		}
	}
	return keys, vals
}

func TestBasic(t *testing.T) {
	m := NewTable()
	require.NotNil(t, m)
	m.Put([]byte("somekey"), []byte("hohoho"), 0)
	m.Put([]byte("somekey"), []byte("hahaha"), 0)
	k, v := extract(m)
	require.EqualValues(t, []string{"somekey"}, k)
	require.EqualValues(t, []string{"hahaha"}, v)

	m.Put([]byte("akey"), nil, value.BitDelete)
	m.Put([]byte("somekey"), nil, value.BitDelete)
	k, v = extract(m)
	require.EqualValues(t, []string{"akey", "somekey"}, k)
	require.EqualValues(t, []string{"DEL", "DEL"}, v)

	m.Put([]byte("somekey"), []byte("yes"), 0)
	k, v = extract(m)
	require.EqualValues(t, []string{"akey", "somekey"}, k)
	require.EqualValues(t, []string{"DEL", "yes"}, v)
}

func TestSize(t *testing.T) {
	m := NewTable()
	for i := 0; i < 10000; i++ {
		m.Put([]byte(fmt.Sprintf("k%05d", i)), []byte(fmt.Sprintf("v%05d", i)), 0)
	}
	expected := 10000 * (6 + 6 + 1)
	require.InEpsilon(t, expected, m.Size(), 0.1)
}

func TestConcurrentWrite(t *testing.T) {
	tbl := NewTable()
	n := 20
	m := 50000
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			for j := 0; j < m; j++ {
				tbl.Put([]byte(fmt.Sprintf("k%05d_%08d", i, j)),
					[]byte(fmt.Sprintf("v%05d_%08d", i, j)), 5)
			}
		}(i)
	}
	wg.Wait()
	it := tbl.NewIterator()
	it.SeekToFirst()
	for i := 0; i < n; i++ {
		for j := 0; j < m; j++ {
			require.True(t, it.Valid(), "%d %d", i, j)
			k, v := it.KeyValue()
			require.EqualValues(t, fmt.Sprintf("k%05d_%08d", i, j), string(k))
			require.EqualValues(t, fmt.Sprintf("v%05d_%08d", i, j), string(v[1:]))
			it.Next()
		}
	}
}

func TestMergeIterator(t *testing.T) {
	m := NewTable()
	it := m.NewIterator()
	mergeIt := y.NewMergeIterator([]y.Iterator{it})
	require.False(t, mergeIt.Valid())
}

func BenchmarkAdd(b *testing.B) {
	m := NewTable()
	for i := 0; i < b.N; i++ {
		m.Put([]byte(fmt.Sprintf("k%09d", i)), []byte(fmt.Sprintf("v%09d", i)), 0)
	}
}

func randomKey() []byte {
	b := make([]byte, 8)
	key := rand.Uint32()
	key2 := rand.Uint32()
	binary.LittleEndian.PutUint32(b, key)
	binary.LittleEndian.PutUint32(b[4:], key2)
	return b
}

func BenchmarkReadWrite(b *testing.B) {
	for i := 0; i <= 10; i++ {
		readFrac := float32(i) / 10.0
		b.Run(fmt.Sprintf("frac_%d", i), func(b *testing.B) {
			list := NewTable()
			b.ResetTimer()
			var count int
			b.RunParallel(func(pb *testing.PB) {
				for pb.Next() {
					if rand.Float32() < readFrac {
						if list.Get(randomKey()) != nil {
							count++
						}
					} else {
						newVal := make([]byte, 10)
						list.Put(randomKey(), newVal, 0)
					}
				}
			})
		})
	}
}
<commit_msg>Minor update to TestConcurrentWrite<commit_after>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package mem

import (
	"encoding/binary"
	"fmt"
	"math/rand"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/dgraph-io/badger/value"
	"github.com/dgraph-io/badger/y"
)

func extract(m *Table) ([]string, []string) {
	var keys, vals []string
	it := m.NewIterator()
	for it.SeekToFirst(); it.Valid(); it.Next() {
		keys = append(keys, string(it.Key()))
		v := it.Value()
		if (value.BitDelete & v[0]) != 0 {
			vals = append(vals, "DEL")
		} else {
			vals = append(vals, string(v[1:]))
		}
	}
	return keys, vals
}

func TestBasic(t *testing.T) {
	m := NewTable()
	require.NotNil(t, m)
	m.Put([]byte("somekey"), []byte("hohoho"), 0)
	m.Put([]byte("somekey"), []byte("hahaha"), 0)
	k, v := extract(m)
	require.EqualValues(t, []string{"somekey"}, k)
	require.EqualValues(t, []string{"hahaha"}, v)

	m.Put([]byte("akey"), nil, value.BitDelete)
	m.Put([]byte("somekey"), nil, value.BitDelete)
	k, v = extract(m)
	require.EqualValues(t, []string{"akey", "somekey"}, k)
	require.EqualValues(t, []string{"DEL", "DEL"}, v)

	m.Put([]byte("somekey"), []byte("yes"), 0)
	k, v = extract(m)
	require.EqualValues(t, []string{"akey", "somekey"}, k)
	require.EqualValues(t, []string{"DEL", "yes"}, v)
}

func TestSize(t *testing.T) {
	m := NewTable()
	for i := 0; i < 10000; i++ {
		m.Put([]byte(fmt.Sprintf("k%05d", i)), []byte(fmt.Sprintf("v%05d", i)), 0)
	}
	expected := 10000 * (6 + 6 + 1)
	require.InEpsilon(t, expected, m.Size(), 0.1)
}

func TestConcurrentWrite(t *testing.T) {
	tbl := NewTable()
	n := 20
	m := 50000
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			for j := 0; j < m; j++ {
				tbl.Put([]byte(fmt.Sprintf("k%05d_%08d", i, j)),
					[]byte(fmt.Sprintf("v%05d_%08d", i, j)), 5)
			}
		}(i)
	}
	wg.Wait()
	it := tbl.NewIterator()
	it.SeekToFirst()
	for i := 0; i < n; i++ {
		for j := 0; j < m; j++ {
			require.True(t, it.Valid(), "%d %d", i, j)
			k, v := it.KeyValue()
			require.EqualValues(t, fmt.Sprintf("k%05d_%08d", i, j), string(k))
			require.EqualValues(t, fmt.Sprintf("v%05d_%08d", i, j), string(v[1:]))
			it.Next()
		}
	}
	require.False(t, it.Valid())
}

func TestMergeIterator(t *testing.T) {
	m := NewTable()
	it := m.NewIterator()
	mergeIt := y.NewMergeIterator([]y.Iterator{it})
	require.False(t, mergeIt.Valid())
}

func BenchmarkAdd(b *testing.B) {
	m := NewTable()
	for i := 0; i < b.N; i++ {
		m.Put([]byte(fmt.Sprintf("k%09d", i)), []byte(fmt.Sprintf("v%09d", i)), 0)
	}
}

func randomKey() []byte {
	b := make([]byte, 8)
	key := rand.Uint32()
	key2 := rand.Uint32()
	binary.LittleEndian.PutUint32(b, key)
	binary.LittleEndian.PutUint32(b[4:], key2)
	return b
}

func BenchmarkReadWrite(b *testing.B) {
	for i := 0; i <= 10; i++ {
		readFrac := float32(i) / 10.0
		b.Run(fmt.Sprintf("frac_%d", i), func(b *testing.B) {
			list := NewTable()
			b.ResetTimer()
			var count int
			b.RunParallel(func(pb *testing.PB) {
				for pb.Next() {
					if rand.Float32() < readFrac {
						if list.Get(randomKey()) != nil {
							count++
						}
					} else {
						newVal := make([]byte, 10)
						list.Put(randomKey(), newVal, 0)
					}
				}
			})
		})
	}
}
<|endoftext|>


<commit_before>package wkhtmltopdf

import (
	"reflect"
	"testing"
)

func TestGlobalOptions(t *testing.T) {

	testcases := []struct {
		Options []Option
		Args    []string
	}{
		{[]Option{}, []string{}},
		{[]Option{NoCollate()}, []string{"--no-collate"}},
		{[]Option{DPI(300), Grayscale()}, []string{"--dpi", "300", "--grayscale"}},
		{[]Option{ImageDPI(500), ImageQuality(60)}, []string{"--image-dpi", "500", "--image-quality", "60"}},
		{[]Option{LowQuality(), CookieJar("testpath")}, []string{"--low-quality", "--cookie-jar", "testpath"}},
		{[]Option{MarginBottom("1cm"), MarginTop("2cm")}, []string{"--margin-bottom", "1cm", "--margin-top", "2cm"}},
		{[]Option{MarginLeft("15mm"), MarginRight("1.5cm")}, []string{"--margin-left", "15mm", "--margin-right", "1.5cm"}},
		{[]Option{Landscape(), PageHeight("20cm")}, []string{"--orientation", "landscape", "--page-height", "20cm"}},
		{[]Option{PageWidth("35cm")}, []string{"--page-width", "35cm"}},
		{[]Option{NoPDFCompression(), Quiet()}, []string{"--no-pdf-compression", "--quiet"}},
		{[]Option{Title("An excellent pdf")}, []string{"--title", "An excellent pdf"}},
		{[]Option{PageSize("A4")}, []string{"--page-size", "A4"}},
	}

	for _, tc := range testcases {
		doc := NewDocument(tc.Options...)
		if !reflect.DeepEqual(doc.options, tc.Args) {
			t.Errorf("Wrong arguments created. Expected: %v, Got: %v", tc.Args, doc.options)
		}
	}
}

func TestPageOptions(t *testing.T) {

	testcases := []struct {
		Options []PageOption
		Args    []string
	}{
		{[]PageOption{Allow("test/")}, []string{"--allow", "test/"}},
		{[]PageOption{NoBackground()}, []string{"--no-background"}},
		{[]PageOption{BypassProxy("test.com")}, []string{"--bypass-proxy-for", "test.com"}},
		{[]PageOption{CacheDir("cache/")}, []string{"--cache-dir", "cache/"}},
		{[]PageOption{CheckboxCheckedSVG("path")}, []string{"--checkbox-checked-svg", "path"}},
		{[]PageOption{CheckboxSVG("path")}, []string{"--checkbox-svg", "path"}},
		{[]PageOption{Cookie("name", "value")}, []string{"--cookie", "name", "value"}},
		{[]PageOption{CustomHeader("name", "value")}, []string{"--custom-header", "name", "value"}},
		{[]PageOption{CustomHeaderPropagation()}, []string{"--custom-header-propagation"}},
		{[]PageOption{NoCustomHeaderPropagation()}, []string{"--no-custom-header-propagation"}},
		{[]PageOption{DefaultHeader()}, []string{"--default-header"}},
		{[]PageOption{Encoding("UTF8")}, []string{"--encoding", "UTF8"}},
		{[]PageOption{DisableExternalLinks()}, []string{"--disable-external-links"}},
		{[]PageOption{EnableExternalLinks()}, []string{"--enable-external-links"}},
		{[]PageOption{DisableForms()}, []string{"--disable-forms"}},
		{[]PageOption{EnableForms()}, []string{"--enable-forms"}},
		{[]PageOption{Images()}, []string{"--images"}},
		{[]PageOption{NoImages()}, []string{"--no-images"}},
		{[]PageOption{DisableInternalLinks()}, []string{"--disable-internal-links"}},
		{[]PageOption{EnableInternalLinks()}, []string{"--enable-internal-links"}},
	}

	for _, tc := range testcases {
		pg := NewPage("", tc.Options...)
		if !reflect.DeepEqual(pg.options, tc.Args) {
			t.Errorf("Wrong arguments created. Expected: %v, Got: %v", tc.Args, pg.options)
		}
	}

}
<commit_msg>Added missing test for background<commit_after>package wkhtmltopdf

import (
	"reflect"
	"testing"
)

func TestGlobalOptions(t *testing.T) {

	testcases := []struct {
		Options []Option
		Args    []string
	}{
		{[]Option{}, []string{}},
		{[]Option{NoCollate()}, []string{"--no-collate"}},
		{[]Option{DPI(300), Grayscale()}, []string{"--dpi", "300", "--grayscale"}},
		{[]Option{ImageDPI(500), ImageQuality(60)}, []string{"--image-dpi", "500", "--image-quality", "60"}},
		{[]Option{LowQuality(), CookieJar("testpath")}, []string{"--low-quality", "--cookie-jar", "testpath"}},
		{[]Option{MarginBottom("1cm"), MarginTop("2cm")}, []string{"--margin-bottom", "1cm", "--margin-top", "2cm"}},
		{[]Option{MarginLeft("15mm"), MarginRight("1.5cm")}, []string{"--margin-left", "15mm", "--margin-right", "1.5cm"}},
		{[]Option{Landscape(), PageHeight("20cm")}, []string{"--orientation", "landscape", "--page-height", "20cm"}},
		{[]Option{PageWidth("35cm")}, []string{"--page-width", "35cm"}},
		{[]Option{NoPDFCompression(), Quiet()}, []string{"--no-pdf-compression", "--quiet"}},
		{[]Option{Title("An excellent pdf")}, []string{"--title", "An excellent pdf"}},
		{[]Option{PageSize("A4")}, []string{"--page-size", "A4"}},
	}

	for _, tc := range testcases {
		doc := NewDocument(tc.Options...)
		if !reflect.DeepEqual(doc.options, tc.Args) {
			t.Errorf("Wrong arguments created. Expected: %v, Got: %v", tc.Args, doc.options)
		}
	}
}

func TestPageOptions(t *testing.T) {

	testcases := []struct {
		Options []PageOption
		Args    []string
	}{
		{[]PageOption{Allow("test/")}, []string{"--allow", "test/"}},
		{[]PageOption{Background()}, []string{"--background"}},
		{[]PageOption{NoBackground()}, []string{"--no-background"}},
		{[]PageOption{BypassProxy("test.com")}, []string{"--bypass-proxy-for", "test.com"}},
		{[]PageOption{CacheDir("cache/")}, []string{"--cache-dir", "cache/"}},
		{[]PageOption{CheckboxCheckedSVG("path")}, []string{"--checkbox-checked-svg", "path"}},
		{[]PageOption{CheckboxSVG("path")}, []string{"--checkbox-svg", "path"}},
		{[]PageOption{Cookie("name", "value")}, []string{"--cookie", "name", "value"}},
		{[]PageOption{CustomHeader("name", "value")}, []string{"--custom-header", "name", "value"}},
		{[]PageOption{CustomHeaderPropagation()}, []string{"--custom-header-propagation"}},
		{[]PageOption{NoCustomHeaderPropagation()}, []string{"--no-custom-header-propagation"}},
		{[]PageOption{DefaultHeader()}, []string{"--default-header"}},
		{[]PageOption{Encoding("UTF8")}, []string{"--encoding", "UTF8"}},
		{[]PageOption{DisableExternalLinks()}, []string{"--disable-external-links"}},
		{[]PageOption{EnableExternalLinks()}, []string{"--enable-external-links"}},
		{[]PageOption{DisableForms()}, []string{"--disable-forms"}},
		{[]PageOption{EnableForms()}, []string{"--enable-forms"}},
		{[]PageOption{Images()}, []string{"--images"}},
		{[]PageOption{NoImages()}, []string{"--no-images"}},
		{[]PageOption{DisableInternalLinks()}, []string{"--disable-internal-links"}},
		{[]PageOption{EnableInternalLinks()}, []string{"--enable-internal-links"}},
	}

	for _, tc := range testcases {
		pg := NewPage("", tc.Options...)
		if !reflect.DeepEqual(pg.options, tc.Args) {
			t.Errorf("Wrong arguments created. Expected: %v, Got: %v", tc.Args, pg.options)
		}
	}

}
<|endoftext|>


<commit_before>package orbital

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gregbuehler/orbital/orbital/database"
)

// Feed object
type Feed struct {
	ID          string
	`json:"ID"
	 description:"The ID of the Feed"
	`
	Title       string       `json:"Title,required" description:"A descriptive name for the Feed"`
	Private     string       `json:"Private" description:"Whether or not the Feed is private to the creator of the Feed"`
	Tags        []string     `json:"Tags" description:"Tagged metadata about the Feed"`
	Description string       `json:"Description" description:"A description of the Feed."`
	FeedURL     string       `json:"Feed" description:"The URL of the Feed"`
	Status      string       `json:"Status" description:"Feed 'live' or 'frozen'"`
	Updated     string       `json:"Updated" description:"The time at which this Feed was last updated"`
	Created     string       `json:"Created" description:"The date the Feed was created"`
	Creator     string       `json:"Creator" description:"A URL referencing the creator of the Feed"`
	Version     string       `json:"Version" description:"Version of the data format Feed returned"`
	Location    string       `json:"Location" description:"Location metadata"`
	Website     string       `json:"Website" description:"The URL of a website which is relevant to this Feed"`
	Icon        string       `json:"Icon" description:"The URL of an icon which is relevant to this Feed"`
	UserLogin   string       `json:"User Login" description:"The user who owns the Feed"`
	Datastreams []Datastream `json:"Datastreams" description:"A collection of the Datastreams in this Feed"`
}

// FeedSimpleListHandler lists feeds in summary format
func FeedSimpleListHandler(w http.ResponseWriter, r *http.Request) {
	q := `
		SELECT
			feedid,
			title,
			Private		-- keep Private Private
			Tags,
			Description,
			FeedURL,
			Status,
			Updated,
			Created,
			Creator,
			Version,
			Location,
			Website,
			Icon,
			UserLogin
		FROM
			feeds
		WHERE
			private > 0
		;
	`
	rows, err := database.DBCon.Query(q)
	if err != nil {
		log.Fatal(err)
	}
	for rows.Next() {
		f := Feed{}
		if err := rows.Scan(&f.ID, &f.Title, &f.Private, &f.Tags,&f.Description, &f.FeedURL, &f.Status, &f.Updated, &f.Created, &f.Version, &f.Location, &f.Website, &f.Icon, &f.UserLogin); err != nil {
			log.Fatal(err)
		}
		fmt.Fprintf(w, "%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s\n", f.ID, f.Title, f.Private, f.Tags, f.Description, f.FeedURL, f.Status, f.Updated, f.Created, f.Version, f.Location, f.Website, f.Icon, f.UserLogin)
	}
}

// FeedVerboseListHandler lists feeds in verbose format
func FeedVerboseListHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "FeedVerboseListHandler!\n\n%s", r.RequestURI)
}

// FeedDetailHandler details a feed
func FeedDetailHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "FeedDetailHandler!\n\n%s", r.RequestURI)
}

// FeedUpdateHandler updates details of a feed
func FeedUpdateHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "FeedUpdateHandler!\n\n%s", r.RequestURI)
}
<commit_msg>Revert failed multiline annotation expirament<commit_after>package orbital

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gregbuehler/orbital/orbital/database"
)

// Feed object
type Feed struct {
	ID          string	`json:"ID" description:"The ID of the Feed"`
	Title       string       `json:"Title,required" description:"A descriptive name for the Feed"`
	Private     string       `json:"Private" description:"Whether or not the Feed is private to the creator of the Feed"`
	Tags        []string     `json:"Tags" description:"Tagged metadata about the Feed"`
	Description string       `json:"Description" description:"A description of the Feed."`
	FeedURL     string       `json:"Feed" description:"The URL of the Feed"`
	Status      string       `json:"Status" description:"Feed 'live' or 'frozen'"`
	Updated     string       `json:"Updated" description:"The time at which this Feed was last updated"`
	Created     string       `json:"Created" description:"The date the Feed was created"`
	Creator     string       `json:"Creator" description:"A URL referencing the creator of the Feed"`
	Version     string       `json:"Version" description:"Version of the data format Feed returned"`
	Location    string       `json:"Location" description:"Location metadata"`
	Website     string       `json:"Website" description:"The URL of a website which is relevant to this Feed"`
	Icon        string       `json:"Icon" description:"The URL of an icon which is relevant to this Feed"`
	UserLogin   string       `json:"User Login" description:"The user who owns the Feed"`
	Datastreams []Datastream `json:"Datastreams" description:"A collection of the Datastreams in this Feed"`
}

// FeedSimpleListHandler lists feeds in summary format
func FeedSimpleListHandler(w http.ResponseWriter, r *http.Request) {
	q := `
		SELECT
			feedid,
			title,
			Private		-- keep Private Private
			Tags,
			Description,
			FeedURL,
			Status,
			Updated,
			Created,
			Creator,
			Version,
			Location,
			Website,
			Icon,
			UserLogin
		FROM
			feeds
		WHERE
			private > 0
		;
	`
	rows, err := database.DBCon.Query(q)
	if err != nil {
		log.Fatal(err)
	}
	for rows.Next() {
		f := Feed{}
		if err := rows.Scan(&f.ID, &f.Title, &f.Private, &f.Tags,&f.Description, &f.FeedURL, &f.Status, &f.Updated, &f.Created, &f.Version, &f.Location, &f.Website, &f.Icon, &f.UserLogin); err != nil {
			log.Fatal(err)
		}
		fmt.Fprintf(w, "%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s\n", f.ID, f.Title, f.Private, f.Tags, f.Description, f.FeedURL, f.Status, f.Updated, f.Created, f.Version, f.Location, f.Website, f.Icon, f.UserLogin)
	}
}

// FeedVerboseListHandler lists feeds in verbose format
func FeedVerboseListHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "FeedVerboseListHandler!\n\n%s", r.RequestURI)
}

// FeedDetailHandler details a feed
func FeedDetailHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "FeedDetailHandler!\n\n%s", r.RequestURI)
}

// FeedUpdateHandler updates details of a feed
func FeedUpdateHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "FeedUpdateHandler!\n\n%s", r.RequestURI)
}
<|endoftext|>


<commit_before>package gorp

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// A TypeDefSwitcher is a TypeDeffer except that it won't have its type
// defined on table creation.  Instead, its type will be changed on
// the first migration run.  Useful for foreign keys when the target
// table may be created later than the column.
type TypeDefSwitcher interface {
	// TypeDefSwitch should return the same thing as
	// TypeDeffer.TypeDef.
	TypeDefSwitch() string
}

// A TypeCaster includes TypeDeffer logic but can also return the SQL
// to cast old types to its new type.
type TypeCaster interface {
	// TypeCast should return a string with zero or one '%s'
	// sequences.  If the string contains a '%s', it will be replaced
	// with the old value; otherwise, the return value will simply be
	// used as the type to cast to in the database.
	TypeCast() string
}

type pgJSON []byte

func (j pgJSON) Value() (driver.Value, error) {
	return []byte(j), nil
}

func (j *pgJSON) Scan(src interface{}) error {
	var source []byte
	switch src.(type) {
	case string:
		source = []byte(src.(string))
	case []byte:
		source = src.([]byte)
	default:
		return errors.New("Incompatible type for pgJSON")
	}
	*j = pgJSON(append((*j)[0:0], source...))
	return nil
}

func (j pgJSON) TypeDef() string {
	return "json"
}

type columnLayout struct {
	FieldName  string `json:"field_name"`
	ColumnName string `json:"column_name"`
	TypeDef    string `json:"type_def"`

	// Values that are only used on the new layout, but are
	// unnecessary for old types.
	needsNotNull bool         `json:"-"`
	isPK         bool         `json:"-"`
	hasReference bool         `json:"-"`
	gotype       reflect.Type `json:"-"`
	typeCast     string       `json:"-"`
}

type tableRecord struct {
	ID          int
	Schemaname  string
	Tablename   string
	Layout      pgJSON
	tableLayout []columnLayout `db:"-"`
}

func (t *tableRecord) TableLayout() []columnLayout {
	if t.tableLayout == nil {
		t.tableLayout = []columnLayout{}
		if err := json.Unmarshal([]byte(t.Layout), &t.tableLayout); err != nil {
			panic(err)
		}
	}
	return t.tableLayout
}

func (t *tableRecord) SetTableLayout(l []columnLayout) {
	t.tableLayout = l
	b, err := json.Marshal(l)
	if err != nil {
		panic(err)
	}
	t.Layout = pgJSON(b)
}

func (t *tableRecord) Merge(l []columnLayout) {
	if t.tableLayout == nil {
		t.SetTableLayout(l)
		return
	}
	for _, newCol := range l {
		shouldAppend := true
		for _, oldCol := range t.tableLayout {
			if newCol.ColumnName == oldCol.ColumnName {
				shouldAppend = false
			}
		}
		if shouldAppend {
			t.tableLayout = append(t.tableLayout, newCol)
		}
	}
}

func ptrToVal(t reflect.Type) reflect.Value {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return reflect.New(t)
}

type MigrationManager struct {
	schemaname string
	tablename  string
	dbMap      *DbMap
	oldTables  []*tableRecord
	newTables  []*tableRecord
}

func (m *MigrationManager) layoutFor(t *TableMap) []columnLayout {
	l := make([]columnLayout, 0, len(t.Columns))
	for _, colMap := range t.Columns {
		if colMap.Transient {
			continue
		}
		stype, notNullIgnored := colMap.TypeDefNoNotNull()
		orig := ptrToVal(colMap.origtype).Interface()
		dbValue := ptrToVal(colMap.gotype).Interface()
		cast := "%s"
		typeCaster, ok := orig.(TypeCaster)
		if !ok {
			typeCaster, ok = dbValue.(TypeCaster)
		}
		if ok {
			cast = typeCaster.TypeCast()
			if !strings.Contains(cast, "%s") {
				cast = "%s::" + cast
			}
		}

		col := columnLayout{
			FieldName:    colMap.fieldName,
			ColumnName:   colMap.ColumnName,
			TypeDef:      stype,
			needsNotNull: notNullIgnored,
			isPK:         colMap.isPK,
			hasReference: colMap.References() != nil,
			typeCast:     cast,
			gotype:       colMap.gotype,
		}
		l = append(l, col)
	}
	return l
}

func (m *MigrationManager) addTable(t *TableMap) {
	l := m.layoutFor(t)
	for _, r := range m.newTables {
		if r.Schemaname == t.SchemaName && r.Tablename == t.TableName {
			r.Merge(l)
			return
		}
	}
	r := &tableRecord{
		Schemaname: t.SchemaName,
		Tablename:  t.TableName,
	}
	r.SetTableLayout(l)
	m.newTables = append(m.newTables, r)
}

func (m *MigrationManager) newTableRecords() []*tableRecord {
	if m.newTables == nil {
		m.newTables = make([]*tableRecord, 0, len(m.dbMap.tables))
		for _, tableMap := range m.dbMap.tables {
			m.addTable(tableMap)
		}
	}
	return m.newTables
}

func (m *MigrationManager) Migrate() (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch src := r.(type) {
			case error:
				err = src
			default:
				err = fmt.Errorf("Recovered from panic: %v", src)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(m.schemaname, m.tablename)
	_, err = m.dbMap.Select(&m.oldTables, "select * from "+quotedTable)
	if err != nil {
		return err
	}
	return m.run()
}

func (m *MigrationManager) run() error {
	for _, newTable := range m.newTableRecords() {
		found := false
		for _, oldTable := range m.oldTables {
			if oldTable.Schemaname == newTable.Schemaname && oldTable.Tablename == newTable.Tablename {
				found = true
				if err := m.migrateTable(oldTable, newTable); err != nil {
					return err
				}
			}
		}
		if !found {
			if err := m.dbMap.Insert(newTable); err != nil {
				return err
			}
		}
	}
	return nil
}

func (m *MigrationManager) changeType(quotedTable string, newCol columnLayout, tx *Transaction) error {
	if newCol.isPK {
		// Ehrm.  Backward compatibility issue, here.  Just ignore
		// PKeys for now.
		return nil
	}
	quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
	oldQuotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName + "_type_change_bak")
	sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + quotedColumn + " TO " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "UPDATE " + quotedTable + " SET " + quotedColumn + " = " + fmt.Sprintf(newCol.typeCast, oldQuotedColumn)
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	if newCol.needsNotNull {
		// The NOT NULL constraint needs to be set *after* updating data.
		sql = "ALTER TABLE " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
		if _, err := tx.Exec(sql); err != nil {
			return err
		}
	}
	sql = "ALTER TABLE " + quotedTable + " DROP COLUMN " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	return nil
}

func (m *MigrationManager) migrateTable(oldTable, newTable *tableRecord) (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	if oldTable.Schemaname != newTable.Schemaname || oldTable.Tablename != newTable.Tablename {
		return fmt.Errorf("Unsupported operation: table name change (%s.%s to %s.%s)",
			oldTable.Schemaname,
			oldTable.Tablename,
			newTable.Schemaname,
			newTable.Tablename,
		)
	}
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(newTable.Schemaname, newTable.Tablename)
	for _, newCol := range newTable.TableLayout() {
		quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
		found := false
		for _, oldCol := range oldTable.TableLayout() {
			found = strings.ToLower(oldCol.ColumnName) == strings.ToLower(newCol.ColumnName) ||
				oldCol.FieldName != "" && oldCol.FieldName == newCol.FieldName
			if found {
				// Take care of the case where the difference in
				// TypeDef is just due to gorp versions.
				if !strings.HasPrefix(oldCol.TypeDef, m.dbMap.Dialect.QuoteField(oldCol.ColumnName)) {
					oldCol.TypeDef = fmt.Sprintf("%s %s", m.dbMap.Dialect.QuoteField(oldCol.ColumnName), oldCol.TypeDef)
				}
				if oldCol.TypeDef != newCol.TypeDef {
					if err := m.changeType(quotedTable, newCol, tx); err != nil {
						return err
					}
				}
				break
			}
		}
		if !found {
			sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
			if _, err := tx.Exec(sql); err != nil {
				return err
			}
			// As of the time of this writing, we don't have a way
			// to know whether or not the foreign key column is
			// nillable.  If it was set to be NOT NULL, then the next
			// operation (adding the NOT NULL constraint) will fail.
			if !newCol.hasReference {
				defaultVal := reflect.New(newCol.gotype).Interface()
				sql = "UPDATE " + quotedTable + " SET " + quotedColumn + "=" + m.dbMap.Dialect.BindVar(0)
				if _, err := tx.Exec(sql, defaultVal); err != nil {
					return err
				}
			}
			if newCol.needsNotNull {
				// Most likely, the default value above was not null.
				//
				// TODO: support data binding, somehow.
				sql = "ALTER " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
				if _, err := tx.Exec(sql); err != nil {
					return err
				}
			}
		}
	}
	newTable.ID = oldTable.ID
	if count, err := tx.Update(newTable); err != nil || count != 1 {
		return err
	}
	return nil
}

// Migrater returns a MigrationManager using the given tablename as
// the migration table.
func (m *DbMap) Migrater(schemaname, tablename string) (*MigrationManager, error) {
	added := false
	for _, t := range m.tables {
		if t.SchemaName == schemaname && t.TableName == tablename && t.gotype == reflect.TypeOf(tableRecord{}) {
			added = true
			break
		}
	}
	if !added {
		// Just run the create table statement for the migration table,
		// using a temporary DbMap
		tmpM := &DbMap{
			Db:      m.Db,
			Dialect: m.Dialect,
		}
		tmpM.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
		if err := tmpM.CreateTablesIfNotExists(); err != nil {
			return nil, err
		}
		m.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	}
	return &MigrationManager{
		schemaname: schemaname,
		tablename:  tablename,
		dbMap:      m,
	}, nil
}
<commit_msg>Migration Backward Compatibility, NOT NULL updates<commit_after>package gorp

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// A TypeDefSwitcher is a TypeDeffer except that it won't have its type
// defined on table creation.  Instead, its type will be changed on
// the first migration run.  Useful for foreign keys when the target
// table may be created later than the column.
type TypeDefSwitcher interface {
	// TypeDefSwitch should return the same thing as
	// TypeDeffer.TypeDef.
	TypeDefSwitch() string
}

// A TypeCaster includes TypeDeffer logic but can also return the SQL
// to cast old types to its new type.
type TypeCaster interface {
	// TypeCast should return a string with zero or one '%s'
	// sequences.  If the string contains a '%s', it will be replaced
	// with the old value; otherwise, the return value will simply be
	// used as the type to cast to in the database.
	TypeCast() string
}

type pgJSON []byte

func (j pgJSON) Value() (driver.Value, error) {
	return []byte(j), nil
}

func (j *pgJSON) Scan(src interface{}) error {
	var source []byte
	switch src.(type) {
	case string:
		source = []byte(src.(string))
	case []byte:
		source = src.([]byte)
	default:
		return errors.New("Incompatible type for pgJSON")
	}
	*j = pgJSON(append((*j)[0:0], source...))
	return nil
}

func (j pgJSON) TypeDef() string {
	return "json"
}

type columnLayout struct {
	FieldName  string `json:"field_name"`
	ColumnName string `json:"column_name"`
	OldTypeDef string `json:"type_def"`
	TypeDef    string `json:"type_def_v2"`
	IsNotNull  bool   `json:"is_not_null"`

	// Values that are only used on the new layout, but are
	// unnecessary for old types.
	isPK         bool         `json:"-"`
	hasReference bool         `json:"-"`
	gotype       reflect.Type `json:"-"`
	typeCast     string       `json:"-"`
	oldFieldName string       `json:"-"`
}

type tableRecord struct {
	ID          int
	Schemaname  string
	Tablename   string
	Layout      pgJSON
	tableLayout []columnLayout `db:"-"`
}

func (t *tableRecord) TableLayout() []columnLayout {
	if t.tableLayout == nil {
		t.tableLayout = []columnLayout{}
		if err := json.Unmarshal([]byte(t.Layout), &t.tableLayout); err != nil {
			panic(err)
		}
	}
	return t.tableLayout
}

func (t *tableRecord) SetTableLayout(l []columnLayout) {
	t.tableLayout = l
	b, err := json.Marshal(l)
	if err != nil {
		panic(err)
	}
	t.Layout = pgJSON(b)
}

func (t *tableRecord) Merge(l []columnLayout) {
	if t.tableLayout == nil {
		t.SetTableLayout(l)
		return
	}
	for _, newCol := range l {
		shouldAppend := true
		for _, oldCol := range t.tableLayout {
			if newCol.ColumnName == oldCol.ColumnName {
				shouldAppend = false
			}
		}
		if shouldAppend {
			t.tableLayout = append(t.tableLayout, newCol)
		}
	}
}

func ptrToVal(t reflect.Type) reflect.Value {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return reflect.New(t)
}

type MigrationManager struct {
	schemaname string
	tablename  string
	dbMap      *DbMap
	oldTables  []*tableRecord
	newTables  []*tableRecord
}

func (m *MigrationManager) layoutFor(t *TableMap) []columnLayout {
	l := make([]columnLayout, 0, len(t.Columns))
	for _, colMap := range t.Columns {
		if colMap.Transient {
			continue
		}
		stype, notNullIgnored := colMap.TypeDefNoNotNull()
		orig := ptrToVal(colMap.origtype).Interface()
		dbValue := ptrToVal(colMap.gotype).Interface()
		cast := "%s"
		typeCaster, ok := orig.(TypeCaster)
		if !ok {
			typeCaster, ok = dbValue.(TypeCaster)
		}
		if ok {
			cast = typeCaster.TypeCast()
			if !strings.Contains(cast, "%s") {
				cast = "%s::" + cast
			}
		}

		col := columnLayout{
			FieldName:    colMap.fieldName,
			ColumnName:   colMap.ColumnName,
			TypeDef:      stype,
			IsNotNull:    notNullIgnored,
			isPK:         colMap.isPK,
			hasReference: colMap.References() != nil,
			typeCast:     cast,
			gotype:       colMap.gotype,
		}
		l = append(l, col)
	}
	return l
}

func (m *MigrationManager) addTable(t *TableMap) {
	l := m.layoutFor(t)
	for _, r := range m.newTables {
		if r.Schemaname == t.SchemaName && r.Tablename == t.TableName {
			r.Merge(l)
			return
		}
	}
	r := &tableRecord{
		Schemaname: t.SchemaName,
		Tablename:  t.TableName,
	}
	r.SetTableLayout(l)
	m.newTables = append(m.newTables, r)
}

func (m *MigrationManager) newTableRecords() []*tableRecord {
	if m.newTables == nil {
		m.newTables = make([]*tableRecord, 0, len(m.dbMap.tables))
		for _, tableMap := range m.dbMap.tables {
			m.addTable(tableMap)
		}
	}
	return m.newTables
}

func (m *MigrationManager) Migrate() (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if r := recover(); r != nil {
			switch src := r.(type) {
			case error:
				err = src
			default:
				err = fmt.Errorf("Recovered from panic: %v", src)
			}
		}
		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(m.schemaname, m.tablename)
	_, err = m.dbMap.Select(&m.oldTables, "select * from "+quotedTable)
	if err != nil {
		return err
	}
	// Check the old table layout for missing current versions
	for _, t := range m.oldTables {
		for _, oldCol := range t.tableLayout {
			if oldCol.TypeDef == "" {
				quotedColumn := m.dbMap.Dialect.QuoteField(oldCol.ColumnName)
				oldCol.TypeDef = quotedColumn + " " + oldCol.OldTypeDef
			}
		}
	}
	return m.run(tx)
}

func (m *MigrationManager) run(tx *Transaction) error {
	for _, newTable := range m.newTableRecords() {
		found := false
		for _, oldTable := range m.oldTables {
			if oldTable.Schemaname == newTable.Schemaname && oldTable.Tablename == newTable.Tablename {
				found = true
				if err := m.migrateTable(oldTable, newTable, tx); err != nil {
					return err
				}
			}
		}
		if !found {
			if err := tx.Insert(newTable); err != nil {
				return err
			}
		}
	}
	return nil
}

func (m *MigrationManager) changeType(quotedTable string, newCol, oldCol columnLayout, tx *Transaction) error {
	if newCol.isPK {
		// Ehrm.  Backward compatibility issue, here.  Just ignore
		// PKeys for now.
		return nil
	}
	quotedNewCol := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
	quotedOldCol := m.dbMap.Dialect.QuoteField(oldCol.ColumnName)
	quotedBakCol := quotedOldCol
	if quotedOldCol == quotedNewCol {
		quotedBakCol = m.dbMap.Dialect.QuoteField(oldCol.ColumnName + "_type_change_bak")
		sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + quotedOldCol + " TO " + quotedBakCol
		if _, err := tx.Exec(sql); err != nil {
			return err
		}
	}
	sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "UPDATE " + quotedTable + " SET " + quotedNewCol + " = " + fmt.Sprintf(newCol.typeCast, quotedBakCol)
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	if newCol.IsNotNull && oldCol.IsNotNull {
		// If the not null setting has changed, it will be handled
		// elsewhere; but if it is set to true on both the old column
		// and the new column, we have to handle it here.
		sql = "ALTER TABLE " + quotedTable + " ALTER COLUMN " + quotedNewCol + " SET NOT NULL"
		if _, err := tx.Exec(sql); err != nil {
			return err
		}
	}
	sql = "ALTER TABLE " + quotedTable + " DROP COLUMN " + quotedBakCol
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	return nil
}

func (m *MigrationManager) migrateTable(oldTable, newTable *tableRecord, tx *Transaction) (err error) {
	if oldTable.Schemaname != newTable.Schemaname || oldTable.Tablename != newTable.Tablename {
		return fmt.Errorf("Unsupported operation: table name change (%s.%s to %s.%s)",
			oldTable.Schemaname,
			oldTable.Tablename,
			newTable.Schemaname,
			newTable.Tablename,
		)
	}
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(newTable.Schemaname, newTable.Tablename)
	for _, newCol := range newTable.TableLayout() {
		quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
		found := false
		for _, oldCol := range oldTable.TableLayout() {
			found = strings.ToLower(oldCol.ColumnName) == strings.ToLower(newCol.ColumnName) ||
				oldCol.FieldName != "" && oldCol.FieldName == newCol.FieldName
			if !found {
				// This *only* handles conversion from the old foreign
				// key structure to the new one.
				found = strings.HasPrefix(newCol.FieldName, oldCol.FieldName+".")
			}
			if found {
				if oldCol.TypeDef != newCol.TypeDef {
					if err := m.changeType(quotedTable, newCol, oldCol, tx); err != nil {
						return err
					}
				}
				notNullOp := ""
				if !oldCol.IsNotNull && newCol.IsNotNull {
					notNullOp = "SET"
				} else if oldCol.IsNotNull && !newCol.IsNotNull {
					notNullOp = "DROP"
				}
				if notNullOp != "" {
					sql := "ALTER TABLE " + quotedTable + " ALTER COLUMN " + quotedColumn + " " + notNullOp + " NOT NULL"
					if _, err := tx.Exec(sql); err != nil {
						return err
					}
				}
				break
			}
		}
		if !found {
			sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
			if _, err := tx.Exec(sql); err != nil {
				return err
			}
			// As of the time of this writing, we don't have a way to
			// generate non-nil data for a foreign key column.  If it
			// was set to be NOT NULL, then the next operation (adding
			// the NOT NULL constraint) will fail.
			if !newCol.hasReference {
				defaultVal := reflect.New(newCol.gotype).Interface()
				sql = "UPDATE " + quotedTable + " SET " + quotedColumn + "=" + m.dbMap.Dialect.BindVar(0)
				if _, err := tx.Exec(sql, defaultVal); err != nil {
					return err
				}
			}
			if newCol.IsNotNull {
				// Most likely, the default value above was not null.
				//
				// TODO: support data binding, for when this new
				// column can be populated with data from existing
				// columns or some other method.
				sql = "ALTER " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
				if _, err := tx.Exec(sql); err != nil {
					return err
				}
			}
		}
	}
	newTable.ID = oldTable.ID
	if count, err := tx.Update(newTable); err != nil || count != 1 {
		return err
	}
	return nil
}

// Migrater returns a MigrationManager using the given tablename as
// the migration table.
func (m *DbMap) Migrater(schemaname, tablename string) (*MigrationManager, error) {
	added := false
	for _, t := range m.tables {
		if t.SchemaName == schemaname && t.TableName == tablename && t.gotype == reflect.TypeOf(tableRecord{}) {
			added = true
			break
		}
	}
	if !added {
		// Just run the create table statement for the migration table,
		// using a temporary DbMap
		tmpM := &DbMap{
			Db:      m.Db,
			Dialect: m.Dialect,
		}
		tmpM.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
		if err := tmpM.CreateTablesIfNotExists(); err != nil {
			return nil, err
		}
		m.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	}
	return &MigrationManager{
		schemaname: schemaname,
		tablename:  tablename,
		dbMap:      m,
	}, nil
}
<|endoftext|>


<commit_before>package gorp

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// A TypeDefSwitcher is a TypeDeffer except that it won't have its type
// defined on table creation.  Instead, its type will be changed on
// the first migration run.  Useful for foreign keys when the target
// table may be created later than the column.
type TypeDefSwitcher interface {
	// TypeDefSwitch should return the same thing as
	// TypeDeffer.TypeDef.
	TypeDefSwitch() string
}

// A TypeCaster includes TypeDeffer logic but can also return the SQL
// to cast old types to its new type.
type TypeCaster interface {
	// TypeCast should return a string with zero or one '%s'
	// sequences.  If the string contains a '%s', it will be replaced
	// with the old value; otherwise, the return value will simply be
	// used as the type to cast to in the database.
	TypeCast() string
}

type pgJSON []byte

func (j pgJSON) Value() (driver.Value, error) {
	return []byte(j), nil
}

func (j *pgJSON) Scan(src interface{}) error {
	var source []byte
	switch src.(type) {
	case string:
		source = []byte(src.(string))
	case []byte:
		source = src.([]byte)
	default:
		return errors.New("Incompatible type for pgJSON")
	}
	*j = pgJSON(append((*j)[0:0], source...))
	return nil
}

func (j pgJSON) TypeDef() string {
	return "json"
}

type columnLayout struct {
	FieldName  string `json:"field_name"`
	ColumnName string `json:"column_name"`
	TypeDef    string `json:"type_def"`

	// Values that are only used on the new layout, but are
	// unnecessary for old types.
	needsNotNull bool         `json:"-"`
	gotype       reflect.Type `json:"-"`
	typeCast     string       `json:"-"`
}

type tableRecord struct {
	ID          int
	Schemaname  string
	Tablename   string
	Layout      pgJSON
	tableLayout []columnLayout `db:"-"`
}

func (t *tableRecord) TableLayout() []columnLayout {
	if t.tableLayout == nil {
		t.tableLayout = []columnLayout{}
		if err := json.Unmarshal([]byte(t.Layout), &t.tableLayout); err != nil {
			panic(err)
		}
	}
	return t.tableLayout
}

func (t *tableRecord) SetTableLayout(l []columnLayout) {
	t.tableLayout = l
	b, err := json.Marshal(l)
	if err != nil {
		panic(err)
	}
	t.Layout = pgJSON(b)
}

func (t *tableRecord) Merge(l []columnLayout) {
	if t.tableLayout == nil {
		t.SetTableLayout(l)
		return
	}
	for _, newCol := range l {
		shouldAppend := true
		for _, oldCol := range t.tableLayout {
			if newCol.ColumnName == oldCol.ColumnName {
				shouldAppend = false
			}
		}
		if shouldAppend {
			t.tableLayout = append(t.tableLayout, newCol)
		}
	}
}

func ptrToVal(t reflect.Type) reflect.Value {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return reflect.New(t)
}

type MigrationManager struct {
	schemaname string
	tablename  string
	dbMap      *DbMap
	oldTables  []*tableRecord
	newTables  []*tableRecord
}

func (m *MigrationManager) layoutFor(t *TableMap) []columnLayout {
	l := make([]columnLayout, 0, len(t.Columns))
	for _, colMap := range t.Columns {
		if colMap.Transient {
			continue
		}
		stype, notNullIgnored := colMap.TypeDefNoNotNull()
		orig := ptrToVal(colMap.origtype).Interface()
		dbValue := ptrToVal(colMap.gotype).Interface()
		cast := "%s"
		typeCaster, ok := orig.(TypeCaster)
		if !ok {
			typeCaster, ok = dbValue.(TypeCaster)
		}
		if ok {
			cast = typeCaster.TypeCast()
			if !strings.Contains(cast, "%s") {
				cast = "%s::" + cast
			}
		}

		col := columnLayout{
			FieldName:    colMap.fieldName,
			ColumnName:   colMap.ColumnName,
			TypeDef:      stype,
			needsNotNull: notNullIgnored,
			typeCast:     cast,
			gotype:       colMap.gotype,
		}
		l = append(l, col)
	}
	return l
}

func (m *MigrationManager) addTable(t *TableMap) {
	l := m.layoutFor(t)
	for _, r := range m.newTables {
		if r.Schemaname == t.SchemaName && r.Tablename == t.TableName {
			r.Merge(l)
			return
		}
	}
	r := &tableRecord{
		Schemaname: t.SchemaName,
		Tablename:  t.TableName,
	}
	r.SetTableLayout(l)
	m.newTables = append(m.newTables, r)
}

func (m *MigrationManager) newTableRecords() []*tableRecord {
	if m.newTables == nil {
		m.newTables = make([]*tableRecord, 0, len(m.dbMap.tables))
		for _, tableMap := range m.dbMap.tables {
			m.addTable(tableMap)
		}
	}
	return m.newTables
}

func (m *MigrationManager) Migrate() (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch src := r.(type) {
			case error:
				err = src
			default:
				err = fmt.Errorf("Recovered from panic: %v", src)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(m.schemaname, m.tablename)
	_, err = m.dbMap.Select(&m.oldTables, "select * from "+quotedTable)
	if err != nil {
		return err
	}
	return m.run()
}

func (m *MigrationManager) run() error {
	for _, newTable := range m.newTableRecords() {
		found := false
		for _, oldTable := range m.oldTables {
			if oldTable.Schemaname == newTable.Schemaname && oldTable.Tablename == newTable.Tablename {
				found = true
				if err := m.migrateTable(oldTable, newTable); err != nil {
					return err
				}
			}
		}
		if !found {
			if err := m.dbMap.Insert(newTable); err != nil {
				return err
			}
		}
	}
	return nil
}

func (m *MigrationManager) changeType(quotedTable string, newCol columnLayout, tx *Transaction) error {
	quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
	oldQuotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName + "_type_change_bak")
	sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + quotedColumn + " TO " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "UPDATE " + quotedTable + " SET " + quotedColumn + " = " + fmt.Sprintf(newCol.typeCast, oldQuotedColumn)
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	if newCol.needsNotNull {
		// The NOT NULL constraint needs to be set *after* updating data.
		sql = "ALTER TABLE " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
		if _, err := tx.Exec(sql); err != nil {
			return err
		}
	}
	sql = "ALTER TABLE " + quotedTable + " DROP COLUMN " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	return nil
}

func (m *MigrationManager) migrateTable(oldTable, newTable *tableRecord) (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	if oldTable.Schemaname != newTable.Schemaname || oldTable.Tablename != newTable.Tablename {
		return fmt.Errorf("Unsupported operation: table name change (%s.%s to %s.%s)",
			oldTable.Schemaname,
			oldTable.Tablename,
			newTable.Schemaname,
			newTable.Tablename,
		)
	}
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(newTable.Schemaname, newTable.Tablename)
	for _, newCol := range newTable.TableLayout() {
		quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
		found := false
		for _, oldCol := range oldTable.TableLayout() {
			found = strings.ToLower(oldCol.ColumnName) == strings.ToLower(newCol.ColumnName) ||
				oldCol.FieldName != "" && oldCol.FieldName == newCol.FieldName
			if found {
				if oldCol.TypeDef != newCol.TypeDef {
					if err := m.changeType(quotedTable, newCol, tx); err != nil {
						return err
					}
				}
				break
			}
		}
		if !found {
			sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
			if _, err := tx.Exec(sql); err != nil {
				return err
			}
			defaultVal := reflect.New(newCol.gotype).Interface()
			sql = "UPDATE " + quotedTable + " SET " + quotedColumn + "=" + m.dbMap.Dialect.BindVar(0)
			if _, err := tx.Exec(sql, defaultVal); err != nil {
				return err
			}
			if newCol.needsNotNull {
				// Most likely, the default value above was not null.
				//
				// TODO: support data binding, somehow.
				sql = "ALTER " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
				if _, err := tx.Exec(sql); err != nil {
					return err
				}
			}
		}
	}
	newTable.ID = oldTable.ID
	if count, err := tx.Update(newTable); err != nil || count != 1 {
		return err
	}
	return nil
}

// Migrater returns a MigrationManager using the given tablename as
// the migration table.
func (m *DbMap) Migrater(schemaname, tablename string) (*MigrationManager, error) {
	added := false
	for _, t := range m.tables {
		if t.SchemaName == schemaname && t.TableName == tablename && t.gotype == reflect.TypeOf(tableRecord{}) {
			added = true
			break
		}
	}
	if !added {
		// Just run the create table statement for the migration table,
		// using a temporary DbMap
		tmpM := &DbMap{
			Db:      m.Db,
			Dialect: m.Dialect,
		}
		tmpM.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
		if err := tmpM.CreateTablesIfNotExists(); err != nil {
			return nil, err
		}
		m.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	}
	return &MigrationManager{
		schemaname: schemaname,
		tablename:  tablename,
		dbMap:      m,
	}, nil
}
<commit_msg>Handle a backward compatibility issue with PKeys in the migration manager<commit_after>package gorp

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

// A TypeDefSwitcher is a TypeDeffer except that it won't have its type
// defined on table creation.  Instead, its type will be changed on
// the first migration run.  Useful for foreign keys when the target
// table may be created later than the column.
type TypeDefSwitcher interface {
	// TypeDefSwitch should return the same thing as
	// TypeDeffer.TypeDef.
	TypeDefSwitch() string
}

// A TypeCaster includes TypeDeffer logic but can also return the SQL
// to cast old types to its new type.
type TypeCaster interface {
	// TypeCast should return a string with zero or one '%s'
	// sequences.  If the string contains a '%s', it will be replaced
	// with the old value; otherwise, the return value will simply be
	// used as the type to cast to in the database.
	TypeCast() string
}

type pgJSON []byte

func (j pgJSON) Value() (driver.Value, error) {
	return []byte(j), nil
}

func (j *pgJSON) Scan(src interface{}) error {
	var source []byte
	switch src.(type) {
	case string:
		source = []byte(src.(string))
	case []byte:
		source = src.([]byte)
	default:
		return errors.New("Incompatible type for pgJSON")
	}
	*j = pgJSON(append((*j)[0:0], source...))
	return nil
}

func (j pgJSON) TypeDef() string {
	return "json"
}

type columnLayout struct {
	FieldName  string `json:"field_name"`
	ColumnName string `json:"column_name"`
	TypeDef    string `json:"type_def"`

	// Values that are only used on the new layout, but are
	// unnecessary for old types.
	needsNotNull bool         `json:"-"`
	isPK         bool         `json:"-"`
	gotype       reflect.Type `json:"-"`
	typeCast     string       `json:"-"`
}

type tableRecord struct {
	ID          int
	Schemaname  string
	Tablename   string
	Layout      pgJSON
	tableLayout []columnLayout `db:"-"`
}

func (t *tableRecord) TableLayout() []columnLayout {
	if t.tableLayout == nil {
		t.tableLayout = []columnLayout{}
		if err := json.Unmarshal([]byte(t.Layout), &t.tableLayout); err != nil {
			panic(err)
		}
	}
	return t.tableLayout
}

func (t *tableRecord) SetTableLayout(l []columnLayout) {
	t.tableLayout = l
	b, err := json.Marshal(l)
	if err != nil {
		panic(err)
	}
	t.Layout = pgJSON(b)
}

func (t *tableRecord) Merge(l []columnLayout) {
	if t.tableLayout == nil {
		t.SetTableLayout(l)
		return
	}
	for _, newCol := range l {
		shouldAppend := true
		for _, oldCol := range t.tableLayout {
			if newCol.ColumnName == oldCol.ColumnName {
				shouldAppend = false
			}
		}
		if shouldAppend {
			t.tableLayout = append(t.tableLayout, newCol)
		}
	}
}

func ptrToVal(t reflect.Type) reflect.Value {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return reflect.New(t)
}

type MigrationManager struct {
	schemaname string
	tablename  string
	dbMap      *DbMap
	oldTables  []*tableRecord
	newTables  []*tableRecord
}

func (m *MigrationManager) layoutFor(t *TableMap) []columnLayout {
	l := make([]columnLayout, 0, len(t.Columns))
	for _, colMap := range t.Columns {
		if colMap.Transient {
			continue
		}
		stype, notNullIgnored := colMap.TypeDefNoNotNull()
		orig := ptrToVal(colMap.origtype).Interface()
		dbValue := ptrToVal(colMap.gotype).Interface()
		cast := "%s"
		typeCaster, ok := orig.(TypeCaster)
		if !ok {
			typeCaster, ok = dbValue.(TypeCaster)
		}
		if ok {
			cast = typeCaster.TypeCast()
			if !strings.Contains(cast, "%s") {
				cast = "%s::" + cast
			}
		}

		col := columnLayout{
			FieldName:    colMap.fieldName,
			ColumnName:   colMap.ColumnName,
			TypeDef:      stype,
			needsNotNull: notNullIgnored,
			isPK:         colMap.isPK,
			typeCast:     cast,
			gotype:       colMap.gotype,
		}
		l = append(l, col)
	}
	return l
}

func (m *MigrationManager) addTable(t *TableMap) {
	l := m.layoutFor(t)
	for _, r := range m.newTables {
		if r.Schemaname == t.SchemaName && r.Tablename == t.TableName {
			r.Merge(l)
			return
		}
	}
	r := &tableRecord{
		Schemaname: t.SchemaName,
		Tablename:  t.TableName,
	}
	r.SetTableLayout(l)
	m.newTables = append(m.newTables, r)
}

func (m *MigrationManager) newTableRecords() []*tableRecord {
	if m.newTables == nil {
		m.newTables = make([]*tableRecord, 0, len(m.dbMap.tables))
		for _, tableMap := range m.dbMap.tables {
			m.addTable(tableMap)
		}
	}
	return m.newTables
}

func (m *MigrationManager) Migrate() (err error) {
	defer func() {
		if r := recover(); r != nil {
			switch src := r.(type) {
			case error:
				err = src
			default:
				err = fmt.Errorf("Recovered from panic: %v", src)
			}
		}
	}()
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(m.schemaname, m.tablename)
	_, err = m.dbMap.Select(&m.oldTables, "select * from "+quotedTable)
	if err != nil {
		return err
	}
	return m.run()
}

func (m *MigrationManager) run() error {
	for _, newTable := range m.newTableRecords() {
		found := false
		for _, oldTable := range m.oldTables {
			if oldTable.Schemaname == newTable.Schemaname && oldTable.Tablename == newTable.Tablename {
				found = true
				if err := m.migrateTable(oldTable, newTable); err != nil {
					return err
				}
			}
		}
		if !found {
			if err := m.dbMap.Insert(newTable); err != nil {
				return err
			}
		}
	}
	return nil
}

func (m *MigrationManager) changeType(quotedTable string, newCol columnLayout, tx *Transaction) error {
	if newCol.isPK {
		// Ehrm.  Backward compatibility issue, here.  Just ignore
		// PKeys for now.
		return nil
	}
	quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
	oldQuotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName + "_type_change_bak")
	sql := "ALTER TABLE " + quotedTable + " RENAME COLUMN " + quotedColumn + " TO " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	sql = "UPDATE " + quotedTable + " SET " + quotedColumn + " = " + fmt.Sprintf(newCol.typeCast, oldQuotedColumn)
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	if newCol.needsNotNull {
		// The NOT NULL constraint needs to be set *after* updating data.
		sql = "ALTER TABLE " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
		if _, err := tx.Exec(sql); err != nil {
			return err
		}
	}
	sql = "ALTER TABLE " + quotedTable + " DROP COLUMN " + oldQuotedColumn
	if _, err := tx.Exec(sql); err != nil {
		return err
	}
	return nil
}

func (m *MigrationManager) migrateTable(oldTable, newTable *tableRecord) (err error) {
	tx, err := m.dbMap.Begin()
	if err != nil {
		panic(err)
	}
	defer func() {
		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				panic(rollbackErr)
			}
		}
	}()
	if oldTable.Schemaname != newTable.Schemaname || oldTable.Tablename != newTable.Tablename {
		return fmt.Errorf("Unsupported operation: table name change (%s.%s to %s.%s)",
			oldTable.Schemaname,
			oldTable.Tablename,
			newTable.Schemaname,
			newTable.Tablename,
		)
	}
	quotedTable := m.dbMap.Dialect.QuotedTableForQuery(newTable.Schemaname, newTable.Tablename)
	for _, newCol := range newTable.TableLayout() {
		quotedColumn := m.dbMap.Dialect.QuoteField(newCol.ColumnName)
		found := false
		for _, oldCol := range oldTable.TableLayout() {
			found = strings.ToLower(oldCol.ColumnName) == strings.ToLower(newCol.ColumnName) ||
				oldCol.FieldName != "" && oldCol.FieldName == newCol.FieldName
			if found {
				if oldCol.TypeDef != newCol.TypeDef {
					if err := m.changeType(quotedTable, newCol, tx); err != nil {
						return err
					}
				}
				break
			}
		}
		if !found {
			sql := "ALTER TABLE " + quotedTable + " ADD COLUMN " + newCol.TypeDef
			if _, err := tx.Exec(sql); err != nil {
				return err
			}
			defaultVal := reflect.New(newCol.gotype).Interface()
			sql = "UPDATE " + quotedTable + " SET " + quotedColumn + "=" + m.dbMap.Dialect.BindVar(0)
			if _, err := tx.Exec(sql, defaultVal); err != nil {
				return err
			}
			if newCol.needsNotNull {
				// Most likely, the default value above was not null.
				//
				// TODO: support data binding, somehow.
				sql = "ALTER " + quotedTable + " ALTER COLUMN " + quotedColumn + " SET NOT NULL"
				if _, err := tx.Exec(sql); err != nil {
					return err
				}
			}
		}
	}
	newTable.ID = oldTable.ID
	if count, err := tx.Update(newTable); err != nil || count != 1 {
		return err
	}
	return nil
}

// Migrater returns a MigrationManager using the given tablename as
// the migration table.
func (m *DbMap) Migrater(schemaname, tablename string) (*MigrationManager, error) {
	added := false
	for _, t := range m.tables {
		if t.SchemaName == schemaname && t.TableName == tablename && t.gotype == reflect.TypeOf(tableRecord{}) {
			added = true
			break
		}
	}
	if !added {
		// Just run the create table statement for the migration table,
		// using a temporary DbMap
		tmpM := &DbMap{
			Db:      m.Db,
			Dialect: m.Dialect,
		}
		tmpM.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
		if err := tmpM.CreateTablesIfNotExists(); err != nil {
			return nil, err
		}
		m.AddTableWithNameAndSchema(tableRecord{}, schemaname, tablename).SetKeys(true, "ID")
	}
	return &MigrationManager{
		schemaname: schemaname,
		tablename:  tablename,
		dbMap:      m,
	}, nil
}
<|endoftext|>


<commit_before>package migrator

import (
	"github.com/mattes/migrate/migrate"
	"github.com/resourced/resourced-master/config"
)

func New(generalConfig config.GeneralConfig) *Migrator {
	m := &Migrator{}
	m.AppGeneralConfig = generalConfig

	return m
}

type Migrator struct {
	AppGeneralConfig config.GeneralConfig
}

func (m *Migrator) CoreMigrateUp() ([]error, bool) {
	return migrate.UpSync(m.AppGeneralConfig.DSN, "./migrations/core")
}

func (m *Migrator) TSWatchersMigrateUp() ([]error, bool) {
	for _, dsn := range m.AppGeneralConfig.Watchers.DSNs {
		errs, ok := migrate.UpSync(dsn, "./migrations/ts-watchers")
		if errs != nil && len(errs) > 0 {
			return errs, ok
		}
	}

	return nil, true
}

func (m *Migrator) TSMetricsMigrateUp() ([]error, bool) {
	for _, dsn := range m.AppGeneralConfig.Metrics.DSNs {
		errs, ok := migrate.UpSync(dsn, "./migrations/ts-metrics")
		if errs != nil && len(errs) > 0 {
			return errs, ok
		}
	}

	return nil, true
}
<commit_msg>Upgrade migrate library to avoid using sqlite lib.<commit_after>package migrator

import (
	_ "github.com/mattes/migrate/driver/postgres"
	"github.com/mattes/migrate/migrate"
	"github.com/resourced/resourced-master/config"
)

func New(generalConfig config.GeneralConfig) *Migrator {
	m := &Migrator{}
	m.AppGeneralConfig = generalConfig

	return m
}

type Migrator struct {
	AppGeneralConfig config.GeneralConfig
}

func (m *Migrator) CoreMigrateUp() ([]error, bool) {
	return migrate.UpSync(m.AppGeneralConfig.DSN, "./migrations/core")
}

func (m *Migrator) TSWatchersMigrateUp() ([]error, bool) {
	for _, dsn := range m.AppGeneralConfig.Watchers.DSNs {
		errs, ok := migrate.UpSync(dsn, "./migrations/ts-watchers")
		if errs != nil && len(errs) > 0 {
			return errs, ok
		}
	}

	return nil, true
}

func (m *Migrator) TSMetricsMigrateUp() ([]error, bool) {
	for _, dsn := range m.AppGeneralConfig.Metrics.DSNs {
		errs, ok := migrate.UpSync(dsn, "./migrations/ts-metrics")
		if errs != nil && len(errs) > 0 {
			return errs, ok
		}
	}

	return nil, true
}
<|endoftext|>


<commit_before>package main_test

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/concourse/time-resource/models"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("Out", func() {
	var tmpdir string
	var source string

	var outCmd *exec.Cmd

	BeforeEach(func() {
		var err error

		tmpdir, err = ioutil.TempDir("", "out-source")
		Expect(err).NotTo(HaveOccurred())

		source = path.Join(tmpdir, "out-dir")

		outCmd = exec.Command(outPath, source)
	})

	AfterEach(func() {
		os.RemoveAll(tmpdir)
	})

	Context("when executed", func() {
		var request models.OutRequest
		var response models.OutResponse

		BeforeEach(func() {
			request = models.OutRequest{
				Source: models.Source{Interval: "1s"},
			}

			response = models.OutResponse{}
		})

		JustBeforeEach(func() {
			stdin, err := outCmd.StdinPipe()
			Expect(err).NotTo(HaveOccurred())

			session, err := gexec.Start(outCmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			err = json.NewEncoder(stdin).Encode(request)
			Expect(err).NotTo(HaveOccurred())

			Eventually(session).Should(gexec.Exit(0))

			err = json.Unmarshal(session.Out.Contents(), &response)
			Expect(err).NotTo(HaveOccurred())
		})

		It("reports the current time as the version", func() {
			Expect(response.Version.Time).To(BeTemporally("~", time.Now(), time.Second))
		})
	})
})
<commit_msg>fix race when command quickly exits<commit_after>package main_test

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/concourse/time-resource/models"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("Out", func() {
	var tmpdir string
	var source string

	var outCmd *exec.Cmd

	BeforeEach(func() {
		var err error

		tmpdir, err = ioutil.TempDir("", "out-source")
		Expect(err).NotTo(HaveOccurred())

		source = path.Join(tmpdir, "out-dir")

		outCmd = exec.Command(outPath, source)
	})

	AfterEach(func() {
		os.RemoveAll(tmpdir)
	})

	Context("when executed", func() {
		var request models.OutRequest
		var response models.OutResponse

		BeforeEach(func() {
			request = models.OutRequest{
				Source: models.Source{Interval: "1s"},
			}

			response = models.OutResponse{}
		})

		JustBeforeEach(func() {
			stdin := new(bytes.Buffer)

			err := json.NewEncoder(stdin).Encode(request)
			Expect(err).NotTo(HaveOccurred())

			outCmd.Stdin = stdin

			session, err := gexec.Start(outCmd, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			<-session.Exited
			Expect(session.ExitCode()).To(Equal(0))

			err = json.Unmarshal(session.Out.Contents(), &response)
			Expect(err).NotTo(HaveOccurred())
		})

		It("reports the current time as the version", func() {
			Expect(response.Version.Time).To(BeTemporally("~", time.Now(), time.Second))
		})
	})
})
<|endoftext|>


<commit_before>/*
 Copyright 2020 The GoPlus Authors (goplus.org)

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

package parser

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/qiniu/goplus/ast"
	"github.com/qiniu/goplus/scanner"
	"github.com/qiniu/goplus/token"
)

// -----------------------------------------------------------------------------

// FileSystem represents a file system.
type FileSystem interface {
	ReadDir(dirname string) ([]os.FileInfo, error)
	ReadFile(filename string) ([]byte, error)
	Join(elem ...string) string
}

type localFS struct{}

func (p localFS) ReadDir(dirname string) ([]os.FileInfo, error) {
	return ioutil.ReadDir(dirname)
}

func (p localFS) ReadFile(filename string) ([]byte, error) {
	return ioutil.ReadFile(filename)
}

func (p localFS) Join(elem ...string) string {
	return filepath.Join(elem...)
}

var local FileSystem = localFS{}

// ParseGOPFiles parses the Go+ source files under directory or single Go+ source file.
// The target specifies the directory or single Go+ source file.
//
// The ParseGOPFiles should return the map of packages to run Go+ script, even the target is single file.
//
// If the file or directory couldn't be read, a nil map and the respective error are
// returned.
// If the target is directory and a parse error occurred, a non-nil but incomplete map and the first error encountered are returned.
func ParseGopFiles(fset *token.FileSet, target string, isDir bool, mode Mode) (pkgs map[string]*ast.Package, err error) {
	if !isDir {
		file, err := ParseFile(fset, target, nil, mode)
		if err != nil {
			return pkgs, err
		}
		pkgs = make(map[string]*ast.Package)
		pkg := astFileToPkg(file, target)
		pkgs[file.Name.Name] = pkg
		return pkgs, nil
	}
	return ParseDir(fset, target, nil, mode)
}

// -----------------------------------------------------------------------------
// ParseDir calls ParseFSDir by passing a local filesystem.
//
func ParseDir(fset *token.FileSet, path string, filter func(os.FileInfo) bool, mode Mode) (pkgs map[string]*ast.Package, first error) {
	return ParseFSDir(fset, local, path, filter, mode)
}

// ParseFSDir calls ParseFile for all files with names ending in ".gop" in the
// directory specified by path and returns a map of package name -> package
// AST with all the packages found.
//
// If filter != nil, only the files with os.FileInfo entries passing through
// the filter (and ending in ".gop") are considered. The mode bits are passed
// to ParseFile unchanged. Position information is recorded in fset, which
// must not be nil.
//
// If the directory couldn't be read, a nil map and the respective error are
// returned. If a parse error occurred, a non-nil but incomplete map and the
// first error encountered are returned.
//
func ParseFSDir(fset *token.FileSet, fs FileSystem, path string, filter func(os.FileInfo) bool, mode Mode) (pkgs map[string]*ast.Package, first error) {
	list, err := fs.ReadDir(path)
	if err != nil {
		return nil, err
	}
	pkgs = make(map[string]*ast.Package)
	for _, d := range list {
		if filter == nil || filter(d) {
			filename := fs.Join(path, d.Name())
			if filedata, err := fs.ReadFile(filename); err == nil {
				if src, err := ParseFSFile(fset, fs, filename, filedata, mode); err == nil {
					name := src.Name.Name
					pkg, found := pkgs[name]
					if !found {
						pkg = &ast.Package{
							Name:  name,
							Files: make(map[string]*ast.File),
						}
						pkgs[name] = pkg
					}
					pkg.Files[filename] = src
				} else if first == nil {
					first = err
				}
			} else if first == nil {
				first = err
			}
		}
	}
	return
}

// -----------------------------------------------------------------------------

// ParseFile parses the source code of a single Go+ source file and returns the corresponding ast.File node.
func ParseFile(fset *token.FileSet, filename string, src interface{}, mode Mode) (f *ast.File, err error) {
	return ParseFSFile(fset, local, filename, src, mode)
}

// ParseFSFile parses the source code of a single Go+ source file and returns the corresponding ast.File node.
func ParseFSFile(fset *token.FileSet, fs FileSystem, filename string, src interface{}, mode Mode) (f *ast.File, err error) {
	var code []byte
	if src == nil {
		code, err = fs.ReadFile(filename)
	} else {
		code, err = readSource(src)
	}
	if err != nil {
		return
	}
	return parseFileEx(fset, filename, code, mode)
}

// TODO: should not add package info and init|main function.
// If do this, parsing will display error line number when error occur
func parseFileEx(fset *token.FileSet, filename string, code []byte, mode Mode) (f *ast.File, err error) {
	var b []byte
	var isMod, hasUnnamed bool
	var fsetTmp = token.NewFileSet()
	f, err = parseFile(fsetTmp, filename, code, PackageClauseOnly)
	if err != nil {
		n := len(code)
		b = make([]byte, n+28)
		copy(b, "package main;")
		copy(b[13:], code)
		code = b[:n+13]
	} else {
		isMod = f.Name.Name != "main"
	}
	_, err = parseFile(fsetTmp, filename, code, mode)
	if err != nil {
		if errlist, ok := err.(scanner.ErrorList); ok {
			if e := errlist[0]; strings.HasPrefix(e.Msg, "expected declaration") {
				n := len(code)
				idx := e.Pos.Offset
				if b == nil {
					b = make([]byte, n+14)
					copy(b, code[:idx])
				}
				copy(b[idx+12:], code[idx:n])
				if isMod {
					copy(b[idx:], "func init(){")
				} else {
					copy(b[idx:], "func main(){")
				}
				b[n+12] = '\n'
				b[n+13] = '}'
				code = b[:n+14]
				hasUnnamed = true
				err = nil
			}
		}
	}
	if err == nil {
		f, err = parseFile(fset, filename, code, mode)
		if err == nil {
			f.HasUnnamed = hasUnnamed
		}
	}
	return
}

var (
	errInvalidSource = errors.New("invalid source")
)

func readSource(src interface{}) ([]byte, error) {
	switch s := src.(type) {
	case string:
		return []byte(s), nil
	case []byte:
		return s, nil
	case *bytes.Buffer:
		// is io.Reader, but src is already available in []byte form
		if s != nil {
			return s.Bytes(), nil
		}
	case io.Reader:
		return ioutil.ReadAll(s)
	}
	return nil, errInvalidSource
}

// -----------------------------------------------------------------------------

func (p *parser) parseSliceLit(lbrack token.Pos, len ast.Expr) ast.Expr {
	elts := make([]ast.Expr, 1, 8)
	elts[0] = len
	for p.tok == token.COMMA {
		p.next()
		elt := p.parseRHS()
		elts = append(elts, elt)
	}
	rbrack := p.expect(token.RBRACK)
	return &ast.SliceLit{Lbrack: lbrack, Elts: elts, Rbrack: rbrack}
}

func newSliceLit(lbrack, rbrack token.Pos, len ast.Expr) ast.Expr {
	var elts []ast.Expr
	if len != nil {
		elts = []ast.Expr{len}
	}
	return &ast.SliceLit{Lbrack: lbrack, Elts: elts, Rbrack: rbrack}
}

// astFileToPkg translate ast.File to ast.Package
func astFileToPkg(file *ast.File, fileName string) (pkg *ast.Package) {
	pkg = &ast.Package{
		Name:  file.Name.Name,
		Files: make(map[string]*ast.File),
	}
	pkg.Files[fileName] = file
	return
}

// -----------------------------------------------------------------------------
<commit_msg>remove gop suffix for singnle file<commit_after>/*
 Copyright 2020 The GoPlus Authors (goplus.org)

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

package parser

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/qiniu/goplus/ast"
	"github.com/qiniu/goplus/scanner"
	"github.com/qiniu/goplus/token"
)

// -----------------------------------------------------------------------------

// FileSystem represents a file system.
type FileSystem interface {
	ReadDir(dirname string) ([]os.FileInfo, error)
	ReadFile(filename string) ([]byte, error)
	Join(elem ...string) string
}

type localFS struct{}

func (p localFS) ReadDir(dirname string) ([]os.FileInfo, error) {
	return ioutil.ReadDir(dirname)
}

func (p localFS) ReadFile(filename string) ([]byte, error) {
	return ioutil.ReadFile(filename)
}

func (p localFS) Join(elem ...string) string {
	return filepath.Join(elem...)
}

var local FileSystem = localFS{}

// ParseGopFiles parses the Go+ source files under directory or single Go+ source file.
// The target specifies the directory or single Go+ source file.
//
// The ParseGopFiles should return the map of packages to run Go+ script, even the target is single file.
//
// If the file or directory couldn't be read, a nil map and the respective error are
// returned.
// If the target is directory and a parse error occurred, a non-nil but incomplete map and the first error encountered are returned.
func ParseGopFiles(fset *token.FileSet, target string, isDir bool, mode Mode) (pkgs map[string]*ast.Package, err error) {
	if !isDir {
		file, err := ParseFile(fset, target, nil, mode)
		if err != nil {
			return pkgs, err
		}
		pkgs = make(map[string]*ast.Package)
		pkg := astFileToPkg(file, target)
		pkgs[file.Name.Name] = pkg
		return pkgs, nil
	}
	return ParseDir(fset, target, nil, mode)
}

// -----------------------------------------------------------------------------
// ParseDir calls ParseFSDir by passing a local filesystem.
//
func ParseDir(fset *token.FileSet, path string, filter func(os.FileInfo) bool, mode Mode) (pkgs map[string]*ast.Package, first error) {
	return ParseFSDir(fset, local, path, filter, mode)
}

// ParseFSDir calls ParseFile for all files with names ending in ".gop" in the
// directory specified by path and returns a map of package name -> package
// AST with all the packages found.
//
// If filter != nil, only the files with os.FileInfo entries passing through
// the filter (and ending in ".gop") are considered. The mode bits are passed
// to ParseFile unchanged. Position information is recorded in fset, which
// must not be nil.
//
// If the directory couldn't be read, a nil map and the respective error are
// returned. If a parse error occurred, a non-nil but incomplete map and the
// first error encountered are returned.
//
func ParseFSDir(fset *token.FileSet, fs FileSystem, path string, filter func(os.FileInfo) bool, mode Mode) (pkgs map[string]*ast.Package, first error) {
	list, err := fs.ReadDir(path)
	if err != nil {
		return nil, err
	}
	pkgs = make(map[string]*ast.Package)
	for _, d := range list {
		if strings.HasSuffix(d.Name(), ".gop") && (filter == nil || filter(d)) {
			filename := fs.Join(path, d.Name())
			if filedata, err := fs.ReadFile(filename); err == nil {
				if src, err := ParseFSFile(fset, fs, filename, filedata, mode); err == nil {
					name := src.Name.Name
					pkg, found := pkgs[name]
					if !found {
						pkg = &ast.Package{
							Name:  name,
							Files: make(map[string]*ast.File),
						}
						pkgs[name] = pkg
					}
					pkg.Files[filename] = src
				} else if first == nil {
					first = err
				}
			} else if first == nil {
				first = err
			}
		}
	}
	return
}

// -----------------------------------------------------------------------------

// ParseFile parses the source code of a single Go+ source file and returns the corresponding ast.File node.
func ParseFile(fset *token.FileSet, filename string, src interface{}, mode Mode) (f *ast.File, err error) {
	return ParseFSFile(fset, local, filename, src, mode)
}

// ParseFSFile parses the source code of a single Go+ source file and returns the corresponding ast.File node.
func ParseFSFile(fset *token.FileSet, fs FileSystem, filename string, src interface{}, mode Mode) (f *ast.File, err error) {
	var code []byte
	if src == nil {
		code, err = fs.ReadFile(filename)
	} else {
		code, err = readSource(src)
	}
	if err != nil {
		return
	}
	return parseFileEx(fset, filename, code, mode)
}

// TODO: should not add package info and init|main function.
// If do this, parsing will display error line number when error occur
func parseFileEx(fset *token.FileSet, filename string, code []byte, mode Mode) (f *ast.File, err error) {
	var b []byte
	var isMod, hasUnnamed bool
	var fsetTmp = token.NewFileSet()
	f, err = parseFile(fsetTmp, filename, code, PackageClauseOnly)
	if err != nil {
		n := len(code)
		b = make([]byte, n+28)
		copy(b, "package main;")
		copy(b[13:], code)
		code = b[:n+13]
	} else {
		isMod = f.Name.Name != "main"
	}
	_, err = parseFile(fsetTmp, filename, code, mode)
	if err != nil {
		if errlist, ok := err.(scanner.ErrorList); ok {
			if e := errlist[0]; strings.HasPrefix(e.Msg, "expected declaration") {
				n := len(code)
				idx := e.Pos.Offset
				if b == nil {
					b = make([]byte, n+14)
					copy(b, code[:idx])
				}
				copy(b[idx+12:], code[idx:n])
				if isMod {
					copy(b[idx:], "func init(){")
				} else {
					copy(b[idx:], "func main(){")
				}
				b[n+12] = '\n'
				b[n+13] = '}'
				code = b[:n+14]
				hasUnnamed = true
				err = nil
			}
		}
	}
	if err == nil {
		f, err = parseFile(fset, filename, code, mode)
		if err == nil {
			f.HasUnnamed = hasUnnamed
		}
	}
	return
}

var (
	errInvalidSource = errors.New("invalid source")
)

func readSource(src interface{}) ([]byte, error) {
	switch s := src.(type) {
	case string:
		return []byte(s), nil
	case []byte:
		return s, nil
	case *bytes.Buffer:
		// is io.Reader, but src is already available in []byte form
		if s != nil {
			return s.Bytes(), nil
		}
	case io.Reader:
		return ioutil.ReadAll(s)
	}
	return nil, errInvalidSource
}

// -----------------------------------------------------------------------------

func (p *parser) parseSliceLit(lbrack token.Pos, len ast.Expr) ast.Expr {
	elts := make([]ast.Expr, 1, 8)
	elts[0] = len
	for p.tok == token.COMMA {
		p.next()
		elt := p.parseRHS()
		elts = append(elts, elt)
	}
	rbrack := p.expect(token.RBRACK)
	return &ast.SliceLit{Lbrack: lbrack, Elts: elts, Rbrack: rbrack}
}

func newSliceLit(lbrack, rbrack token.Pos, len ast.Expr) ast.Expr {
	var elts []ast.Expr
	if len != nil {
		elts = []ast.Expr{len}
	}
	return &ast.SliceLit{Lbrack: lbrack, Elts: elts, Rbrack: rbrack}
}

// astFileToPkg translate ast.File to ast.Package
func astFileToPkg(file *ast.File, fileName string) (pkg *ast.Package) {
	pkg = &ast.Package{
		Name:  file.Name.Name,
		Files: make(map[string]*ast.File),
	}
	pkg.Files[fileName] = file
	return
}

// -----------------------------------------------------------------------------
<|endoftext|>