<commit_before>// Copyright 2020 The Ebiten Authors
//
// 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 shaderir

import (
	"fmt"
	"strings"
)

func isValidSwizzling(s string) bool {
	if len(s) < 1 || 4 < len(s) {
		return false
	}

	const (
		xyzw = "xyzw"
		rgba = "rgba"
		strq = "strq"
	)

	switch {
	case strings.IndexByte(xyzw, s[0]) >= 0:
		for _, c := range s {
			if strings.IndexRune(xyzw, c) == -1 {
				return false
			}
		}
	case strings.IndexByte(rgba, s[0]) >= 0:
		for _, c := range s {
			if strings.IndexRune(rgba, c) == -1 {
				return false
			}
		}
	case strings.IndexByte(strq, s[0]) >= 0:
		for _, c := range s {
			if strings.IndexRune(strq, c) == -1 {
				return false
			}
		}
	}
	return false
}

func (p *Program) structName(t *Type) string {
	if t.Main != Struct {
		panic("shaderir: the given type at structName must be a struct")
	}
	s := t.serialize()
	if n, ok := p.structNames[s]; ok {
		return n
	}
	n := fmt.Sprintf("S%d", len(p.structNames))
	p.structNames[s] = n
	p.structTypes = append(p.structTypes, *t)
	return n
}

func (p *Program) Glsl() string {
	p.structNames = map[string]string{}
	p.structTypes = nil

	var lines []string
	for i, t := range p.Uniforms {
		lines = append(lines, fmt.Sprintf("uniform %s;", p.glslVarDecl(&t, fmt.Sprintf("U%d", i))))
	}
	for i, t := range p.Attributes {
		lines = append(lines, fmt.Sprintf("attribute %s;", p.glslVarDecl(&t, fmt.Sprintf("A%d", i))))
	}
	for i, t := range p.Varyings {
		lines = append(lines, fmt.Sprintf("varying %s;", p.glslVarDecl(&t, fmt.Sprintf("V%d", i))))
	}
	for _, f := range p.Funcs {
		lines = append(lines, p.glslFunc(&f)...)
	}

	// Vertex func
	if len(p.VertexFunc.Block.Stmts) > 0 {
		lines = append(lines, "#if defined(COMPILING_VERTEX_SHADER)")
		lines = append(lines, "void main(void) {")
		lines = append(lines, p.glslBlock(&p.VertexFunc.Block, 0, 0)...)
		lines = append(lines, "}")
		lines = append(lines, "#endif")
	}

	// Fragment func
	if len(p.FragmentFunc.Block.Stmts) > 0 {
		lines = append(lines, "#if defined(COMPILING_FRAGMENT_SHADER)")
		lines = append(lines, "void main(void) {")
		lines = append(lines, p.glslBlock(&p.FragmentFunc.Block, 0, 0)...)
		lines = append(lines, "}")
		lines = append(lines, "#endif")
	}

	var stLines []string
	for i, t := range p.structTypes {
		stLines = append(stLines, fmt.Sprintf("struct S%d {", i))
		for j, st := range t.Sub {
			stLines = append(stLines, fmt.Sprintf("\t%s;", p.glslVarDecl(&st, fmt.Sprintf("M%d", j))))
		}
		stLines = append(stLines, "};")
	}
	lines = append(stLines, lines...)

	return strings.Join(lines, "\n") + "\n"
}

func (p *Program) glslType(t *Type) string {
	switch t.Main {
	case None:
		return "void"
	case Array:
		panic("not implemented")
	case Struct:
		return p.structName(t)
	default:
		return t.Main.Glsl()
	}
}

func (p *Program) glslVarDecl(t *Type, varname string) string {
	switch t.Main {
	case None:
		return "?(none)"
	case Array:
		panic("not implemented")
	case Struct:
		return fmt.Sprintf("%s %s", p.structName(t), varname)
	default:
		return fmt.Sprintf("%s %s", t.Main.Glsl(), varname)
	}
}

func (p *Program) glslFunc(f *Func) []string {
	var args []string
	var idx int
	for _, t := range f.InParams {
		args = append(args, "in "+p.glslVarDecl(&t, fmt.Sprintf("l%d", idx)))
		idx++
	}
	for _, t := range f.InOutParams {
		args = append(args, "inout "+p.glslVarDecl(&t, fmt.Sprintf("l%d", idx)))
		idx++
	}
	for _, t := range f.OutParams {
		args = append(args, "out "+p.glslVarDecl(&t, fmt.Sprintf("l%d", idx)))
		idx++
	}
	argsstr := "void"
	if len(args) > 0 {
		argsstr = strings.Join(args, ", ")
	}

	var lines []string
	lines = append(lines, fmt.Sprintf("%s %s(%s) {", p.glslType(&f.Return), f.Name, argsstr))
	lines = append(lines, p.glslBlock(&f.Block, 0, idx)...)
	lines = append(lines, "}")

	return lines
}

func (p *Program) glslBlock(b *Block, level int, localVarIndex int) []string {
	idt := strings.Repeat("\t", level+1)

	var lines []string
	for _, t := range b.LocalVars {
		lines = append(lines, fmt.Sprintf("%s%s;", idt, p.glslVarDecl(&t, fmt.Sprintf("l%d", localVarIndex))))
		localVarIndex++
	}

	var glslExpr func(e *Expr) string
	glslExpr = func(e *Expr) string {
		switch e.Type {
		case IntExpr:
			return fmt.Sprintf("%d", e.Int)
		case FloatExpr:
			return fmt.Sprintf("%.9e", e.Float)
		case VarName:
			switch e.Variable.Type {
			case Uniform:
				return fmt.Sprintf("U%d", e.Variable.Index)
			case Local:
				idx := e.Variable.Index
				switch b {
				case &p.VertexFunc.Block:
					na := len(p.Attributes)
					nv := len(p.Varyings)
					switch {
					case idx < na:
						return fmt.Sprintf("A%d", idx)
					case idx < na+nv:
						return fmt.Sprintf("V%d", idx-na)
					case idx == na+nv:
						return "gl_Position"
					default:
						return fmt.Sprintf("l%d", idx-(na+nv+1))
					}
				case &p.FragmentFunc.Block:
					nv := len(p.Varyings)
					switch {
					case idx < nv:
						return fmt.Sprintf("V%d", idx)
					case idx == nv:
						return "gl_FragCoord"
					default:
						return fmt.Sprintf("l%d", idx-(nv+1))
					}
				default:
					return fmt.Sprintf("l%d", idx)
				}
			default:
				return fmt.Sprintf("?(unexpected variable type: %d)", e.Variable.Type)
			}
		case BuiltinFuncExpr:
			return string(e.BuiltinFunc)
		case SwizzlingExpr:
			if isValidSwizzling(e.Swizzling) {
				return fmt.Sprintf("?(unexpected swizzling: %s)", e.Swizzling)
			}
			return e.Swizzling
		case Ident:
			return e.Ident
		case Unary:
			var op string
			switch e.Op {
			case Add, Sub, Neg:
				op = string(e.Op)
			default:
				op = fmt.Sprintf("?(unexpected op: %s)", string(e.Op))
			}
			return fmt.Sprintf("%s(%s)", op, glslExpr(&e.Exprs[0]))
		case Binary:
			return fmt.Sprintf("(%s) %s (%s)", glslExpr(&e.Exprs[0]), e.Op, glslExpr(&e.Exprs[1]))
		case Selection:
			return fmt.Sprintf("(%s) ? (%s) : (%s)", glslExpr(&e.Exprs[0]), glslExpr(&e.Exprs[1]), glslExpr(&e.Exprs[2]))
		case Call:
			var args []string
			for _, exp := range e.Exprs[1:] {
				args = append(args, glslExpr(&exp))
			}
			return fmt.Sprintf("(%s)(%s)", glslExpr(&e.Exprs[0]), strings.Join(args, ", "))
		case FieldSelector:
			return fmt.Sprintf("(%s).%s", glslExpr(&e.Exprs[0]), glslExpr(&e.Exprs[1]))
		case Index:
			return fmt.Sprintf("(%s)[%s]", glslExpr(&e.Exprs[0]), glslExpr(&e.Exprs[1]))
		default:
			return fmt.Sprintf("?(unexpected expr: %d)", e.Type)
		}
	}

	for _, s := range b.Stmts {
		switch s.Type {
		case ExprStmt:
			lines = append(lines, fmt.Sprintf("%s%s;", idt, glslExpr(&s.Exprs[0])))
		case BlockStmt:
			lines = append(lines, idt+"{")
			lines = append(lines, p.glslBlock(&s.Blocks[0], level+1, localVarIndex)...)
			lines = append(lines, idt+"}")
		case Assign:
			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, glslExpr(&s.Exprs[0]), glslExpr(&s.Exprs[1])))
		case If:
			lines = append(lines, fmt.Sprintf("%sif (%s) {", idt, glslExpr(&s.Exprs[0])))
			lines = append(lines, p.glslBlock(&s.Blocks[0], level+1, localVarIndex)...)
			if len(s.Blocks) > 1 {
				lines = append(lines, fmt.Sprintf("%s} else {", idt))
				lines = append(lines, p.glslBlock(&s.Blocks[1], level+1, localVarIndex)...)
			}
			lines = append(lines, fmt.Sprintf("%s}", idt))
		case For:
			v := localVarIndex
			localVarIndex++
			var delta string
			switch s.ForDelta {
			case 0:
				delta = fmt.Sprintf("?(unexpected delta: %d)", s.ForDelta)
			case 1:
				delta = fmt.Sprintf("l%d++", v)
			case -1:
				delta = fmt.Sprintf("l%d--", v)
			default:
				if s.ForDelta > 0 {
					delta = fmt.Sprintf("l%d += %d", v, s.ForDelta)
				} else {
					delta = fmt.Sprintf("l%d -= %d", v, -s.ForDelta)
				}
			}
			var op string
			switch s.ForOp {
			case LessThanOp, LessThanEqualOp, GreaterThanOp, GreaterThanEqualOp, EqualOp, NotEqualOp:
				op = string(s.ForOp)
			default:
				op = fmt.Sprintf("?(unexpected op: %s)", string(s.ForOp))
			}
			lines = append(lines, fmt.Sprintf("%sfor (int l%d = %d; l%d %s %d; %s) {", idt, v, s.ForInit, v, op, s.ForEnd, delta))
			lines = append(lines, p.glslBlock(&s.Blocks[0], level+1, localVarIndex)...)
			lines = append(lines, fmt.Sprintf("%s}", idt))
		case Continue:
			lines = append(lines, idt+"continue;")
		case Break:
			lines = append(lines, idt+"break;")
		case Return:
			if len(s.Exprs) == 0 {
				lines = append(lines, idt+"return;")
			} else {
				lines = append(lines, fmt.Sprintf("%sreturn %s;", idt, glslExpr(&s.Exprs[0])))
			}
		case Discard:
			lines = append(lines, idt+"discard;")
		default:
			lines = append(lines, fmt.Sprintf("%s?(unexpected stmt: %d)", idt, s.Type))
		}
	}

	return lines
}
<commit_msg>shaderir: Bug fix: Wrong implementation of isValidSwizzling<commit_after>// Copyright 2020 The Ebiten Authors
//
// 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 shaderir

import (
	"fmt"
	"strings"
)

func isValidSwizzling(s string) bool {
	if len(s) < 1 || 4 < len(s) {
		return false
	}

	const (
		xyzw = "xyzw"
		rgba = "rgba"
		strq = "strq"
	)

	switch {
	case strings.IndexByte(xyzw, s[0]) >= 0:
		for _, c := range s {
			if strings.IndexRune(xyzw, c) == -1 {
				return false
			}
		}
		return true
	case strings.IndexByte(rgba, s[0]) >= 0:
		for _, c := range s {
			if strings.IndexRune(rgba, c) == -1 {
				return false
			}
		}
		return true
	case strings.IndexByte(strq, s[0]) >= 0:
		for _, c := range s {
			if strings.IndexRune(strq, c) == -1 {
				return false
			}
		}
		return true
	}
	return false
}

func (p *Program) structName(t *Type) string {
	if t.Main != Struct {
		panic("shaderir: the given type at structName must be a struct")
	}
	s := t.serialize()
	if n, ok := p.structNames[s]; ok {
		return n
	}
	n := fmt.Sprintf("S%d", len(p.structNames))
	p.structNames[s] = n
	p.structTypes = append(p.structTypes, *t)
	return n
}

func (p *Program) Glsl() string {
	p.structNames = map[string]string{}
	p.structTypes = nil

	var lines []string
	for i, t := range p.Uniforms {
		lines = append(lines, fmt.Sprintf("uniform %s;", p.glslVarDecl(&t, fmt.Sprintf("U%d", i))))
	}
	for i, t := range p.Attributes {
		lines = append(lines, fmt.Sprintf("attribute %s;", p.glslVarDecl(&t, fmt.Sprintf("A%d", i))))
	}
	for i, t := range p.Varyings {
		lines = append(lines, fmt.Sprintf("varying %s;", p.glslVarDecl(&t, fmt.Sprintf("V%d", i))))
	}
	for _, f := range p.Funcs {
		lines = append(lines, p.glslFunc(&f)...)
	}

	// Vertex func
	if len(p.VertexFunc.Block.Stmts) > 0 {
		lines = append(lines, "#if defined(COMPILING_VERTEX_SHADER)")
		lines = append(lines, "void main(void) {")
		lines = append(lines, p.glslBlock(&p.VertexFunc.Block, 0, 0)...)
		lines = append(lines, "}")
		lines = append(lines, "#endif")
	}

	// Fragment func
	if len(p.FragmentFunc.Block.Stmts) > 0 {
		lines = append(lines, "#if defined(COMPILING_FRAGMENT_SHADER)")
		lines = append(lines, "void main(void) {")
		lines = append(lines, p.glslBlock(&p.FragmentFunc.Block, 0, 0)...)
		lines = append(lines, "}")
		lines = append(lines, "#endif")
	}

	var stLines []string
	for i, t := range p.structTypes {
		stLines = append(stLines, fmt.Sprintf("struct S%d {", i))
		for j, st := range t.Sub {
			stLines = append(stLines, fmt.Sprintf("\t%s;", p.glslVarDecl(&st, fmt.Sprintf("M%d", j))))
		}
		stLines = append(stLines, "};")
	}
	lines = append(stLines, lines...)

	return strings.Join(lines, "\n") + "\n"
}

func (p *Program) glslType(t *Type) string {
	switch t.Main {
	case None:
		return "void"
	case Array:
		panic("not implemented")
	case Struct:
		return p.structName(t)
	default:
		return t.Main.Glsl()
	}
}

func (p *Program) glslVarDecl(t *Type, varname string) string {
	switch t.Main {
	case None:
		return "?(none)"
	case Array:
		panic("not implemented")
	case Struct:
		return fmt.Sprintf("%s %s", p.structName(t), varname)
	default:
		return fmt.Sprintf("%s %s", t.Main.Glsl(), varname)
	}
}

func (p *Program) glslFunc(f *Func) []string {
	var args []string
	var idx int
	for _, t := range f.InParams {
		args = append(args, "in "+p.glslVarDecl(&t, fmt.Sprintf("l%d", idx)))
		idx++
	}
	for _, t := range f.InOutParams {
		args = append(args, "inout "+p.glslVarDecl(&t, fmt.Sprintf("l%d", idx)))
		idx++
	}
	for _, t := range f.OutParams {
		args = append(args, "out "+p.glslVarDecl(&t, fmt.Sprintf("l%d", idx)))
		idx++
	}
	argsstr := "void"
	if len(args) > 0 {
		argsstr = strings.Join(args, ", ")
	}

	var lines []string
	lines = append(lines, fmt.Sprintf("%s %s(%s) {", p.glslType(&f.Return), f.Name, argsstr))
	lines = append(lines, p.glslBlock(&f.Block, 0, idx)...)
	lines = append(lines, "}")

	return lines
}

func (p *Program) glslBlock(b *Block, level int, localVarIndex int) []string {
	idt := strings.Repeat("\t", level+1)

	var lines []string
	for _, t := range b.LocalVars {
		lines = append(lines, fmt.Sprintf("%s%s;", idt, p.glslVarDecl(&t, fmt.Sprintf("l%d", localVarIndex))))
		localVarIndex++
	}

	var glslExpr func(e *Expr) string
	glslExpr = func(e *Expr) string {
		switch e.Type {
		case IntExpr:
			return fmt.Sprintf("%d", e.Int)
		case FloatExpr:
			return fmt.Sprintf("%.9e", e.Float)
		case VarName:
			switch e.Variable.Type {
			case Uniform:
				return fmt.Sprintf("U%d", e.Variable.Index)
			case Local:
				idx := e.Variable.Index
				switch b {
				case &p.VertexFunc.Block:
					na := len(p.Attributes)
					nv := len(p.Varyings)
					switch {
					case idx < na:
						return fmt.Sprintf("A%d", idx)
					case idx < na+nv:
						return fmt.Sprintf("V%d", idx-na)
					case idx == na+nv:
						return "gl_Position"
					default:
						return fmt.Sprintf("l%d", idx-(na+nv+1))
					}
				case &p.FragmentFunc.Block:
					nv := len(p.Varyings)
					switch {
					case idx < nv:
						return fmt.Sprintf("V%d", idx)
					case idx == nv:
						return "gl_FragCoord"
					default:
						return fmt.Sprintf("l%d", idx-(nv+1))
					}
				default:
					return fmt.Sprintf("l%d", idx)
				}
			default:
				return fmt.Sprintf("?(unexpected variable type: %d)", e.Variable.Type)
			}
		case BuiltinFuncExpr:
			return string(e.BuiltinFunc)
		case SwizzlingExpr:
			if !isValidSwizzling(e.Swizzling) {
				return fmt.Sprintf("?(unexpected swizzling: %s)", e.Swizzling)
			}
			return e.Swizzling
		case Ident:
			return e.Ident
		case Unary:
			var op string
			switch e.Op {
			case Add, Sub, Neg:
				op = string(e.Op)
			default:
				op = fmt.Sprintf("?(unexpected op: %s)", string(e.Op))
			}
			return fmt.Sprintf("%s(%s)", op, glslExpr(&e.Exprs[0]))
		case Binary:
			return fmt.Sprintf("(%s) %s (%s)", glslExpr(&e.Exprs[0]), e.Op, glslExpr(&e.Exprs[1]))
		case Selection:
			return fmt.Sprintf("(%s) ? (%s) : (%s)", glslExpr(&e.Exprs[0]), glslExpr(&e.Exprs[1]), glslExpr(&e.Exprs[2]))
		case Call:
			var args []string
			for _, exp := range e.Exprs[1:] {
				args = append(args, glslExpr(&exp))
			}
			return fmt.Sprintf("(%s)(%s)", glslExpr(&e.Exprs[0]), strings.Join(args, ", "))
		case FieldSelector:
			return fmt.Sprintf("(%s).%s", glslExpr(&e.Exprs[0]), glslExpr(&e.Exprs[1]))
		case Index:
			return fmt.Sprintf("(%s)[%s]", glslExpr(&e.Exprs[0]), glslExpr(&e.Exprs[1]))
		default:
			return fmt.Sprintf("?(unexpected expr: %d)", e.Type)
		}
	}

	for _, s := range b.Stmts {
		switch s.Type {
		case ExprStmt:
			lines = append(lines, fmt.Sprintf("%s%s;", idt, glslExpr(&s.Exprs[0])))
		case BlockStmt:
			lines = append(lines, idt+"{")
			lines = append(lines, p.glslBlock(&s.Blocks[0], level+1, localVarIndex)...)
			lines = append(lines, idt+"}")
		case Assign:
			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, glslExpr(&s.Exprs[0]), glslExpr(&s.Exprs[1])))
		case If:
			lines = append(lines, fmt.Sprintf("%sif (%s) {", idt, glslExpr(&s.Exprs[0])))
			lines = append(lines, p.glslBlock(&s.Blocks[0], level+1, localVarIndex)...)
			if len(s.Blocks) > 1 {
				lines = append(lines, fmt.Sprintf("%s} else {", idt))
				lines = append(lines, p.glslBlock(&s.Blocks[1], level+1, localVarIndex)...)
			}
			lines = append(lines, fmt.Sprintf("%s}", idt))
		case For:
			v := localVarIndex
			localVarIndex++
			var delta string
			switch s.ForDelta {
			case 0:
				delta = fmt.Sprintf("?(unexpected delta: %d)", s.ForDelta)
			case 1:
				delta = fmt.Sprintf("l%d++", v)
			case -1:
				delta = fmt.Sprintf("l%d--", v)
			default:
				if s.ForDelta > 0 {
					delta = fmt.Sprintf("l%d += %d", v, s.ForDelta)
				} else {
					delta = fmt.Sprintf("l%d -= %d", v, -s.ForDelta)
				}
			}
			var op string
			switch s.ForOp {
			case LessThanOp, LessThanEqualOp, GreaterThanOp, GreaterThanEqualOp, EqualOp, NotEqualOp:
				op = string(s.ForOp)
			default:
				op = fmt.Sprintf("?(unexpected op: %s)", string(s.ForOp))
			}
			lines = append(lines, fmt.Sprintf("%sfor (int l%d = %d; l%d %s %d; %s) {", idt, v, s.ForInit, v, op, s.ForEnd, delta))
			lines = append(lines, p.glslBlock(&s.Blocks[0], level+1, localVarIndex)...)
			lines = append(lines, fmt.Sprintf("%s}", idt))
		case Continue:
			lines = append(lines, idt+"continue;")
		case Break:
			lines = append(lines, idt+"break;")
		case Return:
			if len(s.Exprs) == 0 {
				lines = append(lines, idt+"return;")
			} else {
				lines = append(lines, fmt.Sprintf("%sreturn %s;", idt, glslExpr(&s.Exprs[0])))
			}
		case Discard:
			lines = append(lines, idt+"discard;")
		default:
			lines = append(lines, fmt.Sprintf("%s?(unexpected stmt: %d)", idt, s.Type))
		}
	}

	return lines
}
<|endoftext|>


<commit_before>package virt

import (
	"context"
	"errors"
	"io"

	"github.com/jimmyfrasche/etlite/internal/internal/errint"
	"github.com/jimmyfrasche/etlite/internal/internal/errusr"
	"github.com/jimmyfrasche/etlite/internal/token"
	"github.com/jimmyfrasche/etlite/internal/virt/internal/builder"
)

type ImportSpec struct {
	Pos            token.Position
	Internal, Temp bool
	Table, Frame   string
	Header         []string
	DDL            string
	Limit, Offset  int
}

func (s ImportSpec) Valid() error {
	if s.DDL != "" {
		if s.Internal {
			return errint.New("CREATE TABLE FROM marked as internal table")
		}
		if s.Temp {
			return errint.New("CREATE TABLE FROM marked as temporary table (by system, not user)")
		}
		if len(s.Header) > 0 {
			return errint.New("CREATE TABLE FROM provides user defined header")
		}
		if s.Table == "" {
			return errint.New("CREATE TABLE FROM provides no table name")
		}
	} else if s.Internal && s.Temp {
		return errint.New("internal table marked as temporary (by system, not user)")
	}
	if s.Table == "" {
		return errint.New("table did not provide name")
	}
	return nil
}

//XXX have compiler write all SQL and just pass select/insert to export/import? what about derived table names then?

func Import(s ImportSpec) Instruction {
	return func(ctx context.Context, m *Machine) error {
		if err := s.Valid(); err != nil {
			return err
		}

		//CREATE TABLE FROM
		if s.DDL != "" {
			//build table and compute header
			if err := m.exec(s.DDL); err != nil {
				return errusr.Wrap(s.Pos, err)
			}
			p, err := m.conn.Prepare("SELECT * FROM " + s.Table)
			if err != nil {
				return err
			}
			s.Header = p.Columns()
			if len(s.Header) == 0 {
				return errint.New("could not retrieve columns in CREATE TABLE FROM")
			}
			if err := p.Close(); err != nil {
				return err
			}
		}

		//prep the decoder
		inHeader, err := m.readHeader(s.Table, s.Frame, s.Header)
		if err != nil {
			return err
		}

		//no header specified, use the one from the decoder
		if len(s.Header) == 0 {
			if len(inHeader) == 0 {
				return errors.New("no header specified and none returned by " + m.decoder.Name() + " format")
			}
			s.Header = inHeader
		}

		//internal tables can be created en masse so have to take care of their own savepointing.
		if !s.Internal {
			if err := m.savepoint(); err != nil {
				return err
			}
		}

		//not a CREATE TABLE FROM so we need to make the table
		if s.DDL == "" {
			if err := m.createTable(s.Temp || s.Internal, s.Table, s.Header); err != nil {
				return err
			}
		}

		if err := m.bulkInsert(s.Table, s.Header, s.Limit, s.Offset); err != nil {
			return err
		}

		if !s.Internal {
			if err := m.release(); err != nil {
				return err
			}
		}

		return nil
	}
}

func (m *Machine) readHeader(table, frame string, header []string) ([]string, error) {
	dec := m.decoder
	if dec == nil {
		return nil, errint.New("no decoder available when importing")
	}

	if frame != "" {
		table = frame
	}

	return dec.ReadHeader(table, header)
}

//CreateTable initializes the decoder then creates table name with header.
//
//InitDecoder must be called before this.
//
//See MkCreateTableFrom and BulkInsert.
func (m *Machine) createTable(temp bool, name string, header []string) error {
	//create ddl
	b := builder.New("CREATE")

	if temp {
		b.Push("TEMPORARY")
	}

	b.Push("TABLE", name, "(")

	b.CSV(header, func(h string) {
		b.Push(h, "TEXT")
	})

	b.Push(")")

	return m.exec(b.Join(" "))
}

//BulkInsert from the current decoder into table name with header.
//
//Table name must exist and should have been created by
//either CreateTableFrom or CreateTable, with no intervening
//reads to or changes of the decoder.
//
//Before that DecodeHeader must be called.
func (m *Machine) bulkInsert(name string, header []string, limit, offset int) error {
	//make sure we have a decoder
	dec := m.decoder
	if dec == nil {
		return errint.Newf("no decoder when attempting to import %s", name)
	}

	//build the bulk loader
	b := builder.New("INSERT INTO", name, "(")

	b.CSV(header, func(h string) {
		b.Push(h)
	})

	b.Push(") VALUES (")

	b.CSV(header, func(string) {
		b.Push("?")
	})
	b.Push(")")

	p, err := m.conn.Prepare(b.Join(" "))
	if err != nil {
		return err
	}
	defer p.Close()

	bulk, err := p.Loader()
	if err != nil {
		return err
	}

	if offset > 0 {
		if err := dec.Skip(offset); err != nil {
			return err
		}
	}

	for rows := 0; limit > 0 && rows == limit || true; rows++ {
		row, err := dec.ReadRow()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		if err := bulk.Load(row); err != nil {
			return err
		}
	}

	if err := bulk.Close(); err != nil {
		return err
	}

	return dec.Reset()
}
<commit_msg>remove incorrect helper<commit_after>package virt

import (
	"context"
	"errors"
	"io"

	"github.com/jimmyfrasche/etlite/internal/internal/errint"
	"github.com/jimmyfrasche/etlite/internal/internal/errusr"
	"github.com/jimmyfrasche/etlite/internal/token"
	"github.com/jimmyfrasche/etlite/internal/virt/internal/builder"
)

type ImportSpec struct {
	Pos            token.Position
	Internal, Temp bool
	Table, Frame   string
	Header         []string
	DDL            string
	Limit, Offset  int
}

func (s ImportSpec) Valid() error {
	if s.DDL != "" {
		if s.Internal {
			return errint.New("CREATE TABLE FROM marked as internal table")
		}
		if s.Temp {
			return errint.New("CREATE TABLE FROM marked as temporary table (by system, not user)")
		}
		if len(s.Header) > 0 {
			return errint.New("CREATE TABLE FROM provides user defined header")
		}
		if s.Table == "" {
			return errint.New("CREATE TABLE FROM provides no table name")
		}
	} else if s.Internal && s.Temp {
		return errint.New("internal table marked as temporary (by system, not user)")
	}
	if s.Table == "" {
		return errint.New("table did not provide name")
	}
	return nil
}

//XXX have compiler write all SQL and just pass select/insert to export/import? what about derived table names then?

func Import(s ImportSpec) Instruction {
	return func(ctx context.Context, m *Machine) error {
		if err := s.Valid(); err != nil {
			return err
		}

		//CREATE TABLE FROM
		if s.DDL != "" {
			//build table and compute header
			if err := m.exec(s.DDL); err != nil {
				return errusr.Wrap(s.Pos, err)
			}
			p, err := m.conn.Prepare("SELECT * FROM " + s.Table)
			if err != nil {
				return err
			}
			s.Header = p.Columns()
			if len(s.Header) == 0 {
				return errint.New("could not retrieve columns in CREATE TABLE FROM")
			}
			if err := p.Close(); err != nil {
				return err
			}
		}

		//prep the decoder
		dec := m.decoder
		if dec == nil {
			return errint.New("no decoder available when importing")
		}
		inHeader, err := dec.ReadHeader(s.Frame, s.Header)
		if err != nil {
			return err
		}

		//no header specified, use the one from the decoder
		if len(s.Header) == 0 {
			if len(inHeader) == 0 {
				return errors.New("no header specified and none returned by " + m.decoder.Name() + " format")
			}
			s.Header = inHeader
		}

		//internal tables can be created en masse so have to take care of their own savepointing.
		if !s.Internal {
			if err := m.savepoint(); err != nil {
				return err
			}
		}

		//not a CREATE TABLE FROM so we need to make the table
		if s.DDL == "" {
			if err := m.createTable(s.Temp || s.Internal, s.Table, s.Header); err != nil {
				return err
			}
		}

		if err := m.bulkInsert(s.Table, s.Header, s.Limit, s.Offset); err != nil {
			return err
		}

		if !s.Internal {
			if err := m.release(); err != nil {
				return err
			}
		}

		return nil
	}
}

//CreateTable initializes the decoder then creates table name with header.
//
//InitDecoder must be called before this.
//
//See MkCreateTableFrom and BulkInsert.
func (m *Machine) createTable(temp bool, name string, header []string) error {
	//create ddl
	b := builder.New("CREATE")

	if temp {
		b.Push("TEMPORARY")
	}

	b.Push("TABLE", name, "(")

	b.CSV(header, func(h string) {
		b.Push(h, "TEXT")
	})

	b.Push(")")

	return m.exec(b.Join(" "))
}

//BulkInsert from the current decoder into table name with header.
//
//Table name must exist and should have been created by
//either CreateTableFrom or CreateTable, with no intervening
//reads to or changes of the decoder.
//
//Before that DecodeHeader must be called.
func (m *Machine) bulkInsert(name string, header []string, limit, offset int) error {
	//make sure we have a decoder
	dec := m.decoder
	if dec == nil {
		return errint.Newf("no decoder when attempting to import %s", name)
	}

	//build the bulk loader
	b := builder.New("INSERT INTO", name, "(")

	b.CSV(header, func(h string) {
		b.Push(h)
	})

	b.Push(") VALUES (")

	b.CSV(header, func(string) {
		b.Push("?")
	})
	b.Push(")")

	p, err := m.conn.Prepare(b.Join(" "))
	if err != nil {
		return err
	}
	defer p.Close()

	bulk, err := p.Loader()
	if err != nil {
		return err
	}

	if offset > 0 {
		if err := dec.Skip(offset); err != nil {
			return err
		}
	}

	for rows := 0; limit > 0 && rows == limit || true; rows++ {
		row, err := dec.ReadRow()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		if err := bulk.Load(row); err != nil {
			return err
		}
	}

	if err := bulk.Close(); err != nil {
		return err
	}

	return dec.Reset()
}
<|endoftext|>


<commit_before>package session

import (
	"encoding/binary"
	"encoding/hex"
	"errors"
	"io"
	"net"
	"strconv"

	"github.com/garyburd/redigo/redis"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/rs/xid"
	"github.com/sirupsen/logrus"
	"github.com/superfly/wormhole/messages"
	"golang.org/x/crypto/ssh"

	wnet "github.com/superfly/wormhole/net"
)

const (
	sshRemoteForwardRequest      = "tcpip-forward"
	sshForwardedTCPReturnRequest = "forwarded-tcpip"
)

var (
	openSessionsMetric = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "open_sessions_total",
			Help:      "Number of active sessions, partitioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this session belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	openChannelsMetric = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "open_channels_total",
			Help:      "Number of active channels, partitioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	ingressConnDurationMetric = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace:  "wormhole",
			Subsystem:  "ssh_session",
			Name:       "ingress_conn_duration_seconds",
			Help:       "Duration in seconds of ingress connections, paritioned by backend, node and cluster.",
			Objectives: map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	ingressConnRcvdBytesMetric = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "ingress_conn_rcvd_bytes",
			Help:      "Number of bytes received from ingress connections, paritioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	ingressConnSentBytesMetric = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "ingress_conn_sent_bytes",
			Help:      "Number of bytes sent to ingress connections, paritioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)
)

func init() {
	prometheus.MustRegister(openSessionsMetric)
	prometheus.MustRegister(openChannelsMetric)
	prometheus.MustRegister(ingressConnDurationMetric)
	prometheus.MustRegister(ingressConnRcvdBytesMetric)
	prometheus.MustRegister(ingressConnSentBytesMetric)
}

// SSHSession extends information about connected client stored in Session.
// It also includes SSH-specific information like the SSH conn, SSH server config, etc.
type SSHSession struct {
	baseSession

	sshSessionID string
	config       *ssh.ServerConfig
	tcpConn      net.Conn
	conn         *ssh.ServerConn
	reqs         <-chan *ssh.Request
	chans        <-chan ssh.NewChannel
}

type tcpipForward struct {
	Host string
	Port uint32
}

type directForward struct {
	Host1 string
	Port1 uint32
	Host2 string
	Port2 uint32
}

type instrumentedIO struct {
	rwc        io.ReadWriteCloser
	metricFunc func()
}

func (io instrumentedIO) Close() error {
	go io.metricFunc()
	return io.rwc.Close()
}

func (io instrumentedIO) Read(p []byte) (int, error) {
	return io.rwc.Read(p)
}

func (io instrumentedIO) Write(p []byte) (int, error) {
	return io.rwc.Write(p)
}

// NewSSHSession creates new SshSession struct
func NewSSHSession(logger *logrus.Logger, clusterURL, nodeID string, redisPool *redis.Pool, tcpConn net.Conn, config *ssh.ServerConfig) *SSHSession {
	base := baseSession{
		id:         xid.New().String(),
		nodeID:     nodeID,
		store:      NewRedisStore(redisPool),
		ClusterURL: clusterURL,
		logger:     logger.WithFields(logrus.Fields{"prefix": "SSHSession"}),
	}
	s := &SSHSession{
		tcpConn:     tcpConn,
		baseSession: base,
	}
	config.PasswordCallback = s.authFromToken
	s.config = config
	return s
}

// RequireStream performs SSH handshake and ensures SSHSession is ready to receive
// and send data
func (s *SSHSession) RequireStream() error {
	// Before use, a handshake must be performed on the incoming net.Conn.
	sshConn, chans, reqs, err := ssh.NewServerConn(s.tcpConn, s.config)
	if err != nil {
		return err
	}
	s.conn = sshConn
	s.chans = chans
	s.reqs = reqs
	go handleChannels(chans)
	go openSessionsMetric.With(labels(s)).Add(1)
	return nil
}

// HandleRequests handles all requests coming over the SSH connection from the client.
// The main function is to accept ingress traffic (from the listener) once the remote port
// forwarding is set up.
// It also handles out-of-band SSH request types, like the keepalive or register-release.
func (s *SSHSession) HandleRequests(ln net.Listener) {
	for req := range s.reqs {
		switch req.Type {
		case sshRemoteForwardRequest:
			go func() {
				s.handleRemoteForward(req, ln)
			}()
		case "register-release":
			go s.registerRelease(req)
		case "keepalive":
			go s.handleKeepalive(req)
		}
	}
}

// RequireAuthentication registers the connection, since authentication is part of the SSH handshake
// TODO: figure out a better interface for Session
func (s *SSHSession) RequireAuthentication() error {
	// done as a hook to ssh handshake
	go s.store.RegisterConnection(s)
	return nil
}

// RegisterEndpoint registers the endpoint and adds it to the current session record
// The endpoint is a particular instance of a running wormhole client
func (s *SSHSession) RegisterEndpoint() error {
	return s.store.RegisterEndpoint(s)
}

// Close closes SSHSession and registers disconnection
func (s *SSHSession) Close() {
	s.store.RegisterDisconnection(s)
	s.logger.Infof("Closed session %s for %s %s (%s).", s.ID(), s.NodeID(), s.Agent(), s.Client())
	go func() {
		openSessionsMetric.With(labels(s)).Sub(1)
	}()
	s.conn.Close()
}

func (s *SSHSession) authFromToken(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
	backendID, err := s.store.BackendIDFromToken(string(pass))
	if err != nil && err != redis.ErrNil {
		return nil, err
	}
	if backendID == "" {
		return nil, errors.New("token '" + string(pass) + "' rejected")
	}

	// assume false if not set
	requiresClientAuth, err := s.store.BackendRequiresClientAuth(backendID)
	if err != nil {
		return nil, err
	}

	s.backendID = backendID
	s.agent = string(c.ClientVersion())
	s.sshSessionID = hex.EncodeToString(c.SessionID())
	s.requiresClientAuth = requiresClientAuth
	s.clientAddr = c.RemoteAddr().String()

	return nil, nil
}

func (s *SSHSession) setSSHPort(req *ssh.Request, ln net.Listener) tcpipForward {
	t := tcpipForward{}
	ssh.Unmarshal(req.Payload, &t)

	reply := (t.Port == 0) && req.WantReply

	if reply { // Client sent port 0. let them know which port is actually being used
		_, port, _ := net.SplitHostPort(ln.Addr().String())
		portNum, _ := strconv.Atoi(port)
		b := make([]byte, 4)
		binary.BigEndian.PutUint32(b, uint32(portNum))
		t.Port = uint32(portNum)
		req.Reply(true, b)
	} else {
		req.Reply(true, nil)
	}

	return t
}

func (s *SSHSession) handleRemoteForward(req *ssh.Request, ln net.Listener) {
	defer func() {
		err := ln.Close()
		if err != nil {
			s.logger.Debugf("Couldn't close ingress conn: %s", err)
			return
		}
		s.logger.Debugf("Closed ingress conn: %s", ln.Addr().String())
	}()

	t := s.setSSHPort(req, ln)
	p := directForward{}
	quit := make(chan bool)
	go func(ln net.Listener) { // Handle incoming connections on this new listener
		for {
			select {
			case <-quit:
				return
			default:
				ingressConn, err := ln.Accept()

				if err != nil {
					netErr, ok := err.(net.Error)

					//If this is a timeout, then continue to wait for
					//new connections
					if ok && netErr.Timeout() && netErr.Temporary() {
						continue
					}
					s.logger.Errorln("Could not accept Ingress TCP conn:", err)
					return
				}
				s.logger.Debugln("Accepted Ingress TCP conn from:", ingressConn.RemoteAddr())

				host, port, err := net.SplitHostPort(ingressConn.RemoteAddr().String())
				if err != nil {
					return
				}
				portnum, err := strconv.Atoi(port)
				if err != nil {
					return
				}

				p.Host1 = t.Host
				p.Port1 = t.Port
				p.Host2 = host
				p.Port2 = uint32(portnum)

				ch, reqs, err := s.conn.OpenChannel(sshForwardedTCPReturnRequest, ssh.Marshal(p))
				if err != nil {
					s.logger.Errorf("Open forwarded Channel error: %s", err.Error())
					return
				}
				go ssh.DiscardRequests(reqs)
				go openChannelsMetric.With(labels(s)).Add(1)
				go func() {
					chWritten, ingressConnWritten, err := wnet.CopyCloseIO(ch, ingressConn)
					openChannelsMetric.With(labels(s)).Sub(1)
					if connWithMetrics, ok := ingressConn.(*wnet.ServerConnTracker); ok {
						connWithMetrics.ReportDataMetrics(ingressConnWritten, chWritten)
					}
					if err != nil && err != io.EOF {
						s.logger.Error(err)
					}
				}()
			}
		}
	}(ln)

	s.conn.Wait()
	quit <- true
}

func handleChannels(chans <-chan ssh.NewChannel) {
	for range chans {
		// nothing for now.
	}
}

func (s *SSHSession) handleKeepalive(req *ssh.Request) {
	if req.WantReply {
		req.Reply(true, nil)
	}
	go func() {
		if err := s.store.RegisterHeartbeat(s); err != nil {
			s.logger.Warnf("Failed to register session heartbeat: %s", err.Error())
		}
	}()
}

func (s *SSHSession) registerRelease(req *ssh.Request) {
	if req.WantReply {
		req.Reply(true, nil)
	}

	msg, err := messages.Unpack(req.Payload)

	if err != nil {
		s.logger.Warnf("Couldn't process release info: %s", err.Error())
		return
	}

	if release, ok := msg.(*messages.Release); ok {
		s.release = release
		s.store.RegisterRelease(s)
	} else {
		s.logger.Warnf("Couldn't process release info: Unexpected message type")
	}
}

func labels(s Session) prometheus.Labels {
	return prometheus.Labels{"backend": s.BackendID(),
		"node":    s.NodeID(),
		"cluster": s.Cluster(),
	}
}
<commit_msg>trim token spaces<commit_after>package session

import (
	"encoding/binary"
	"encoding/hex"
	"errors"
	"io"
	"net"
	"strconv"
	"strings"

	"github.com/garyburd/redigo/redis"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/rs/xid"
	"github.com/sirupsen/logrus"
	"github.com/superfly/wormhole/messages"
	"golang.org/x/crypto/ssh"

	wnet "github.com/superfly/wormhole/net"
)

const (
	sshRemoteForwardRequest      = "tcpip-forward"
	sshForwardedTCPReturnRequest = "forwarded-tcpip"
)

var (
	openSessionsMetric = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "open_sessions_total",
			Help:      "Number of active sessions, partitioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this session belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	openChannelsMetric = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "open_channels_total",
			Help:      "Number of active channels, partitioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	ingressConnDurationMetric = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace:  "wormhole",
			Subsystem:  "ssh_session",
			Name:       "ingress_conn_duration_seconds",
			Help:       "Duration in seconds of ingress connections, paritioned by backend, node and cluster.",
			Objectives: map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	ingressConnRcvdBytesMetric = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "ingress_conn_rcvd_bytes",
			Help:      "Number of bytes received from ingress connections, paritioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)

	ingressConnSentBytesMetric = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "wormhole",
			Subsystem: "ssh_session",
			Name:      "ingress_conn_sent_bytes",
			Help:      "Number of bytes sent to ingress connections, paritioned by backend, node and cluster.",
		},
		[]string{
			// Which backend this channel belongs to?
			"backend",
			// What wormhole instance this session is running on?
			"node",
			// What region this session belongs to?
			"cluster",
		},
	)
)

func init() {
	prometheus.MustRegister(openSessionsMetric)
	prometheus.MustRegister(openChannelsMetric)
	prometheus.MustRegister(ingressConnDurationMetric)
	prometheus.MustRegister(ingressConnRcvdBytesMetric)
	prometheus.MustRegister(ingressConnSentBytesMetric)
}

// SSHSession extends information about connected client stored in Session.
// It also includes SSH-specific information like the SSH conn, SSH server config, etc.
type SSHSession struct {
	baseSession

	sshSessionID string
	config       *ssh.ServerConfig
	tcpConn      net.Conn
	conn         *ssh.ServerConn
	reqs         <-chan *ssh.Request
	chans        <-chan ssh.NewChannel
}

type tcpipForward struct {
	Host string
	Port uint32
}

type directForward struct {
	Host1 string
	Port1 uint32
	Host2 string
	Port2 uint32
}

type instrumentedIO struct {
	rwc        io.ReadWriteCloser
	metricFunc func()
}

func (io instrumentedIO) Close() error {
	go io.metricFunc()
	return io.rwc.Close()
}

func (io instrumentedIO) Read(p []byte) (int, error) {
	return io.rwc.Read(p)
}

func (io instrumentedIO) Write(p []byte) (int, error) {
	return io.rwc.Write(p)
}

// NewSSHSession creates new SshSession struct
func NewSSHSession(logger *logrus.Logger, clusterURL, nodeID string, redisPool *redis.Pool, tcpConn net.Conn, config *ssh.ServerConfig) *SSHSession {
	base := baseSession{
		id:         xid.New().String(),
		nodeID:     nodeID,
		store:      NewRedisStore(redisPool),
		ClusterURL: clusterURL,
		logger:     logger.WithFields(logrus.Fields{"prefix": "SSHSession"}),
	}
	s := &SSHSession{
		tcpConn:     tcpConn,
		baseSession: base,
	}
	config.PasswordCallback = s.authFromToken
	s.config = config
	return s
}

// RequireStream performs SSH handshake and ensures SSHSession is ready to receive
// and send data
func (s *SSHSession) RequireStream() error {
	// Before use, a handshake must be performed on the incoming net.Conn.
	sshConn, chans, reqs, err := ssh.NewServerConn(s.tcpConn, s.config)
	if err != nil {
		return err
	}
	s.conn = sshConn
	s.chans = chans
	s.reqs = reqs
	go handleChannels(chans)
	go openSessionsMetric.With(labels(s)).Add(1)
	return nil
}

// HandleRequests handles all requests coming over the SSH connection from the client.
// The main function is to accept ingress traffic (from the listener) once the remote port
// forwarding is set up.
// It also handles out-of-band SSH request types, like the keepalive or register-release.
func (s *SSHSession) HandleRequests(ln net.Listener) {
	for req := range s.reqs {
		switch req.Type {
		case sshRemoteForwardRequest:
			go func() {
				s.handleRemoteForward(req, ln)
			}()
		case "register-release":
			go s.registerRelease(req)
		case "keepalive":
			go s.handleKeepalive(req)
		}
	}
}

// RequireAuthentication registers the connection, since authentication is part of the SSH handshake
// TODO: figure out a better interface for Session
func (s *SSHSession) RequireAuthentication() error {
	// done as a hook to ssh handshake
	go s.store.RegisterConnection(s)
	return nil
}

// RegisterEndpoint registers the endpoint and adds it to the current session record
// The endpoint is a particular instance of a running wormhole client
func (s *SSHSession) RegisterEndpoint() error {
	return s.store.RegisterEndpoint(s)
}

// Close closes SSHSession and registers disconnection
func (s *SSHSession) Close() {
	s.store.RegisterDisconnection(s)
	s.logger.Infof("Closed session %s for %s %s (%s).", s.ID(), s.NodeID(), s.Agent(), s.Client())
	go func() {
		openSessionsMetric.With(labels(s)).Sub(1)
	}()
	s.conn.Close()
}

func (s *SSHSession) authFromToken(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
	backendID, err := s.store.BackendIDFromToken(strings.TrimSpace(string(pass)))
	if err != nil && err != redis.ErrNil {
		return nil, err
	}
	if backendID == "" {
		return nil, errors.New("token '" + string(pass) + "' rejected")
	}

	// assume false if not set
	requiresClientAuth, err := s.store.BackendRequiresClientAuth(backendID)
	if err != nil {
		return nil, err
	}

	s.backendID = backendID
	s.agent = string(c.ClientVersion())
	s.sshSessionID = hex.EncodeToString(c.SessionID())
	s.requiresClientAuth = requiresClientAuth
	s.clientAddr = c.RemoteAddr().String()

	return nil, nil
}

func (s *SSHSession) setSSHPort(req *ssh.Request, ln net.Listener) tcpipForward {
	t := tcpipForward{}
	ssh.Unmarshal(req.Payload, &t)

	reply := (t.Port == 0) && req.WantReply

	if reply { // Client sent port 0. let them know which port is actually being used
		_, port, _ := net.SplitHostPort(ln.Addr().String())
		portNum, _ := strconv.Atoi(port)
		b := make([]byte, 4)
		binary.BigEndian.PutUint32(b, uint32(portNum))
		t.Port = uint32(portNum)
		req.Reply(true, b)
	} else {
		req.Reply(true, nil)
	}

	return t
}

func (s *SSHSession) handleRemoteForward(req *ssh.Request, ln net.Listener) {
	defer func() {
		err := ln.Close()
		if err != nil {
			s.logger.Debugf("Couldn't close ingress conn: %s", err)
			return
		}
		s.logger.Debugf("Closed ingress conn: %s", ln.Addr().String())
	}()

	t := s.setSSHPort(req, ln)
	p := directForward{}
	quit := make(chan bool)
	go func(ln net.Listener) { // Handle incoming connections on this new listener
		for {
			select {
			case <-quit:
				return
			default:
				ingressConn, err := ln.Accept()

				if err != nil {
					netErr, ok := err.(net.Error)

					//If this is a timeout, then continue to wait for
					//new connections
					if ok && netErr.Timeout() && netErr.Temporary() {
						continue
					}
					s.logger.Errorln("Could not accept Ingress TCP conn:", err)
					return
				}
				s.logger.Debugln("Accepted Ingress TCP conn from:", ingressConn.RemoteAddr())

				host, port, err := net.SplitHostPort(ingressConn.RemoteAddr().String())
				if err != nil {
					return
				}
				portnum, err := strconv.Atoi(port)
				if err != nil {
					return
				}

				p.Host1 = t.Host
				p.Port1 = t.Port
				p.Host2 = host
				p.Port2 = uint32(portnum)

				ch, reqs, err := s.conn.OpenChannel(sshForwardedTCPReturnRequest, ssh.Marshal(p))
				if err != nil {
					s.logger.Errorf("Open forwarded Channel error: %s", err.Error())
					return
				}
				go ssh.DiscardRequests(reqs)
				go openChannelsMetric.With(labels(s)).Add(1)
				go func() {
					chWritten, ingressConnWritten, err := wnet.CopyCloseIO(ch, ingressConn)
					openChannelsMetric.With(labels(s)).Sub(1)
					if connWithMetrics, ok := ingressConn.(*wnet.ServerConnTracker); ok {
						connWithMetrics.ReportDataMetrics(ingressConnWritten, chWritten)
					}
					if err != nil && err != io.EOF {
						s.logger.Error(err)
					}
				}()
			}
		}
	}(ln)

	s.conn.Wait()
	quit <- true
}

func handleChannels(chans <-chan ssh.NewChannel) {
	for range chans {
		// nothing for now.
	}
}

func (s *SSHSession) handleKeepalive(req *ssh.Request) {
	if req.WantReply {
		req.Reply(true, nil)
	}
	go func() {
		if err := s.store.RegisterHeartbeat(s); err != nil {
			s.logger.Warnf("Failed to register session heartbeat: %s", err.Error())
		}
	}()
}

func (s *SSHSession) registerRelease(req *ssh.Request) {
	if req.WantReply {
		req.Reply(true, nil)
	}

	msg, err := messages.Unpack(req.Payload)

	if err != nil {
		s.logger.Warnf("Couldn't process release info: %s", err.Error())
		return
	}

	if release, ok := msg.(*messages.Release); ok {
		s.release = release
		s.store.RegisterRelease(s)
	} else {
		s.logger.Warnf("Couldn't process release info: Unexpected message type")
	}
}

func labels(s Session) prometheus.Labels {
	return prometheus.Labels{"backend": s.BackendID(),
		"node":    s.NodeID(),
		"cluster": s.Cluster(),
	}
}
<|endoftext|>


<commit_before>// +build !jenkins

package people

import (
	"fmt"
	"os"
	"testing"

	"github.com/Financial-Times/base-ft-rw-app-go/baseftrwapp"
	"github.com/Financial-Times/neo-utils-go/neoutils"
	"github.com/jmcvetta/neoism"
	"github.com/stretchr/testify/assert"
)

var peopleDriver baseftrwapp.Service

func TestDelete(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"

	peopleDriver = getPeopleCypherDriver(t)

	personToDelete := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToDelete), "Failed to write person")

	found, err := peopleDriver.Delete(uuid)
	assert.True(found, "Didn't manage to delete person for uuid %", uuid)
	assert.NoError(err, "Error deleting person for uuid %s", uuid)

	p, found, err := peopleDriver.Read(uuid)

	assert.Equal(person{}, p, "Found person %s who should have been deleted", p)
	assert.False(found, "Found person for uuid %s who should have been deleted", uuid)
	assert.NoError(err, "Error trying to find person for uuid %s", uuid)
}

func TestCreateAllValuesPresent(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"
	peopleDriver = getPeopleCypherDriver(t)

	personToWrite := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")

	readPersonForUUIDAndCheckFieldsMatch(t, uuid, personToWrite)

	cleanUp(t, uuid)
}

func TestCreateHandlesSpecialCharacters(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"
	peopleDriver = getPeopleCypherDriver(t)

	personToWrite := person{UUID: uuid, Name: "Thomas M. O'Gara", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")

	readPersonForUUIDAndCheckFieldsMatch(t, uuid, personToWrite)

	cleanUp(t, uuid)
}

func TestCreateNotAllValuesPresent(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"
	peopleDriver = getPeopleCypherDriver(t)

	personToWrite := person{UUID: uuid, Name: "Test",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")

	readPersonForUUIDAndCheckFieldsMatch(t, uuid, personToWrite)

	cleanUp(t, uuid)
}

func TestUpdateWillRemovePropertiesNoLongerPresent(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"
	peopleDriver = getPeopleCypherDriver(t)

	personToWrite := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")
	readPersonForUUIDAndCheckFieldsMatch(t, uuid, personToWrite)

	updatedPerson := person{UUID: uuid, Name: "Test",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}}

	assert.NoError(peopleDriver.Write(updatedPerson), "Failed to write updated person")
	readPersonForUUIDAndCheckFieldsMatch(t, uuid, updatedPerson)

	cleanUp(t, uuid)
}

func TestConnectivityCheck(t *testing.T) {
	assert := assert.New(t)
	peopleDriver = getPeopleCypherDriver(t)
	err := peopleDriver.Check()
	assert.NoError(err, "Unexpected error on connectivity check")
}

func getPeopleCypherDriver(t *testing.T) service {
	assert := assert.New(t)
	url := os.Getenv("NEO4J_TEST_URL")
	if url == "" {
		url = "http://localhost:7474/db/data"
	}

	db, err := neoism.Connect(url)
	assert.NoError(err, "Failed to connect to Neo4j")
	return NewCypherPeopleService(neoutils.StringerDb{db}, db)
}

func readPersonForUUIDAndCheckFieldsMatch(t *testing.T, uuid string, expectedPerson person) {
	assert := assert.New(t)
	storedPerson, found, err := peopleDriver.Read(uuid)

	assert.NoError(err, "Error finding person for uuid %s", uuid)
	assert.True(found, "Didn't find person for uuid %s", uuid)
	assert.Equal(expectedPerson, storedPerson, "people should be the same")
}

func TestWriteAliasesAreWrittenAndAreEqualToList(t *testing.T) {
	assert := assert.New(t)
	peopleDriver := getPeopleCypherDriver(t)
	uuid := "12345"
	personToWrite := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	peopleDriver.Write(personToWrite)

	result := []struct {
		Aliases []string `json:"t.aliases"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:"12345"}) RETURN t.aliases
				`,
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal("alias 1", result[0].Aliases[0], "PrefLabel should be 'alias 1")
	cleanUp(t, uuid)
}

func TestWritePrefLabelIsAlsoWrittenAndIsEqualToName(t *testing.T) {
	assert := assert.New(t)
	peopleDriver := getPeopleCypherDriver(t)
	uuid := "12345"
	personToWrite := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}}

	storedPerson := peopleDriver.Write(personToWrite)

	fmt.Printf("", storedPerson)

	result := []struct {
		PrefLabel string `json:"t.prefLabel"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:"12345"}) RETURN t.prefLabel
				`,
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal("Test", result[0].PrefLabel, "PrefLabel should be 'Test")
	cleanUp(t, uuid)
}

func cleanUp(t *testing.T, uuid string) {
	assert := assert.New(t)
	found, err := peopleDriver.Delete(uuid)
	assert.True(found, "Didn't manage to delete person for uuid %", uuid)
	assert.NoError(err, "Error deleting person for uuid %s", uuid)
}
<commit_msg>Refactored tests<commit_after>// +build !jenkins

package people

import (
	"fmt"
	"os"
	"testing"

	"github.com/Financial-Times/base-ft-rw-app-go/baseftrwapp"
	"github.com/Financial-Times/neo-utils-go/neoutils"
	"github.com/jmcvetta/neoism"
	"github.com/stretchr/testify/assert"
)

const (
	minimalPersonUuid = "180cec41-23fa-4148-806b-0602924e6858"
	fullPersonUuid    = "bbc4f575-edb3-4f51-92f0-5ce6c708d1ea"
)

var minimalPerson = person{
	UUID:        minimalPersonUuid,
	Name:        "Minimal Person",
	Identifiers: []identifier{fsIdentifier},
}

var fullPerson = person{
	UUID:        fullPersonUuid,
	Name:        "Full Person",
	BirthYear:   1900,
	Salutation:  "Dr.",
	Identifiers: []identifier{fsIdentifier, firstTmeIdentifier},
	Aliases:     []string{"Diff Name"},
}

var fsIdentifier = identifier{
	Authority:       fsAuthority,
	IdentifierValue: "012345-E",
}

var firstTmeIdentifier = identifier{
	Authority:       tmeAuthority,
	IdentifierValue: "tmeIdentifier",
}

var secondTmeIdentifier = identifier{
	Authority:       tmeAuthority,
	IdentifierValue: "tmeIdentifier2",
}

var invalidIdentifier = identifier{
	Authority:       "Invalid Authority",
	IdentifierValue: "tmeIdentifier2",
}

var peopleDriver baseftrwapp.Service

func TestCreateAllValuesPresent(t *testing.T) {
	assert := assert.New(t)
	peopleDriver = getPeopleCypherDriver(t)

	assert.NoError(peopleDriver.Write(fullPerson), "Failed to write person")

	readPersonForUUIDAndCheckFieldsMatch(t, fullPersonUuid, fullPerson)

	cleanUp(t, fullPersonUuid)
}

func TestCreateNotAllValuesPresent(t *testing.T) {
	assert := assert.New(t)
	peopleDriver = getPeopleCypherDriver(t)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write person")

	readPersonForUUIDAndCheckFieldsMatch(t, minimalPersonUuid, minimalPerson)

	cleanUp(t, minimalPersonUuid)
}

func TestCreateHandlesSpecialCharacters(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"
	peopleDriver = getPeopleCypherDriver(t)

	personToWrite := person{UUID: uuid, Name: "Thomas M. O'Gara", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")

	readPersonForUUIDAndCheckFieldsMatch(t, uuid, personToWrite)

	cleanUp(t, uuid)
}

func TestUpdateWillRemovePropertiesNoLongerPresent(t *testing.T) {
	assert := assert.New(t)
	peopleDriver = getPeopleCypherDriver(t)

	assert.NoError(peopleDriver.Write(fullPerson), "Failed to write person")
	readPersonForUUIDAndCheckFieldsMatch(t, fullPersonUuid, fullPerson)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write updated person")
	readPersonForUUIDAndCheckFieldsMatch(t, minimalPersonUuid, minimalPerson)

	cleanUp(t, minimalPersonUuid)
}

func TestWritePersonWithUnsupportedAuthority(t *testing.T) {
	assert := assert.New(t)
	uuid := "12345"
	peopleDriver = getPeopleCypherDriver(t)

	personToWrite := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{invalidIdentifier}, Aliases: []string{"alias 1", "alias 2"}}

	assert.Error(peopleDriver.Write(personToWrite))
}

func TestAliasesAreWrittenAreAbleToBeReadAndAreEqualToList(t *testing.T) {
	assert := assert.New(t)
	peopleDriver := getPeopleCypherDriver(t)
	uuid := "12345"
	personToWrite := person{UUID: uuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	peopleDriver.Write(personToWrite)

	result := []struct {
		Aliases []string `json:"t.aliases"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:"12345"}) RETURN t.aliases
				`,
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal("alias 1", result[0].Aliases[0], "PrefLabel should be 'alias 1")
	cleanUp(t, uuid)
}

func TestPrefLabelIsEqualToNameAndAbleToBeRead(t *testing.T) {
	assert := assert.New(t)
	peopleDriver := getPeopleCypherDriver(t)

	storedPerson := peopleDriver.Write(fullPerson)

	fmt.Printf("", storedPerson)

	result := []struct {
		PrefLabel string `json:"t.prefLabel"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:{uuid}}) RETURN t.prefLabel
				`,
		Parameters: map[string]interface{}{
			"uuid": fullPersonUuid,
		},
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal(fullPerson.Name, result[0].PrefLabel, "PrefLabel should be 'Full Person")
	cleanUp(t, fullPersonUuid)
}

func TestDelete(t *testing.T) {
	assert := assert.New(t)
	peopleDriver = getPeopleCypherDriver(t)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write person")

	found, err := peopleDriver.Delete(minimalPersonUuid)
	assert.True(found, "Didn't manage to delete person for uuid %", minimalPersonUuid)
	assert.NoError(err, "Error deleting person for uuid %s", minimalPersonUuid)

	p, found, err := peopleDriver.Read(minimalPersonUuid)

	assert.Equal(person{}, p, "Found person %s who should have been deleted", p)
	assert.False(found, "Found person for uuid %s who should have been deleted", minimalPersonUuid)
	assert.NoError(err, "Error trying to find person for uuid %s", minimalPersonUuid)
}

func TestConnectivityCheck(t *testing.T) {
	assert := assert.New(t)
	peopleDriver = getPeopleCypherDriver(t)
	err := peopleDriver.Check()
	assert.NoError(err, "Unexpected error on connectivity check")
}

func getPeopleCypherDriver(t *testing.T) service {
	assert := assert.New(t)
	url := os.Getenv("NEO4J_TEST_URL")
	if url == "" {
		url = "http://localhost:7474/db/data"
	}

	db, err := neoism.Connect(url)
	assert.NoError(err, "Failed to connect to Neo4j")
	return NewCypherPeopleService(neoutils.StringerDb{db}, db)
}

func readPersonForUUIDAndCheckFieldsMatch(t *testing.T, uuid string, expectedPerson person) {
	assert := assert.New(t)
	storedPerson, found, err := peopleDriver.Read(uuid)

	assert.NoError(err, "Error finding person for uuid %s", uuid)
	assert.True(found, "Didn't find person for uuid %s", uuid)
	assert.Equal(expectedPerson, storedPerson, "people should be the same")
}

func cleanUp(t *testing.T, uuid string) {
	assert := assert.New(t)
	found, err := peopleDriver.Delete(uuid)
	assert.True(found, "Didn't manage to delete person for uuid %", uuid)
	assert.NoError(err, "Error deleting person for uuid %s", uuid)
}
<|endoftext|>


<commit_before>package datafiller

import (
	"encoding/json"
	"fmt"
	"reflect"
	"testing"
	"time"
)

// Actual tests

func TestUnassignedInt(t *testing.T) {
	var i int
	Fill(&i)
	if reflect.ValueOf(i).Kind() != reflect.Int {
		t.Errorf("Type error: %v, want %v", reflect.ValueOf(i).Kind(), reflect.Int)
	}
	fmt.Println(i)
}

func TestSimpleTypes(t *testing.T) {
	// var BoolV bool
	var IntV int
	var Int8V int8
	var Int16V int16
	var Int32V int32
	var Int64V int64
	var UintV uint
	var Uint8V uint8
	var Uint16V uint16
	var Uint32V uint32
	var Uint64V uint64
	var Float32V float32
	var Float64V float64
	var Complex64V complex64
	var Complex128V complex128

	tests := []struct {
		value        interface{}
		expectedType reflect.Kind
	}{
		// {&BoolV, reflect.Bool},
		{&IntV, reflect.Int},
		{&Int8V, reflect.Int8},
		{&Int16V, reflect.Int16},
		{&Int32V, reflect.Int32},
		{&Int64V, reflect.Int64},
		{&UintV, reflect.Uint},
		{&Uint8V, reflect.Uint8},
		{&Uint16V, reflect.Uint16},
		{&Uint32V, reflect.Uint32},
		{&Uint64V, reflect.Uint64},
		{&Float32V, reflect.Float32},
		{&Float64V, reflect.Float64},
		{&Complex64V, reflect.Complex64},
		{&Complex128V, reflect.Complex128},
	}

	for _, test := range tests {
		// Type level checking
		testValue := reflect.Indirect(reflect.ValueOf(test.value))
		if testValue.Kind() != test.expectedType {
			t.Errorf("Type error: %v, want %v", testValue.Kind(), test.expectedType)
		}

		Fill(test.value)

		testValue = reflect.Indirect(reflect.ValueOf(test.value))
		if testValue.Kind() != test.expectedType {
			t.Errorf("Type error: %v, want %v", testValue.Kind(), test.expectedType)
		}

		// Value level checking
		ifc := testValue.Interface()
		fmt.Printf("Type: %v; \t\t Value: %v \n", testValue.Kind(), ifc)
		// TODO(tvi): Figure out mock testing.

		zero := reflect.Zero(testValue.Type())
		if reflect.DeepEqual(ifc, zero.Interface()) {
			t.Errorf("Changed value is zero-value (type: %v): value %v, do not want %v", testValue.Type(), zero.Interface(), ifc)
		}
	}
}

// var ArrayV         array
// var ChanV          chan
// var FuncV          func
// var InterfaceV     interface
// var MapV           map
// var PtrV           ptr
// var SliceV         slice
// var StringV        string
// var StructV        struct

func TestComplexTypes(t *testing.T) {

}

// var InvalidV       invalid
// var UintptrV uintptr
// var UnsafePointerV unsafepointer

func TestPointerTypes(t *testing.T) {

}

func TestSimpleValues(t *testing.T) {
	// var BoolV bool
	var IntV int
	var Int8V int8
	var Int16V int16
	var Int32V int32
	var Int64V int64
	var UintV uint
	var Uint8V uint8
	var Uint16V uint16
	var Uint32V uint32
	var Uint64V uint64
	var Float32V float32
	var Float64V float64
	var Complex64V complex64
	var Complex128V complex128

	tests := []struct {
		value         interface{}
		expectedValue interface{}
	}{
		{&IntV, int(55)},
		{&Int8V, int8(55)},
		{&Int16V, int16(55)},
		{&Int32V, int32(55)},
		{&Int64V, int64(55)},
		{&UintV, uint(55)},
		{&Uint8V, uint8(55)},
		{&Uint16V, uint16(55)},
		{&Uint32V, uint32(55)},
		{&Uint64V, uint64(55)},
		{&Float32V, float32(0.91889215)},
		{&Float64V, float64(0.9188921451568604)},
		{&Complex64V, complex(0.91889215, 0.23150717)},
		{&Complex128V, complex(0.9188921451568604, 0.23150716722011566)},
	}

	for _, test := range tests {
		f := NewFiller(Seed(7))
		f.Fill(test.value)
		testValue := reflect.Indirect(reflect.ValueOf(test.value))
		ifc := testValue.Interface()
		// fmt.Printf("Type: %v; \t\t Value: %v \n", testValue.Kind(), ifc)

		if !reflect.DeepEqual(ifc, test.expectedValue) {
			t.Errorf("Value mismatch (type: %v): value %v, want %v", testValue.Type(), ifc, test.expectedValue)
		}
	}
}

// Tests for debugging

func TestDebugSimpleInt(t *testing.T) {
	i := 1
	Fill(&i)
	fmt.Println(i)
}

type S struct {
	A string
	B []struct {
		Q []struct {
			W int
		}
		C string
		D string
		E int
	}
}

func TestDebugSimpleStruct(t *testing.T) {
	i := S{}
	Fill(&i)

	b, err := json.Marshal(i)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(b))

	fmt.Println(i)
}

type A struct {
	T time.Time
	Q string
}

func TestDebugSimpleTimeStruct(t *testing.T) {
	i := A{}
	Fill(&i)

	b, err := json.Marshal(i)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(b))

	fmt.Println(i)
}

type D struct {
	Q string `json:"-"`
	A string `json:"myName,omitempty"`
	B string `datafiller:"-"`
	C string `datafiller:"name,omitempty"`
}

func TestDebugSimpleTaggedStruct(t *testing.T) {
	i := D{}
	Fill(&i)

	b, err := json.Marshal(i)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(b))

	fmt.Println(i)
}
<commit_msg>New Generator test.<commit_after>package datafiller

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"reflect"
	"testing"
	"time"
)

// Actual tests

func TestUnassignedInt(t *testing.T) {
	var i int
	Fill(&i)
	if reflect.ValueOf(i).Kind() != reflect.Int {
		t.Errorf("Type error: %v, want %v", reflect.ValueOf(i).Kind(), reflect.Int)
	}
	fmt.Println(i)
}

func TestSimpleTypes(t *testing.T) {
	// var BoolV bool
	var IntV int
	var Int8V int8
	var Int16V int16
	var Int32V int32
	var Int64V int64
	var UintV uint
	var Uint8V uint8
	var Uint16V uint16
	var Uint32V uint32
	var Uint64V uint64
	var Float32V float32
	var Float64V float64
	var Complex64V complex64
	var Complex128V complex128

	tests := []struct {
		value        interface{}
		expectedType reflect.Kind
	}{
		// TODO(tvi): Fix.
		// {&BoolV, reflect.Bool},
		{&IntV, reflect.Int},
		{&Int8V, reflect.Int8},
		{&Int16V, reflect.Int16},
		{&Int32V, reflect.Int32},
		{&Int64V, reflect.Int64},
		{&UintV, reflect.Uint},
		{&Uint8V, reflect.Uint8},
		{&Uint16V, reflect.Uint16},
		{&Uint32V, reflect.Uint32},
		{&Uint64V, reflect.Uint64},
		{&Float32V, reflect.Float32},
		{&Float64V, reflect.Float64},
		{&Complex64V, reflect.Complex64},
		{&Complex128V, reflect.Complex128},
	}

	for _, test := range tests {
		// Type level checking
		testValue := reflect.Indirect(reflect.ValueOf(test.value))
		if testValue.Kind() != test.expectedType {
			t.Errorf("Type error: %v, want %v", testValue.Kind(), test.expectedType)
		}

		Fill(test.value)

		testValue = reflect.Indirect(reflect.ValueOf(test.value))
		if testValue.Kind() != test.expectedType {
			t.Errorf("Type error: %v, want %v", testValue.Kind(), test.expectedType)
		}

		// Value level checking
		ifc := testValue.Interface()
		fmt.Printf("Type: %v; \t\t Value: %v \n", testValue.Kind(), ifc)
		// TODO(tvi): Figure out mock testing.

		zero := reflect.Zero(testValue.Type())
		if reflect.DeepEqual(ifc, zero.Interface()) {
			t.Errorf("Changed value is zero-value (type: %v): value %v, do not want %v", testValue.Type(), zero.Interface(), ifc)
		}
	}
}

// var ArrayV         array
// var ChanV          chan
// var FuncV          func
// var InterfaceV     interface
// var MapV           map
// var PtrV           ptr
// var SliceV         slice
// var StringV        string
// var StructV        struct

func TestComplexTypes(t *testing.T) {

}

// var InvalidV       invalid
// var UintptrV uintptr
// var UnsafePointerV unsafepointer

func TestPointerTypes(t *testing.T) {

}

func TestSimpleValues(t *testing.T) {
	// var BoolV bool
	var IntV int
	var Int8V int8
	var Int16V int16
	var Int32V int32
	var Int64V int64
	var UintV uint
	var Uint8V uint8
	var Uint16V uint16
	var Uint32V uint32
	var Uint64V uint64
	var Float32V float32
	var Float64V float64
	// var Complex64V complex64
	// var Complex128V complex128

	tests := []struct {
		value         interface{}
		expectedValue interface{}
	}{
		{&IntV, int(55)},
		{&Int8V, int8(55)},
		{&Int16V, int16(55)},
		{&Int32V, int32(55)},
		{&Int64V, int64(55)},
		{&UintV, uint(55)},
		{&Uint8V, uint8(55)},
		{&Uint16V, uint16(55)},
		{&Uint32V, uint32(55)},
		{&Uint64V, uint64(55)},
		{&Float32V, float32(0.91889215)},
		{&Float64V, float64(0.9188921451568604)},
		// {&Complex64V, complex(0.91889215, 0.23150717)},
		// {&Complex128V, complex(0.9188921451568604, 0.23150716722011566)},
	}

	for _, test := range tests {
		f := NewFiller(Seed(7))
		f.Fill(test.value)
		testValue := reflect.Indirect(reflect.ValueOf(test.value))
		ifc := testValue.Interface()
		fmt.Printf("Type: %v; \t\t Value: %v \n", testValue.Kind(), ifc)

		if !reflect.DeepEqual(ifc, test.expectedValue) {
			t.Errorf("Value mismatch (type: %v): value %v, want %v", testValue.Type(), ifc, test.expectedValue)
		}
	}
}

func clear(r *rand.Rand, s rand.Source) string {
	newRnd := rand.New(s)
	// rnd := reflect.Indirect(reflect.ValueOf(r))
	// rnd.Set(reflect.Indirect(reflect.ValueOf(newRnd)))

	// rnd := reflect.Indirect(reflect.ValueOf(r))
	// rnd.Set(reflect.Indirect(reflect.ValueOf(newRnd)))
	reflect.ValueOf(r).Set(reflect.ValueOf(newRnd))
	// r = rand.New(s)
	return ""
}

func TestSimpleValuesGenerate(t *testing.T) {
	var Complex64V complex64
	var Complex128V complex128

	// s := rand.NewSource(7)
	// r := rand.New(s)
	// TODO(tvi): Hack the rand num generator.

	r1 := rand.New(rand.NewSource(7))
	r2 := rand.New(rand.NewSource(7))
	tests := []struct {
		value         interface{}
		expectedValue interface{}
		// dummy         string
	}{
		// {&Complex64V, complex(r.Float32(), r.Float32()), clear(r,s)},
		// {&Complex128V, complex(r.Float64(), r.Float64()), clear(r,s)},
		{&Complex64V, complex(r1.Float32(), r1.Float32())},
		{&Complex128V, complex(r2.Float64(), r2.Float64())},
	}

	for _, test := range tests {
		f := NewFiller(Seed(7))
		f.Fill(test.value)
		testValue := reflect.Indirect(reflect.ValueOf(test.value))
		ifc := testValue.Interface()
		fmt.Printf("Type: %v; \t\t Value: %v \n", testValue.Kind(), ifc)

		if !reflect.DeepEqual(ifc, test.expectedValue) {
			t.Errorf("Value mismatch (type: %v): value %v, want %v", testValue.Type(), ifc, test.expectedValue)
		}
	}
}

// Tests for debugging

func TestDebugSimpleInt(t *testing.T) {
	i := 1
	Fill(&i)
	fmt.Println(i)
}

type S struct {
	A string
	B []struct {
		Q []struct {
			W int
		}
		C string
		D string
		E int
	}
}

func TestDebugSimpleStruct(t *testing.T) {
	i := S{}
	Fill(&i)

	b, err := json.Marshal(i)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(b))

	fmt.Println(i)
}

type A struct {
	T time.Time
	Q string
}

func TestDebugSimpleTimeStruct(t *testing.T) {
	i := A{}
	Fill(&i)

	b, err := json.Marshal(i)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(b))

	fmt.Println(i)
}

type D struct {
	Q string `json:"-"`
	A string `json:"myName,omitempty"`
	B string `datafiller:"-"`
	C string `datafiller:"name,omitempty"`
}

func TestDebugSimpleTaggedStruct(t *testing.T) {
	i := D{}
	Fill(&i)

	b, err := json.Marshal(i)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(string(b))

	fmt.Println(i)
}
<|endoftext|>


<commit_before>// +build !jenkins

package people

import (
	"fmt"
	"os"
	"testing"

	"github.com/Financial-Times/base-ft-rw-app-go/baseftrwapp"
	"github.com/Financial-Times/neo-utils-go/neoutils"
	"github.com/jmcvetta/neoism"
	"github.com/stretchr/testify/assert"
)

const (
	minimalPersonUuid = "180cec41-23fa-4148-806b-0602924e6858"
	fullPersonUuid    = "bbc4f575-edb3-4f51-92f0-5ce6c708d1ea"
	uniquePersonUuid  = "bb596d64-78c5-4b00-a88f-e8248c956073"
)

var minimalPerson = person{
	UUID:        minimalPersonUuid,
	Name:        "Minimal Person",
	Identifiers: []identifier{fsIdentifier},
}

var fullPerson = person{
	UUID:        fullPersonUuid,
	Name:        "Full Person",
	BirthYear:   1900,
	Salutation:  "Dr.",
	Identifiers: []identifier{fsIdentifier, firstTmeIdentifier, secondTmeIdentifier},
	Aliases:     []string{"Diff Name"},
}

var fsIdentifier = identifier{
	Authority:       fsAuthority,
	IdentifierValue: "012345-E",
}

var firstTmeIdentifier = identifier{
	Authority:       tmeAuthority,
	IdentifierValue: "tmeIdentifier",
}

var secondTmeIdentifier = identifier{
	Authority:       tmeAuthority,
	IdentifierValue: "tmeIdentifier2",
}

var invalidIdentifier = identifier{
	Authority:       "Invalid Authority",
	IdentifierValue: "tmeIdentifier2",
}

var peopleDriver baseftrwapp.Service

func TestCreateAllValuesPresent(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(fullPerson), "Failed to write person")

	storedPerson, _, err := peopleDriver.Read(fullPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedPerson)
}

func TestCreateNotAllValuesPresent(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write person")

	storedPerson, _, err := peopleDriver.Read(minimalPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedPerson)
}

func TestCreateHandlesSpecialCharacters(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	personToWrite := person{UUID: uniquePersonUuid, Name: "Thomas M. O'Gara", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")

	storedPerson, _, err := peopleDriver.Read(uniquePersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedPerson)
}

func TestUpdateWillRemovePropertiesNoLongerPresent(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(fullPerson), "Failed to write person")
	storedFullPerson, _, err := peopleDriver.Read(fullPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedFullPerson)

	var minimalPerson = person{
		UUID:        fullPersonUuid,
		Name:        "Minimal Person",
		Identifiers: []identifier{fsIdentifier},
	}

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write updated person")
	storedMinimalPerson, _, err := peopleDriver.Read(fullPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedMinimalPerson)
}

func TestWritePersonWithUnsupportedAuthority(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	personToWrite := person{UUID: uniquePersonUuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{invalidIdentifier}, Aliases: []string{"alias 1", "alias 2"}}

	assert.Error(peopleDriver.Write(personToWrite))
}

func TestAliasesAreWrittenAndAreAbleToBeReadInOrder(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)
	personToWrite := person{UUID: uniquePersonUuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	peopleDriver.Write(personToWrite)

	result := []struct {
		Aliases []string `json:"t.aliases"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:{uuid}}) RETURN t.aliases
				`,
		Parameters: map[string]interface{}{
			"uuid": uniquePersonUuid,
		},
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal("alias 1", result[0].Aliases[0], "PrefLabel should be 'alias 1")
}

func TestAddingPersonWithExistingIdentifiers(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	cypherDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)
	assert.NoError(cypherDriver.Write(fullPerson))
	err := cypherDriver.Write(minimalPerson)
	assert.Error(err)
	assert.IsType(&neoutils.ConstraintViolationError{}, err)
}

func TestPrefLabelIsEqualToNameAndAbleToBeRead(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	storedPerson := peopleDriver.Write(fullPerson)

	fmt.Printf("", storedPerson)

	result := []struct {
		PrefLabel string `json:"t.prefLabel"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:{uuid}}) RETURN t.prefLabel
				`,
		Parameters: map[string]interface{}{
			"uuid": fullPersonUuid,
		},
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal(fullPerson.Name, result[0].PrefLabel, "PrefLabel should be 'Full Person")
}

func TestDelete(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write person")

	found, err := peopleDriver.Delete(minimalPersonUuid)
	assert.True(found, "Didn't manage to delete person for uuid %", minimalPersonUuid)
	assert.NoError(err, "Error deleting person for uuid %s", minimalPersonUuid)

	p, found, err := peopleDriver.Read(minimalPersonUuid)

	assert.Equal(person{}, p, "Found person %s who should have been deleted", p)
	assert.False(found, "Found person for uuid %s who should have been deleted", minimalPersonUuid)
	assert.NoError(err, "Error trying to find person for uuid %s", minimalPersonUuid)
}

func getDatabaseConnectionAndCheckClean(t *testing.T, assert *assert.Assertions) *neoism.Database {
	db := getDatabaseConnection(assert)
	cleanDB(db, t, assert)
	checkDbClean(db, t)
	return db
}

func getDatabaseConnection(assert *assert.Assertions) *neoism.Database {
	url := os.Getenv("NEO4J_TEST_URL")
	if url == "" {
		url = "http://localhost:7474/db/data"
	}

	db, err := neoism.Connect(url)
	assert.NoError(err, "Failed to connect to Neo4j")
	return db
}

func cleanDB(db *neoism.Database, t *testing.T, assert *assert.Assertions) {
	qs := []*neoism.CypherQuery{
		{
			Statement: fmt.Sprintf("MATCH (mp:Thing {uuid: '%v'})<-[:IDENTIFIES*0..]-(i:Identifier) DETACH DELETE mp, i", minimalPersonUuid),
		},
		{
			Statement: fmt.Sprintf("MATCH (fp:Thing {uuid: '%v'})<-[:IDENTIFIES*0..]-(i:Identifier) DETACH DELETE fp, i", fullPersonUuid),
		},
	}

	err := db.CypherBatch(qs)
	assert.NoError(err)
}

func checkDbClean(db *neoism.Database, t *testing.T) {
	assert := assert.New(t)

	result := []struct {
		Uuid string `json:"org.uuid"`
	}{}

	checkGraph := neoism.CypherQuery{
		Statement: `
			MATCH (org:Thing) WHERE org.uuid in {uuids} RETURN org.uuid
		`,
		Parameters: neoism.Props{
			"uuids": []string{fullPersonUuid, minimalPersonUuid},
		},
		Result: &result,
	}
	err := db.Cypher(&checkGraph)
	assert.NoError(err)
	assert.Empty(result)
}

func getCypherDriver(db *neoism.Database) service {
	cr := NewCypherPeopleService(neoutils.NewBatchCypherRunner(neoutils.StringerDb{db}, 3), db)
	cr.Initialise()
	return cr
}
<commit_msg>Changed name of new test<commit_after>// +build !jenkins

package people

import (
	"fmt"
	"os"
	"testing"

	"github.com/Financial-Times/base-ft-rw-app-go/baseftrwapp"
	"github.com/Financial-Times/neo-utils-go/neoutils"
	"github.com/jmcvetta/neoism"
	"github.com/stretchr/testify/assert"
)

const (
	minimalPersonUuid = "180cec41-23fa-4148-806b-0602924e6858"
	fullPersonUuid    = "bbc4f575-edb3-4f51-92f0-5ce6c708d1ea"
	uniquePersonUuid  = "bb596d64-78c5-4b00-a88f-e8248c956073"
)

var minimalPerson = person{
	UUID:        minimalPersonUuid,
	Name:        "Minimal Person",
	Identifiers: []identifier{fsIdentifier},
}

var fullPerson = person{
	UUID:        fullPersonUuid,
	Name:        "Full Person",
	BirthYear:   1900,
	Salutation:  "Dr.",
	Identifiers: []identifier{fsIdentifier, firstTmeIdentifier, secondTmeIdentifier},
	Aliases:     []string{"Diff Name"},
}

var fsIdentifier = identifier{
	Authority:       fsAuthority,
	IdentifierValue: "012345-E",
}

var firstTmeIdentifier = identifier{
	Authority:       tmeAuthority,
	IdentifierValue: "tmeIdentifier",
}

var secondTmeIdentifier = identifier{
	Authority:       tmeAuthority,
	IdentifierValue: "tmeIdentifier2",
}

var invalidIdentifier = identifier{
	Authority:       "Invalid Authority",
	IdentifierValue: "tmeIdentifier2",
}

var peopleDriver baseftrwapp.Service

func TestCreateAllValuesPresent(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(fullPerson), "Failed to write person")

	storedPerson, _, err := peopleDriver.Read(fullPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedPerson)
}

func TestCreateNotAllValuesPresent(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write person")

	storedPerson, _, err := peopleDriver.Read(minimalPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedPerson)
}

func TestCreateHandlesSpecialCharacters(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	personToWrite := person{UUID: uniquePersonUuid, Name: "Thomas M. O'Gara", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	assert.NoError(peopleDriver.Write(personToWrite), "Failed to write person")

	storedPerson, _, err := peopleDriver.Read(uniquePersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedPerson)
}

func TestUpdateWillRemovePropertiesNoLongerPresent(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(fullPerson), "Failed to write person")
	storedFullPerson, _, err := peopleDriver.Read(fullPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedFullPerson)

	var minimalPerson = person{
		UUID:        fullPersonUuid,
		Name:        "Minimal Person",
		Identifiers: []identifier{fsIdentifier},
	}

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write updated person")
	storedMinimalPerson, _, err := peopleDriver.Read(fullPersonUuid)

	assert.NoError(err)
	assert.NotEmpty(storedMinimalPerson)
}

func TestWritePersonWithUnsupportedAuthority(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	personToWrite := person{UUID: uniquePersonUuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{invalidIdentifier}, Aliases: []string{"alias 1", "alias 2"}}

	assert.Error(peopleDriver.Write(personToWrite))
}

func TestAliasesAreWrittenAndAreAbleToBeReadInOrder(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)
	personToWrite := person{UUID: uniquePersonUuid, Name: "Test", BirthYear: 1974, Salutation: "Mr",
		Identifiers: []identifier{identifier{fsAuthority, "FACTSET_ID"}}, Aliases: []string{"alias 1", "alias 2"}}

	peopleDriver.Write(personToWrite)

	result := []struct {
		Aliases []string `json:"t.aliases"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:{uuid}}) RETURN t.aliases
				`,
		Parameters: map[string]interface{}{
			"uuid": uniquePersonUuid,
		},
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal("alias 1", result[0].Aliases[0], "PrefLabel should be 'alias 1")
}

func TestAddingPersonWithExistingIdentifiersShouldFail(t *testing.T) {
	assert := assert.New(t)

	db := getDatabaseConnectionAndCheckClean(t, assert)
	cypherDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)
	assert.NoError(cypherDriver.Write(fullPerson))
	err := cypherDriver.Write(minimalPerson)
	assert.Error(err)
	assert.IsType(&neoutils.ConstraintViolationError{}, err)
}

func TestPrefLabelIsEqualToNameAndAbleToBeRead(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	storedPerson := peopleDriver.Write(fullPerson)

	fmt.Printf("", storedPerson)

	result := []struct {
		PrefLabel string `json:"t.prefLabel"`
	}{}

	getPrefLabelQuery := &neoism.CypherQuery{
		Statement: `
				MATCH (t:Person {uuid:{uuid}}) RETURN t.prefLabel
				`,
		Parameters: map[string]interface{}{
			"uuid": fullPersonUuid,
		},
		Result: &result,
	}

	err := peopleDriver.cypherRunner.CypherBatch([]*neoism.CypherQuery{getPrefLabelQuery})
	assert.NoError(err)
	assert.Equal(fullPerson.Name, result[0].PrefLabel, "PrefLabel should be 'Full Person")
}

func TestDelete(t *testing.T) {
	assert := assert.New(t)
	db := getDatabaseConnectionAndCheckClean(t, assert)
	peopleDriver := getCypherDriver(db)
	defer cleanDB(db, t, assert)

	assert.NoError(peopleDriver.Write(minimalPerson), "Failed to write person")

	found, err := peopleDriver.Delete(minimalPersonUuid)
	assert.True(found, "Didn't manage to delete person for uuid %", minimalPersonUuid)
	assert.NoError(err, "Error deleting person for uuid %s", minimalPersonUuid)

	p, found, err := peopleDriver.Read(minimalPersonUuid)

	assert.Equal(person{}, p, "Found person %s who should have been deleted", p)
	assert.False(found, "Found person for uuid %s who should have been deleted", minimalPersonUuid)
	assert.NoError(err, "Error trying to find person for uuid %s", minimalPersonUuid)
}

func getDatabaseConnectionAndCheckClean(t *testing.T, assert *assert.Assertions) *neoism.Database {
	db := getDatabaseConnection(assert)
	cleanDB(db, t, assert)
	checkDbClean(db, t)
	return db
}

func getDatabaseConnection(assert *assert.Assertions) *neoism.Database {
	url := os.Getenv("NEO4J_TEST_URL")
	if url == "" {
		url = "http://localhost:7474/db/data"
	}

	db, err := neoism.Connect(url)
	assert.NoError(err, "Failed to connect to Neo4j")
	return db
}

func cleanDB(db *neoism.Database, t *testing.T, assert *assert.Assertions) {
	qs := []*neoism.CypherQuery{
		{
			Statement: fmt.Sprintf("MATCH (mp:Thing {uuid: '%v'})<-[:IDENTIFIES*0..]-(i:Identifier) DETACH DELETE mp, i", minimalPersonUuid),
		},
		{
			Statement: fmt.Sprintf("MATCH (fp:Thing {uuid: '%v'})<-[:IDENTIFIES*0..]-(i:Identifier) DETACH DELETE fp, i", fullPersonUuid),
		},
	}

	err := db.CypherBatch(qs)
	assert.NoError(err)
}

func checkDbClean(db *neoism.Database, t *testing.T) {
	assert := assert.New(t)

	result := []struct {
		Uuid string `json:"org.uuid"`
	}{}

	checkGraph := neoism.CypherQuery{
		Statement: `
			MATCH (org:Thing) WHERE org.uuid in {uuids} RETURN org.uuid
		`,
		Parameters: neoism.Props{
			"uuids": []string{fullPersonUuid, minimalPersonUuid},
		},
		Result: &result,
	}
	err := db.Cypher(&checkGraph)
	assert.NoError(err)
	assert.Empty(result)
}

func getCypherDriver(db *neoism.Database) service {
	cr := NewCypherPeopleService(neoutils.NewBatchCypherRunner(neoutils.StringerDb{db}, 3), db)
	cr.Initialise()
	return cr
}
<|endoftext|>


<commit_before>// Copyright 2014, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

package xz

import "crypto/sha256"
import "encoding/binary"
import "bitbucket.org/rawr/golib/errs"
import "bitbucket.org/rawr/goxz/lib"

// The logic here deals with the sizing and creation of literal blocks for the
// LZMA2 filter. If the size of the compressed data ends up being larger than
// the maximum size needed for a literally "compressed" block, then it makes
// sense to just encode the data ad-verbatim. This places an upper limit on the
// worst-case compression size and also greatly improves decompression speed
// since literal blocks are almost as fast as a memcopy.

const literalChunkSize = 1 << 16
const literalChunkMarker = 1
const literalEndMarker = 0

var literalFilters = func() *lib.Filters {
	opts, _ := lib.NewOptionsLZMA(lib.PRESET_DEFAULT)
	filters := lib.NewFilters()
	filters.Append(lib.FILTER_LZMA2, opts)
	return filters
}()

// Make sure we support the requested check type.
func literalCheckSupport(check int) {
	switch check {
	case CheckNone, CheckCRC32, CheckCRC64, CheckSHA256:
	default:
		panic(optionsError)
	}
}

// Compute the compressed data size for a literal block. This is only the data
// portion and does not contain the block headers, padding, or check.
func literalCompSize(uncompSize int64) int64 {
	// Implementation based on pixz/writer.c
	numChunks := (uncompSize + literalChunkSize - 1) / literalChunkSize
	compSize := uncompSize + int64(numChunks*3) + 1
	return compSize
}

// Compute the total size for a literal block. This includes the block headers,
// data, padding, and check.
func literalBlockSize(check int, uncompSize int64) int64 {
	// Implementation based on pixz/writer.c
	literalCheckSupport(check)
	compSize := literalCompSize(uncompSize)

	block := lib.NewBlockCustom(check, literalFilters, compSize, uncompSize)
	headerSize, err := block.HeaderSize()
	errs.Panic(err)
	padSize := (4 - compSize%4) % 4
	checkSize := lib.CheckSize(check)
	return int64(headerSize) + compSize + padSize + int64(checkSize)
}

// Compresses the input as literal chunks into the output slice. The output
// slice must be large enough, as calculated with literalBlockSize().
func literalBlockEncode(check int, input []byte, output []byte) (cnt, pad int) {
	// Implementation based on pixz/writer.c
	literalCheckSupport(check)
	uncompSize := int64(len(input))
	compSize := literalCompSize(uncompSize)

	// Write the header
	block := lib.NewBlockCustom(check, literalFilters, compSize, uncompSize)
	headerSize, err := block.HeaderSize()
	errs.Panic(err)
	errs.Panic(block.HeaderEncode(output))
	cnt += headerSize

	// Write the literal chunks
	for inCnt := 0; inCnt < len(input); inCnt += literalChunkSize {
		chunk := input[inCnt:]
		if len(chunk) > literalChunkSize {
			chunk = chunk[:literalChunkSize]
		}
		output[cnt] = literalChunkMarker
		cnt++

		// Chunk size
		size := uint16(len(chunk) - 1)
		binary.BigEndian.PutUint16(output[cnt:], size)
		cnt += 2

		// Copy literal data
		cnt += copy(output[cnt:], chunk)
	}
	output[cnt] = literalEndMarker
	cnt++

	// Write the padding
	for ; cnt%4 > 0; cnt++ {
		output[cnt] = 0
		pad++
	}

	// Write the check
	switch check {
	case CheckNone:
		// Do nothing
	case CheckCRC32:
		crc := lib.CRC32(input, 0)
		binary.LittleEndian.PutUint32(output[cnt:], crc)
	case CheckCRC64:
		crc := lib.CRC64(input, 0)
		binary.LittleEndian.PutUint64(output[cnt:], crc)
	case CheckSHA256:
		sha := sha256.Sum256(input)
		copy(output[cnt:], sha[:])
	}
	cnt += lib.CheckSize(check)

	return cnt, pad
}
<commit_msg>Comment style change<commit_after>// Copyright 2014, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

package xz

import "crypto/sha256"
import "encoding/binary"
import "bitbucket.org/rawr/golib/errs"
import "bitbucket.org/rawr/goxz/lib"

// The logic here deals with the sizing and creation of literal blocks for the
// LZMA2 filter. If the size of the compressed data ends up being larger than
// the maximum size needed for a literally "compressed" block, then it makes
// sense to just encode the data ad-verbatim. This places an upper limit on the
// worst-case compression size and also greatly improves decompression speed
// since literal blocks are almost as fast as a memcopy.

const literalChunkSize = 1 << 16
const literalChunkMarker = 1
const literalEndMarker = 0

var literalFilters = func() *lib.Filters {
	opts, _ := lib.NewOptionsLZMA(lib.PRESET_DEFAULT)
	filters := lib.NewFilters()
	filters.Append(lib.FILTER_LZMA2, opts)
	return filters
}()

// Make sure we support the requested check type.
func literalCheckSupport(check int) {
	switch check {
	case CheckNone, CheckCRC32, CheckCRC64, CheckSHA256:
	default:
		panic(optionsError)
	}
}

// Compute the compressed data size for a literal block. This is only the data
// portion and does not contain the block headers, padding, or check.
func literalCompSize(uncompSize int64) int64 {
	// Implementation based on pixz/writer.c
	numChunks := (uncompSize + literalChunkSize - 1) / literalChunkSize
	compSize := uncompSize + int64(numChunks*3) + 1
	return compSize
}

// Compute the total size for a literal block. This includes the block headers,
// data, padding, and check.
func literalBlockSize(check int, uncompSize int64) int64 {
	// Implementation based on pixz/writer.c
	literalCheckSupport(check)
	compSize := literalCompSize(uncompSize)

	block := lib.NewBlockCustom(check, literalFilters, compSize, uncompSize)
	headerSize, err := block.HeaderSize()
	errs.Panic(err)
	padSize := (4 - compSize%4) % 4
	checkSize := lib.CheckSize(check)
	return int64(headerSize) + compSize + padSize + int64(checkSize)
}

// Compresses the input as literal chunks into the output slice. The output
// slice must be large enough, as calculated with literalBlockSize.
func literalBlockEncode(check int, input []byte, output []byte) (cnt, pad int) {
	// Implementation based on pixz/writer.c
	literalCheckSupport(check)
	uncompSize := int64(len(input))
	compSize := literalCompSize(uncompSize)

	// Write the header
	block := lib.NewBlockCustom(check, literalFilters, compSize, uncompSize)
	headerSize, err := block.HeaderSize()
	errs.Panic(err)
	errs.Panic(block.HeaderEncode(output))
	cnt += headerSize

	// Write the literal chunks
	for inCnt := 0; inCnt < len(input); inCnt += literalChunkSize {
		chunk := input[inCnt:]
		if len(chunk) > literalChunkSize {
			chunk = chunk[:literalChunkSize]
		}
		output[cnt] = literalChunkMarker
		cnt++

		// Chunk size
		size := uint16(len(chunk) - 1)
		binary.BigEndian.PutUint16(output[cnt:], size)
		cnt += 2

		// Copy literal data
		cnt += copy(output[cnt:], chunk)
	}
	output[cnt] = literalEndMarker
	cnt++

	// Write the padding
	for ; cnt%4 > 0; cnt++ {
		output[cnt] = 0
		pad++
	}

	// Write the check
	switch check {
	case CheckNone:
		// Do nothing
	case CheckCRC32:
		crc := lib.CRC32(input, 0)
		binary.LittleEndian.PutUint32(output[cnt:], crc)
	case CheckCRC64:
		crc := lib.CRC64(input, 0)
		binary.LittleEndian.PutUint64(output[cnt:], crc)
	case CheckSHA256:
		sha := sha256.Sum256(input)
		copy(output[cnt:], sha[:])
	}
	cnt += lib.CheckSize(check)

	return cnt, pad
}
<|endoftext|>