<commit_before>package container_disk

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"time"

	"kubevirt.io/client-go/log"
	containerdisk "kubevirt.io/kubevirt/pkg/container-disk"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/virt-handler/isolation"

	"k8s.io/apimachinery/pkg/types"

	v1 "kubevirt.io/client-go/api/v1"
)

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

type mounter struct {
	podIsolationDetector   isolation.PodIsolationDetector
	mountStateDir          string
	mountRecords           map[types.UID]*vmiMountTargetRecord
	mountRecordsLock       sync.Mutex
	suppressWarningTimeout time.Duration
	pathGetter             containerdisk.SocketPathGetter
}

type Mounter interface {
	ContainerDisksReady(vmi *v1.VirtualMachineInstance, notInitializedSince time.Time) (bool, error)
	Mount(vmi *v1.VirtualMachineInstance, verify bool) error
	Unmount(vmi *v1.VirtualMachineInstance) error
}

type vmiMountTargetEntry struct {
	TargetFile string `json:"targetFile"`
	SocketFile string `json:"socketFile"`
}

type vmiMountTargetRecord struct {
	MountTargetEntries []vmiMountTargetEntry `json:"mountTargetEntries"`
}

func NewMounter(isoDetector isolation.PodIsolationDetector, mountStateDir string) Mounter {
	return &mounter{
		mountRecords:           make(map[types.UID]*vmiMountTargetRecord),
		podIsolationDetector:   isoDetector,
		mountStateDir:          mountStateDir,
		suppressWarningTimeout: 1 * time.Minute,
		pathGetter:             containerdisk.NewSocketPathGetter(""),
	}
}

func (m *mounter) deleteMountTargetRecord(vmi *v1.VirtualMachineInstance) error {
	if string(vmi.UID) == "" {
		return fmt.Errorf("unable to find container disk mounted directories for vmi without uid")
	}

	recordFile := filepath.Join(m.mountStateDir, string(vmi.UID))

	exists, err := diskutils.FileExists(recordFile)
	if err != nil {
		return err
	}

	if exists {
		record, err := m.getMountTargetRecord(vmi)
		if err != nil {
			return err
		}

		for _, target := range record.MountTargetEntries {
			os.Remove(target.TargetFile)
			os.Remove(target.SocketFile)
		}

		os.Remove(recordFile)
	}

	m.mountRecordsLock.Lock()
	defer m.mountRecordsLock.Unlock()
	delete(m.mountRecords, vmi.UID)

	return nil
}

func (m *mounter) getMountTargetRecord(vmi *v1.VirtualMachineInstance) (*vmiMountTargetRecord, error) {
	var ok bool
	var existingRecord *vmiMountTargetRecord

	if string(vmi.UID) == "" {
		return nil, fmt.Errorf("unable to find container disk mounted directories for vmi without uid")
	}

	m.mountRecordsLock.Lock()
	defer m.mountRecordsLock.Unlock()
	existingRecord, ok = m.mountRecords[vmi.UID]

	// first check memory cache
	if ok {
		return existingRecord, nil
	}

	// if not there, see if record is on disk, this can happen if virt-handler restarts
	recordFile := filepath.Join(m.mountStateDir, filepath.Clean(string(vmi.UID)))

	exists, err := diskutils.FileExists(recordFile)
	if err != nil {
		return nil, err
	}

	if exists {
		record := vmiMountTargetRecord{}
		// #nosec No risk for path injection. Using static base and cleaned filename
		bytes, err := ioutil.ReadFile(recordFile)
		if err != nil {
			return nil, err
		}
		err = json.Unmarshal(bytes, &record)
		if err != nil {
			return nil, err
		}

		m.mountRecords[vmi.UID] = &record
		return &record, nil
	}

	// not found
	return nil, nil
}

func (m *mounter) setMountTargetRecord(vmi *v1.VirtualMachineInstance, record *vmiMountTargetRecord) error {
	if string(vmi.UID) == "" {
		return fmt.Errorf("unable to set container disk mounted directories for vmi without uid")
	}

	recordFile := filepath.Join(m.mountStateDir, string(vmi.UID))
	fileExists, err := diskutils.FileExists(recordFile)
	if err != nil {
		return err
	}

	m.mountRecordsLock.Lock()
	defer m.mountRecordsLock.Unlock()

	existingRecord, ok := m.mountRecords[vmi.UID]
	if ok && fileExists && reflect.DeepEqual(existingRecord, record) {
		// already done
		return nil
	}

	bytes, err := json.Marshal(record)
	if err != nil {
		return err
	}

	err = os.MkdirAll(filepath.Dir(recordFile), 0755)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(recordFile, bytes, 0644)
	if err != nil {
		return err
	}

	m.mountRecords[vmi.UID] = record

	return nil
}

// Mount takes a vmi and mounts all container disks of the VMI, so that they are visible for the qemu process.
// Additionally qcow2 images are validated if "verify" is true. The validation happens with rlimits set, to avoid DOS.
func (m *mounter) Mount(vmi *v1.VirtualMachineInstance, verify bool) error {
	record := vmiMountTargetRecord{}

	for i, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {
			targetFile, err := containerdisk.GetDiskTargetPathFromHostView(vmi, i)
			if err != nil {
				return err
			}

			sock, err := m.pathGetter(vmi, i)
			if err != nil {
				return err
			}

			record.MountTargetEntries = append(record.MountTargetEntries, vmiMountTargetEntry{
				TargetFile: targetFile,
				SocketFile: sock,
			})
		}
	}

	if len(record.MountTargetEntries) > 0 {
		err := m.setMountTargetRecord(vmi, &record)
		if err != nil {
			return err
		}
	}

	for i, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {
			targetFile, err := containerdisk.GetDiskTargetPathFromHostView(vmi, i)
			if err != nil {
				return err
			}

			nodeRes := isolation.NodeIsolationResult()

			if isMounted, err := nodeRes.IsMounted(targetFile); err != nil {
				return fmt.Errorf("failed to determine if %s is already mounted: %v", targetFile, err)
			} else if !isMounted {
				sock, err := m.pathGetter(vmi, i)
				if err != nil {
					return err
				}

				res, err := m.podIsolationDetector.DetectForSocket(vmi, sock)
				if err != nil {
					return fmt.Errorf("failed to detect socket for containerDisk %v: %v", volume.Name, err)
				}
				mountInfo, err := res.MountInfoRoot()
				if err != nil {
					return fmt.Errorf("failed to detect root mount info of containerDisk  %v: %v", volume.Name, err)
				}
				nodeMountInfo, err := nodeRes.ParentMountInfoFor(mountInfo)
				if err != nil {
					return fmt.Errorf("failed to detect root mount point of containerDisk %v on the node: %v", volume.Name, err)
				}
				sourceFile, err := containerdisk.GetImage(filepath.Join(nodeRes.MountRoot(), nodeMountInfo.Root, nodeMountInfo.MountPoint), volume.ContainerDisk.Path)
				if err != nil {
					return fmt.Errorf("failed to find a sourceFile in containerDisk %v: %v", volume.Name, err)
				}
				f, err := os.Create(targetFile)
				if err != nil {
					return fmt.Errorf("failed to create mount point target %v: %v", targetFile, err)
				}
				f.Close()

				log.DefaultLogger().Object(vmi).Infof("Bind mounting container disk at %s to %s", strings.TrimPrefix(sourceFile, nodeRes.MountRoot()), targetFile)
				// #nosec g204 no risk to TrimPref as argument as it just trims two fixed strings
				out, err := exec.Command("/usr/bin/virt-chroot", "--mount", "/proc/1/ns/mnt", "mount", "-o", "ro,bind", strings.TrimPrefix(sourceFile, nodeRes.MountRoot()), targetFile).CombinedOutput()
				if err != nil {
					return fmt.Errorf("failed to bindmount containerDisk %v: %v : %v", volume.Name, string(out), err)
				}
			}
			if verify {
				res, err := m.podIsolationDetector.Detect(vmi)
				if err != nil {
					return fmt.Errorf("failed to detect VMI pod: %v", err)
				}
				imageInfo, err := isolation.GetImageInfo(containerdisk.GetDiskTargetPathFromLauncherView(i), res)
				if err != nil {
					return fmt.Errorf("failed to get image info: %v", err)
				}

				if err := containerdisk.VerifyImage(imageInfo); err != nil {
					return fmt.Errorf("invalid image in containerDisk %v: %v", volume.Name, err)
				}
			}
		}
	}
	return nil
}

// Legacy Unmount unmounts all container disks of a given VMI when the hold HostPath method was in use.
// This exists for backwards compatibility for VMIs running before a KubeVirt update occurs.
func (m *mounter) legacyUnmount(vmi *v1.VirtualMachineInstance) error {
	mountDir := containerdisk.GetLegacyVolumeMountDirOnHost(vmi)

	files, err := ioutil.ReadDir(mountDir)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to list container disk mounts: %v", err)
	}

	if vmi.UID != "" {
		for _, file := range files {
			path := filepath.Join(mountDir, file.Name())
			if strings.HasSuffix(path, ".sock") {
				continue
			}
			if mounted, err := isolation.NodeIsolationResult().IsMounted(path); err != nil {
				return fmt.Errorf("failed to check mount point for containerDisk %v: %v", path, err)
			} else if mounted {
				// #nosec No risk for attacket injection. Parameters are predefined strings
				out, err := exec.Command("/usr/bin/virt-chroot", "--mount", "/proc/1/ns/mnt", "umount", path).CombinedOutput()
				if err != nil {
					return fmt.Errorf("failed to unmount containerDisk %v: %v : %v", path, string(out), err)
				}
			}
		}

		if err := os.RemoveAll(mountDir); err != nil {
			return fmt.Errorf("failed to remove containerDisk files: %v", err)
		}
	}
	return nil
}

// Unmount unmounts all container disks of a given VMI.
func (m *mounter) Unmount(vmi *v1.VirtualMachineInstance) error {
	if vmi.UID != "" {

		// this will catch unmounting a vmi's container disk when
		// an old VMI is left over after a KubeVirt update
		err := m.legacyUnmount(vmi)
		if err != nil {
			return err
		}

		record, err := m.getMountTargetRecord(vmi)
		if err != nil {
			return err
		} else if record == nil {
			// no entries to unmount

			log.DefaultLogger().Object(vmi).Infof("No container disk mount entries found to unmount")
			return nil
		}

		log.DefaultLogger().Object(vmi).Infof("Found container disk mount entries")
		for _, entry := range record.MountTargetEntries {
			path := entry.TargetFile
			log.DefaultLogger().Object(vmi).Infof("Looking to see if containerdisk is mounted at path %s", path)
			if mounted, err := isolation.NodeIsolationResult().IsMounted(path); err != nil {
				return fmt.Errorf("failed to check mount point for containerDisk %v: %v", path, err)
			} else if mounted {
				log.DefaultLogger().Object(vmi).Infof("unmounting container disk at path %s", path)
				// #nosec No risk for attacket injection. Parameters are predefined strings
				out, err := exec.Command("/usr/bin/virt-chroot", "--mount", "/proc/1/ns/mnt", "umount", path).CombinedOutput()
				if err != nil {
					return fmt.Errorf("failed to unmount containerDisk %v: %v : %v", path, string(out), err)
				}
			}

		}
		err = m.deleteMountTargetRecord(vmi)
		if err != nil {
			return err
		}
	}
	return nil
}

func (m *mounter) ContainerDisksReady(vmi *v1.VirtualMachineInstance, notInitializedSince time.Time) (bool, error) {
	for i, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {
			_, err := m.pathGetter(vmi, i)
			if err != nil {
				log.DefaultLogger().Object(vmi).Infof("containerdisk %s not yet ready", volume.Name)
				if time.Now().After(notInitializedSince.Add(m.suppressWarningTimeout)) {
					return false, fmt.Errorf("containerdisk %s still not ready after one minute", volume.Name)
				}
				return false, nil
			}
		}
	}
	log.DefaultLogger().Object(vmi).V(4).Info("all containerdisks are ready")
	return true, nil
}
<commit_msg>Append rootfs mount to containerdisk base path<commit_after>package container_disk

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"time"

	"kubevirt.io/client-go/log"
	containerdisk "kubevirt.io/kubevirt/pkg/container-disk"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/virt-handler/isolation"

	"k8s.io/apimachinery/pkg/types"

	v1 "kubevirt.io/client-go/api/v1"
)

//go:generate mockgen -source $GOFILE -package=$GOPACKAGE -destination=generated_mock_$GOFILE

type mounter struct {
	podIsolationDetector   isolation.PodIsolationDetector
	mountStateDir          string
	mountRecords           map[types.UID]*vmiMountTargetRecord
	mountRecordsLock       sync.Mutex
	suppressWarningTimeout time.Duration
	pathGetter             containerdisk.SocketPathGetter
}

type Mounter interface {
	ContainerDisksReady(vmi *v1.VirtualMachineInstance, notInitializedSince time.Time) (bool, error)
	Mount(vmi *v1.VirtualMachineInstance, verify bool) error
	Unmount(vmi *v1.VirtualMachineInstance) error
}

type vmiMountTargetEntry struct {
	TargetFile string `json:"targetFile"`
	SocketFile string `json:"socketFile"`
}

type vmiMountTargetRecord struct {
	MountTargetEntries []vmiMountTargetEntry `json:"mountTargetEntries"`
}

func NewMounter(isoDetector isolation.PodIsolationDetector, mountStateDir string) Mounter {
	return &mounter{
		mountRecords:           make(map[types.UID]*vmiMountTargetRecord),
		podIsolationDetector:   isoDetector,
		mountStateDir:          mountStateDir,
		suppressWarningTimeout: 1 * time.Minute,
		pathGetter:             containerdisk.NewSocketPathGetter(""),
	}
}

func (m *mounter) deleteMountTargetRecord(vmi *v1.VirtualMachineInstance) error {
	if string(vmi.UID) == "" {
		return fmt.Errorf("unable to find container disk mounted directories for vmi without uid")
	}

	recordFile := filepath.Join(m.mountStateDir, string(vmi.UID))

	exists, err := diskutils.FileExists(recordFile)
	if err != nil {
		return err
	}

	if exists {
		record, err := m.getMountTargetRecord(vmi)
		if err != nil {
			return err
		}

		for _, target := range record.MountTargetEntries {
			os.Remove(target.TargetFile)
			os.Remove(target.SocketFile)
		}

		os.Remove(recordFile)
	}

	m.mountRecordsLock.Lock()
	defer m.mountRecordsLock.Unlock()
	delete(m.mountRecords, vmi.UID)

	return nil
}

func (m *mounter) getMountTargetRecord(vmi *v1.VirtualMachineInstance) (*vmiMountTargetRecord, error) {
	var ok bool
	var existingRecord *vmiMountTargetRecord

	if string(vmi.UID) == "" {
		return nil, fmt.Errorf("unable to find container disk mounted directories for vmi without uid")
	}

	m.mountRecordsLock.Lock()
	defer m.mountRecordsLock.Unlock()
	existingRecord, ok = m.mountRecords[vmi.UID]

	// first check memory cache
	if ok {
		return existingRecord, nil
	}

	// if not there, see if record is on disk, this can happen if virt-handler restarts
	recordFile := filepath.Join(m.mountStateDir, filepath.Clean(string(vmi.UID)))

	exists, err := diskutils.FileExists(recordFile)
	if err != nil {
		return nil, err
	}

	if exists {
		record := vmiMountTargetRecord{}
		// #nosec No risk for path injection. Using static base and cleaned filename
		bytes, err := ioutil.ReadFile(recordFile)
		if err != nil {
			return nil, err
		}
		err = json.Unmarshal(bytes, &record)
		if err != nil {
			return nil, err
		}

		m.mountRecords[vmi.UID] = &record
		return &record, nil
	}

	// not found
	return nil, nil
}

func (m *mounter) setMountTargetRecord(vmi *v1.VirtualMachineInstance, record *vmiMountTargetRecord) error {
	if string(vmi.UID) == "" {
		return fmt.Errorf("unable to set container disk mounted directories for vmi without uid")
	}

	recordFile := filepath.Join(m.mountStateDir, string(vmi.UID))
	fileExists, err := diskutils.FileExists(recordFile)
	if err != nil {
		return err
	}

	m.mountRecordsLock.Lock()
	defer m.mountRecordsLock.Unlock()

	existingRecord, ok := m.mountRecords[vmi.UID]
	if ok && fileExists && reflect.DeepEqual(existingRecord, record) {
		// already done
		return nil
	}

	bytes, err := json.Marshal(record)
	if err != nil {
		return err
	}

	err = os.MkdirAll(filepath.Dir(recordFile), 0755)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(recordFile, bytes, 0644)
	if err != nil {
		return err
	}

	m.mountRecords[vmi.UID] = record

	return nil
}

// Mount takes a vmi and mounts all container disks of the VMI, so that they are visible for the qemu process.
// Additionally qcow2 images are validated if "verify" is true. The validation happens with rlimits set, to avoid DOS.
func (m *mounter) Mount(vmi *v1.VirtualMachineInstance, verify bool) error {
	record := vmiMountTargetRecord{}

	for i, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {
			targetFile, err := containerdisk.GetDiskTargetPathFromHostView(vmi, i)
			if err != nil {
				return err
			}

			sock, err := m.pathGetter(vmi, i)
			if err != nil {
				return err
			}

			record.MountTargetEntries = append(record.MountTargetEntries, vmiMountTargetEntry{
				TargetFile: targetFile,
				SocketFile: sock,
			})
		}
	}

	if len(record.MountTargetEntries) > 0 {
		err := m.setMountTargetRecord(vmi, &record)
		if err != nil {
			return err
		}
	}

	for i, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {
			targetFile, err := containerdisk.GetDiskTargetPathFromHostView(vmi, i)
			if err != nil {
				return err
			}

			nodeRes := isolation.NodeIsolationResult()

			if isMounted, err := nodeRes.IsMounted(targetFile); err != nil {
				return fmt.Errorf("failed to determine if %s is already mounted: %v", targetFile, err)
			} else if !isMounted {
				sock, err := m.pathGetter(vmi, i)
				if err != nil {
					return err
				}

				res, err := m.podIsolationDetector.DetectForSocket(vmi, sock)
				if err != nil {
					return fmt.Errorf("failed to detect socket for containerDisk %v: %v", volume.Name, err)
				}
				mountInfo, err := res.MountInfoRoot()
				if err != nil {
					return fmt.Errorf("failed to detect root mount info of containerDisk  %v: %v", volume.Name, err)
				}
				nodeMountInfo, err := nodeRes.ParentMountInfoFor(mountInfo)
				if err != nil {
					return fmt.Errorf("failed to detect root mount point of containerDisk %v on the node: %v", volume.Name, err)
				}
				sourceFile, err := containerdisk.GetImage(filepath.Join(nodeRes.MountRoot(), nodeMountInfo.Root, nodeMountInfo.MountPoint, mountInfo.Root), volume.ContainerDisk.Path)
				if err != nil {
					return fmt.Errorf("failed to find a sourceFile in containerDisk %v: %v", volume.Name, err)
				}
				f, err := os.Create(targetFile)
				if err != nil {
					return fmt.Errorf("failed to create mount point target %v: %v", targetFile, err)
				}
				f.Close()

				log.DefaultLogger().Object(vmi).Infof("Bind mounting container disk at %s to %s", strings.TrimPrefix(sourceFile, nodeRes.MountRoot()), targetFile)
				// #nosec g204 no risk to TrimPref as argument as it just trims two fixed strings
				out, err := exec.Command("/usr/bin/virt-chroot", "--mount", "/proc/1/ns/mnt", "mount", "-o", "ro,bind", strings.TrimPrefix(sourceFile, nodeRes.MountRoot()), targetFile).CombinedOutput()
				if err != nil {
					return fmt.Errorf("failed to bindmount containerDisk %v: %v : %v", volume.Name, string(out), err)
				}
			}
			if verify {
				res, err := m.podIsolationDetector.Detect(vmi)
				if err != nil {
					return fmt.Errorf("failed to detect VMI pod: %v", err)
				}
				imageInfo, err := isolation.GetImageInfo(containerdisk.GetDiskTargetPathFromLauncherView(i), res)
				if err != nil {
					return fmt.Errorf("failed to get image info: %v", err)
				}

				if err := containerdisk.VerifyImage(imageInfo); err != nil {
					return fmt.Errorf("invalid image in containerDisk %v: %v", volume.Name, err)
				}
			}
		}
	}
	return nil
}

// Legacy Unmount unmounts all container disks of a given VMI when the hold HostPath method was in use.
// This exists for backwards compatibility for VMIs running before a KubeVirt update occurs.
func (m *mounter) legacyUnmount(vmi *v1.VirtualMachineInstance) error {
	mountDir := containerdisk.GetLegacyVolumeMountDirOnHost(vmi)

	files, err := ioutil.ReadDir(mountDir)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to list container disk mounts: %v", err)
	}

	if vmi.UID != "" {
		for _, file := range files {
			path := filepath.Join(mountDir, file.Name())
			if strings.HasSuffix(path, ".sock") {
				continue
			}
			if mounted, err := isolation.NodeIsolationResult().IsMounted(path); err != nil {
				return fmt.Errorf("failed to check mount point for containerDisk %v: %v", path, err)
			} else if mounted {
				// #nosec No risk for attacket injection. Parameters are predefined strings
				out, err := exec.Command("/usr/bin/virt-chroot", "--mount", "/proc/1/ns/mnt", "umount", path).CombinedOutput()
				if err != nil {
					return fmt.Errorf("failed to unmount containerDisk %v: %v : %v", path, string(out), err)
				}
			}
		}

		if err := os.RemoveAll(mountDir); err != nil {
			return fmt.Errorf("failed to remove containerDisk files: %v", err)
		}
	}
	return nil
}

// Unmount unmounts all container disks of a given VMI.
func (m *mounter) Unmount(vmi *v1.VirtualMachineInstance) error {
	if vmi.UID != "" {

		// this will catch unmounting a vmi's container disk when
		// an old VMI is left over after a KubeVirt update
		err := m.legacyUnmount(vmi)
		if err != nil {
			return err
		}

		record, err := m.getMountTargetRecord(vmi)
		if err != nil {
			return err
		} else if record == nil {
			// no entries to unmount

			log.DefaultLogger().Object(vmi).Infof("No container disk mount entries found to unmount")
			return nil
		}

		log.DefaultLogger().Object(vmi).Infof("Found container disk mount entries")
		for _, entry := range record.MountTargetEntries {
			path := entry.TargetFile
			log.DefaultLogger().Object(vmi).Infof("Looking to see if containerdisk is mounted at path %s", path)
			if mounted, err := isolation.NodeIsolationResult().IsMounted(path); err != nil {
				return fmt.Errorf("failed to check mount point for containerDisk %v: %v", path, err)
			} else if mounted {
				log.DefaultLogger().Object(vmi).Infof("unmounting container disk at path %s", path)
				// #nosec No risk for attacket injection. Parameters are predefined strings
				out, err := exec.Command("/usr/bin/virt-chroot", "--mount", "/proc/1/ns/mnt", "umount", path).CombinedOutput()
				if err != nil {
					return fmt.Errorf("failed to unmount containerDisk %v: %v : %v", path, string(out), err)
				}
			}

		}
		err = m.deleteMountTargetRecord(vmi)
		if err != nil {
			return err
		}
	}
	return nil
}

func (m *mounter) ContainerDisksReady(vmi *v1.VirtualMachineInstance, notInitializedSince time.Time) (bool, error) {
	for i, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {
			_, err := m.pathGetter(vmi, i)
			if err != nil {
				log.DefaultLogger().Object(vmi).Infof("containerdisk %s not yet ready", volume.Name)
				if time.Now().After(notInitializedSince.Add(m.suppressWarningTimeout)) {
					return false, fmt.Errorf("containerdisk %s still not ready after one minute", volume.Name)
				}
				return false, nil
			}
		}
	}
	log.DefaultLogger().Object(vmi).V(4).Info("all containerdisks are ready")
	return true, nil
}
<|endoftext|>


<commit_before>// go:generate statik -src=./static

package main

import (
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"golang.org/x/crypto/acme/autocert"

	"github.com/gorilla/handlers"
	"github.com/jakdept/dandler"
	_ "github.com/jakdept/sp9k1/statik"
	"github.com/rakyll/statik/fs"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

// now would this be shitposting if there were _tests_?

var serverBanner = `
'______________________________________________________________________________
/                                                                              \
|                 '.'                           .-:::::::::::::::::::::-'      |
|                 -///-'                      '/+++++++++++++++++++++++++-     |
|                ':+++++/-                    -++//////////////////////+++     |
|                /++++++//:.                  -+/----------------------:++     |
|               '/++++///:::.                 -+/---:dddddddds+ymmms----++     |
|             .://+////:::::-                 -+/...:mmmmmdyoymNNNNy...-++     |
|           '://////::::::::::-.              -++::-:mmmdyoymNNNNNNy--:/++     |
|           ://:::::::::::/+++//:'            -+++++ommmhoymNNNNNNNh++++++     |
|           :/syys+:::///++syhyo:.            -+++++ommmmdsodNNNNNNh++++++     |
|         '.+mdo+hmh++++/ommo+yNh-''          -+++++ommmmmmhosmNNNNh++++++     |
|       .//+mN.'''hNs///:mN:'''sNy//-         -+++++ommmmmmmdyohNNNh++++++     |
|      '///+NN.'''hNs::::mM-'''sMy/::-        -++++++dmmmmmmmmdssdNh++++++     |
|      '::::yNd+/yMd::://sNmo/sMm/:::-        -++++++oyyyyyyyyyyo+s+++++++     |
|      '-::::+hmmds//++++/+ydmds/::::-'       -+++++++++++++++++++++++++++     |
|   '-///++++++++++++++++++///:::::////:.     '/+++++++++++++++++++++++++-     |
|  .///////////+hmmmmmmmmmmmh+::://+///::-      .:::::::+++++++++/:::::-'      |
|  ::////////::::+shmmNmmhs+:://++//::::::.             +++++++++:             |
|  -:::::::::::::::::::::///+++///:::::::-'             +++++++++:             |
|   .-:::::::::::::::////+////:::::::::-.'              +++++++++:             |
|      '''''''    ''''''''''''''''''''                  .........'             |
\______________________________________________________________________________/
`

var (
	listenAddress = kingpin.Flag("listen", "addresses to listen for incoming non-TLS connections").
			Short('l').Default("127.0.0.1:8080").TCP()

	enableTLS = kingpin.Flag("tls", "enables TLS for server through Let's Encrypt").Default("false").Bool()
	domain    = kingpin.Flag("hostname", "hostname to register").String()

	imgDir = kingpin.Flag("images", "directory of images to serve").
		Short('i').Default("./").ExistingDir()

	cacheMinDays = kingpin.Flag("cacheMin", "minimum days to cache images in browser").
			Default("30").Int()

	cacheVariation = kingpin.Flag("cacheVariation", "difference between minimum and maximum length to cache images").
			Default("7").Int()

	thumbWidth  = kingpin.Flag("width", "maximum thumbnail width").Default("310").Int()
	thumbHeight = kingpin.Flag("height", "thumbnail height").Default("200").Int()

	staticDir = kingpin.Flag("static", "alternate static directory to serve").Short('s').ExistingDir()

	templateFile = kingpin.Flag("template", "alternate index template to serve").Short('t').ExistingFile()

	canonicalURL        = kingpin.Flag("canonicalURL", "default redirect to serve").Default("localhost:80").String()
	canonicalDisableTLS = kingpin.Flag("canonicalDisableTLS", "force unencrypted protocol").Default("false").Bool()
	canonicalForceTLS   = kingpin.Flag("canonicalForceTLS", "force encrypted protocol").Default("true").Bool()
	canonicalForceHost  = kingpin.Flag("canonicalForceHost", "force a specific hostname").Default("true").Bool()
	canonicalForcePort  = kingpin.Flag("canonicalForcePort", "force a specific port").Default("false").Bool()
)

func parseTemplate(logger *log.Logger, fs http.FileSystem) *template.Template {
	if *templateFile != "" {
		// if an alternate template was provided, i can use that instead
		return template.Must(template.ParseFiles(*templateFile))
	}
	// have to do it the hard way because it comes from fs
	templFile, err := fs.Open("/page.template")
	if err != nil {
		logger.Fatal(err)
	}
	templData, err := ioutil.ReadAll(templFile)
	if err != nil {
		logger.Fatal(err)
	}
	return template.Must(template.New("page.template").Parse(string(templData)))
}

func createStaticFS(logger *log.Logger, path string) http.FileSystem {
	if path != "" {
		return http.Dir(path)
	}
	filesystem, err := fs.New()
	if err != nil {
		logger.Fatal(err)
	}
	return filesystem
}

func buildMuxer(logger *log.Logger,
	fs http.FileSystem,
	templ *template.Template,
	done chan struct{},
) http.Handler {

	day := time.Hour * time.Duration(64)
	var h http.Handler
	mux := http.NewServeMux()

	// building the static handler
	h = http.FileServer(fs)
	// split the main folder off into a redirect
	h = dandler.Split(http.RedirectHandler("/", 302), h)
	// add a prefix before the handler
	h = http.StripPrefix("/static/", h)
	// add some expiration
	h = dandler.ExpiresRange(day*time.Duration(*cacheMinDays),
		day*time.Duration(*cacheVariation), h)
	// add the static handler to the muxer
	mux.Handle("/static/", h)

	// create a caching handler
	h = dandler.ThumbCache(logger, *thumbWidth, *thumbHeight, 32<<20, *imgDir, "thumbs", "jpg")
	// split the folder itself into a redirect
	h = dandler.Split(http.RedirectHandler("/", 302), h)
	// strip the prefix
	h = http.StripPrefix("/thumb/", h)
	// add some expiration
	h = dandler.ExpiresRange(day*time.Duration(*cacheMinDays),
		day*time.Duration(*cacheVariation), h)
	// add the thumbnail handler to the muxer
	mux.Handle("/thumb/", h)

	h = dandler.DirSplit(logger, *imgDir, done,
		dandler.Index(logger, *imgDir, done, templ),
		dandler.ContentType(logger, *imgDir),
	)
	mux.Handle("/", h)

	h = dandler.ASCIIHeader("shit\nposting\n9001", serverBanner, " ", mux)
	h = handlers.CombinedLoggingHandler(os.Stdout, h)

	// add canonical header if required
	if *canonicalForceHost ||
		*canonicalForcePort ||
		*canonicalForceTLS ||
		*canonicalDisableTLS {
		options := 0
		if *canonicalForceHost {
			options += dandler.ForceHost
		}
		if *canonicalForcePort {
			options += dandler.ForcePort
		}
		if *canonicalForceTLS {
			options += dandler.ForceHTTPS
		} else if *canonicalDisableTLS {
			options += dandler.ForceHTTP
		}

		h = dandler.CanonicalHost(*canonicalURL, options, h)
	}

	// compress responses
	h = handlers.CompressHandler(h)

	return h
}

func main() {

	kingpin.CommandLine.HelpFlag.Short('h')
	kingpin.CommandLine.Version("1.0")
	kingpin.CommandLine.Author("jakdept")
	kingpin.Parse()

	logger := log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Llongfile)

	fs := createStaticFS(logger, *staticDir)

	templ := parseTemplate(logger, fs)

	done := make(chan struct{})
	defer close(done)

	srvHandlers := buildMuxer(logger, fs, templ, done)

	var errChan chan error
	go func() {
		errChan <- http.ListenAndServe((*listenAddress).String(), srvHandlers)
	}()
	if *enableTLS {
		go func() {
			errChan <- http.Serve(autocert.NewListener("example.com"), srvHandlers)
		}()
	}
	for e := range errChan {
		log.Fatal(e)
	}
}
<commit_msg>rewrote the binary so the flags are much simpler<commit_after>// go:generate statik -src=./static

package main

import (
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"golang.org/x/crypto/acme/autocert"

	"github.com/gorilla/handlers"
	"github.com/jakdept/dandler"
	_ "github.com/jakdept/sp9k1/statik"
	"github.com/rakyll/statik/fs"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

// now would this be shitposting if there were _tests_?

var serverBanner = `
'______________________________________________________________________________
/                                                                              \
|                 '.'                           .-:::::::::::::::::::::-'      |
|                 -///-'                      '/+++++++++++++++++++++++++-     |
|                ':+++++/-                    -++//////////////////////+++     |
|                /++++++//:.                  -+/----------------------:++     |
|               '/++++///:::.                 -+/---:dddddddds+ymmms----++     |
|             .://+////:::::-                 -+/...:mmmmmdyoymNNNNy...-++     |
|           '://////::::::::::-.              -++::-:mmmdyoymNNNNNNy--:/++     |
|           ://:::::::::::/+++//:'            -+++++ommmhoymNNNNNNNh++++++     |
|           :/syys+:::///++syhyo:.            -+++++ommmmdsodNNNNNNh++++++     |
|         '.+mdo+hmh++++/ommo+yNh-''          -+++++ommmmmmhosmNNNNh++++++     |
|       .//+mN.'''hNs///:mN:'''sNy//-         -+++++ommmmmmmdyohNNNh++++++     |
|      '///+NN.'''hNs::::mM-'''sMy/::-        -++++++dmmmmmmmmdssdNh++++++     |
|      '::::yNd+/yMd::://sNmo/sMm/:::-        -++++++oyyyyyyyyyyo+s+++++++     |
|      '-::::+hmmds//++++/+ydmds/::::-'       -+++++++++++++++++++++++++++     |
|   '-///++++++++++++++++++///:::::////:.     '/+++++++++++++++++++++++++-     |
|  .///////////+hmmmmmmmmmmmh+::://+///::-      .:::::::+++++++++/:::::-'      |
|  ::////////::::+shmmNmmhs+:://++//::::::.             +++++++++:             |
|  -:::::::::::::::::::::///+++///:::::::-'             +++++++++:             |
|   .-:::::::::::::::////+////:::::::::-.'              +++++++++:             |
|      '''''''    ''''''''''''''''''''                  .........'             |
\______________________________________________________________________________/
`

var (
	listenAddress  = kingpin.Flag("listen", "non-TLS listen addresses").Short('l').Default("127.0.0.1:8080").TCPList()
	enableTLS      = kingpin.Flag("tls", "enables auto-TLS and push to https").Default("false").Bool()
	domain         = kingpin.Flag("domain", "domain to register, startup, and redirect to").String()
	imgDir         = kingpin.Flag("images", "directory of images to serve").Short('i').Default("./").ExistingDir()
	cacheMinDays   = kingpin.Flag("cacheMin", "minimum days to cache images in browser").Default("30").Int()
	cacheVariation = kingpin.Flag("cacheDayVariation", "browser cache variation").Default("7").Int()
	thumbWidth     = kingpin.Flag("width", "maximum thumbnail width").Default("310").Int()
	thumbHeight    = kingpin.Flag("height", "thumbnail height").Default("200").Int()
	staticDir      = kingpin.Flag("static", "alternate static directory to serve").Short('s').ExistingDir()
	templateFile   = kingpin.Flag("template", "alternate index template to serve").Short('t').ExistingFile()
)

func parseTemplate(logger *log.Logger, fs http.FileSystem) *template.Template {
	if *templateFile != "" {
		// if an alternate template was provided, i can use that instead
		return template.Must(template.ParseFiles(*templateFile))
	}
	// have to do it the hard way because it comes from fs
	templFile, err := fs.Open("/page.template")
	if err != nil {
		logger.Fatal(err)
	}
	templData, err := ioutil.ReadAll(templFile)
	if err != nil {
		logger.Fatal(err)
	}
	return template.Must(template.New("page.template").Parse(string(templData)))
}

func createStaticFS(logger *log.Logger, path string) http.FileSystem {
	if path != "" {
		return http.Dir(path)
	}
	filesystem, err := fs.New()
	if err != nil {
		logger.Fatal(err)
	}
	return filesystem
}

func buildMuxer(logger *log.Logger,
	fs http.FileSystem,
	templ *template.Template,
	done chan struct{},
) http.Handler {

	day := time.Hour * time.Duration(64)
	var h http.Handler
	mux := http.NewServeMux()

	// building the static handler
	h = http.FileServer(fs)
	// split the main folder off into a redirect
	h = dandler.Split(http.RedirectHandler("/", 302), h)
	// add a prefix before the handler
	h = http.StripPrefix("/static/", h)
	// add some expiration
	h = dandler.ExpiresRange(day*time.Duration(*cacheMinDays),
		day*time.Duration(*cacheVariation), h)
	// add the static handler to the muxer
	mux.Handle("/static/", h)

	// create a caching handler
	h = dandler.ThumbCache(logger, *thumbWidth, *thumbHeight, 32<<20, *imgDir, "thumbs", "jpg")
	// split the folder itself into a redirect
	h = dandler.Split(http.RedirectHandler("/", 302), h)
	// strip the prefix
	h = http.StripPrefix("/thumb/", h)
	// add some expiration
	h = dandler.ExpiresRange(day*time.Duration(*cacheMinDays),
		day*time.Duration(*cacheVariation), h)
	// add the thumbnail handler to the muxer
	mux.Handle("/thumb/", h)

	h = dandler.DirSplit(logger, *imgDir, done,
		dandler.Index(logger, *imgDir, done, templ),
		dandler.ContentType(logger, *imgDir),
	)
	mux.Handle("/", h)

	// add canonical header if required
	if *domain != "" {
		options := 0
		options += dandler.ForceHost
		if *enableTLS {
			options += dandler.ForceHTTPS
			h = dandler.CanonicalHost(fmt.Sprintf("%s:443", *domain), options, h)
		} else {
			h = dandler.CanonicalHost(fmt.Sprintf("%s:%d", *domain,
				(*listenAddress)[0].IP), options, h)
		}
	}

	h = dandler.ASCIIHeader("shit\nposting\n9001", serverBanner, " ", mux)
	h = handlers.CombinedLoggingHandler(os.Stdout, h)

	// compress responses
	h = handlers.CompressHandler(h)

	return h
}

func main() {

	kingpin.CommandLine.HelpFlag.Short('h')
	kingpin.CommandLine.Version("1.0")
	kingpin.CommandLine.Author("jakdept")
	kingpin.Parse()

	if *enableTLS && *domain == "" {
		log.Fatal("failed to start - if you specify --tls you must specify --domain")
	}

	logger := log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Llongfile)

	fs := createStaticFS(logger, *staticDir)

	templ := parseTemplate(logger, fs)

	done := make(chan struct{})
	defer close(done)

	srvHandlers := buildMuxer(logger, fs, templ, done)

	var errChan chan error
	go func() {
		for _, address := range *listenAddress {
			errChan <- http.ListenAndServe(address.String(), srvHandlers)
		}
	}()
	if *enableTLS {
		go func() {
			errChan <- http.Serve(autocert.NewListener(*domain), srvHandlers)
		}()
	}
	for e := range errChan {
		log.Fatal(e)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

var (
	address = flag.String("address", "0.0.0.0", "Listening address")
	port    = flag.String("port", "8080", "Listening port")
	sslPort = flag.String("sslPort", "10433", "SSL listening port")
	path    = flag.String("path", "/", "URL path")
	deny    = flag.String("deny", "", "Sesitive directories or files to be forbidden when listing path (comma sperated)")
	status  = flag.Int("status", 200, "Returned HTTP status code")
	cert    = flag.String("cert", "cert.pem", "SSL certificate path")
	key     = flag.String("key", "key.pem", "SSL private Key path")
)

type bytesHandler []byte

func (h bytesHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(*status)
	w.Write(h)
}

func isDenied(path, denyList string) bool {
	if len(denyList) == 0 {
		return false
	}
	for _, pathElement := range strings.Split(path, string(filepath.Separator)) {
		for _, denyElement := range strings.Split(denyList, ",") {
			match, err := filepath.Match(strings.TrimSpace(denyElement), pathElement)
			if err != nil {
				log.Print("error matching file path element: ", err)
			}
			if match {
				return true
			}
		}
	}
	return false
}

type protectdFileSystem struct {
	fs http.FileSystem
}

func (pfs protectdFileSystem) Open(path string) (http.File, error) {
	if isDenied(path, *deny) {
		return nil, os.ErrPermission
	}
	return pfs.fs.Open(path)
}

func main() {
	flag.Parse()
	listen := *address + ":" + *port
	listenTLS := *address + ":" + *sslPort
	body := flag.Arg(0)
	if body == "" {
		body = "."
	}
	var handler http.Handler
	if fi, err := os.Stat(body); err == nil {
		switch mode := fi.Mode(); {
		case mode.IsDir():
			handler = http.StripPrefix(*path, http.FileServer(protectdFileSystem{http.Dir(body)}))
		case mode.IsRegular():
			if content, err := ioutil.ReadFile(body); err != nil {
				log.Fatal("Error reading file: ", err)
			} else {
				handler = bytesHandler(content)
			}
		}
	} else {
		handler = bytesHandler(body)
	}
	http.Handle(*path, handler)
	go func() {
		if _, err := os.Stat(*cert); err != nil {
			return
		}
		if _, err := os.Stat(*key); err != nil {
			return
		}
		log.Fatal(http.ListenAndServeTLS(listenTLS, *cert, *key, nil))
	}()
	log.Printf("Serving %s on %s%s...", body, listen, *path)
	log.Fatal(http.ListenAndServe(listen, nil))
}
<commit_msg>better flag description<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

var (
	address = flag.String("address", "0.0.0.0", "Listening address")
	port    = flag.String("port", "8080", "Listening port")
	sslPort = flag.String("sslPort", "10433", "SSL listening port")
	path    = flag.String("path", "/", "URL path")
	deny    = flag.String("deny", "", "Sesitive directory or file patterns to be denied when listing path (comma sperated)")
	status  = flag.Int("status", 200, "Returned HTTP status code")
	cert    = flag.String("cert", "cert.pem", "SSL certificate path")
	key     = flag.String("key", "key.pem", "SSL private Key path")
)

type bytesHandler []byte

func (h bytesHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(*status)
	w.Write(h)
}

func isDenied(path, denyList string) bool {
	if len(denyList) == 0 {
		return false
	}
	for _, pathElement := range strings.Split(path, string(filepath.Separator)) {
		for _, denyElement := range strings.Split(denyList, ",") {
			match, err := filepath.Match(strings.TrimSpace(denyElement), pathElement)
			if err != nil {
				log.Print("error matching file path element: ", err)
			}
			if match {
				return true
			}
		}
	}
	return false
}

type protectdFileSystem struct {
	fs http.FileSystem
}

func (pfs protectdFileSystem) Open(path string) (http.File, error) {
	if isDenied(path, *deny) {
		return nil, os.ErrPermission
	}
	return pfs.fs.Open(path)
}

func main() {
	flag.Parse()
	listen := *address + ":" + *port
	listenTLS := *address + ":" + *sslPort
	body := flag.Arg(0)
	if body == "" {
		body = "."
	}
	var handler http.Handler
	if fi, err := os.Stat(body); err == nil {
		switch mode := fi.Mode(); {
		case mode.IsDir():
			handler = http.StripPrefix(*path, http.FileServer(protectdFileSystem{http.Dir(body)}))
		case mode.IsRegular():
			if content, err := ioutil.ReadFile(body); err != nil {
				log.Fatal("Error reading file: ", err)
			} else {
				handler = bytesHandler(content)
			}
		}
	} else {
		handler = bytesHandler(body)
	}
	http.Handle(*path, handler)
	go func() {
		if _, err := os.Stat(*cert); err != nil {
			return
		}
		if _, err := os.Stat(*key); err != nil {
			return
		}
		log.Fatal(http.ListenAndServeTLS(listenTLS, *cert, *key, nil))
	}()
	log.Printf("Serving %s on %s%s...", body, listen, *path)
	log.Fatal(http.ListenAndServe(listen, nil))
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 jdbcio contains cross-language functionality for reading and writing data to JDBC.
// These transforms only work on runners that support cross-language transforms.
//
// Setup
//
// Transforms specified here are cross-language transforms implemented in a
// different SDK (listed below). During pipeline construction, the Go SDK will
// need to connect to an expansion service containing information on these
// transforms in their native SDK.
//
// To use an expansion service, it must be run as a separate process accessible
// during pipeline construction. The address of that process must be passed to
// the transforms in this package.
//
// The version of the expansion service should match the version of the Beam SDK
// being used. For numbered releases of Beam, these expansions services are
// released to the Maven repository as modules. For development versions of
// Beam, it is recommended to build and run it from source using Gradle.
//
// Current supported SDKs, including expansion service modules and reference
// documentation:
// * Java
//    - Vendored Module: beam-sdks-java-extensions-schemaio-expansion-service
//    - Run via Gradle: ./gradlew :sdks:java:extensions:schemaio-expansion-service:build
// 						java -jar <location_of_jar_file_generated_from_above> <port>
//    - Reference Class: org.apache.beam.sdk.io.jdbc.JdbcIO
package jdbcio

import (
	"bytes"
	"fmt"
	"reflect"

	"github.com/apache/beam/sdks/v2/go/pkg/beam"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/xlangx"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
)

func init() {
	beam.RegisterType(reflect.TypeOf((*jdbcConfigSchema)(nil)).Elem())
	beam.RegisterType(reflect.TypeOf((*config)(nil)).Elem())
	beam.RegisterType(reflect.TypeOf((*jdbcConfig)(nil)).Elem())
}

const (
	readURN             = "beam:transform:org.apache.beam:schemaio_jdbc_read:v1"
	writeURN            = "beam:transform:org.apache.beam:schemaio_jdbc_write:v1"
	serviceGradleTarget = ":sdks:java:extensions:schemaio-expansion-service:runExpansionService"
)

var autoStartupAddress string = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget)

// jdbcConfigSchema is the config schema as per the expected corss language payload
// for JDBC IO read and write transform.
type jdbcConfigSchema struct {
	Location   string  `beam:"location"`
	Config     []byte  `beam:"config"`
	DataSchema *[]byte `beam:"dataSchema"`
}

// config is used to set the config field of jdbcConfigSchema. It contains the
// details required to make a connection to the JDBC database.
type config struct {
	DriverClassName       string    `beam:"driverClassName"`
	JDBCUrl               string    `beam:"jdbcUrl"`
	Username              string    `beam:"username"`
	Password              string    `beam:"password"`
	ConnectionProperties  *string   `beam:"connectionProperties"`
	ConnectionInitSQLs    *[]string `beam:"connectionInitSqls"`
	ReadQuery             *string   `beam:"readQuery"`
	WriteStatement        *string   `beam:"writeStatement"`
	FetchSize             *int16    `beam:"fetchSize"`
	OutputParallelization *bool     `beam:"outputParallelization"`
}

// jdbcConfig stores the expansion service and configuration for JDBC IO.
type jdbcConfig struct {
	classpaths    []string
	expansionAddr string
	config        *config
}

// TODO(riteshghorse): update the IO to use wrapper created in BigQueryIO.
func toRow(pl interface{}) []byte {
	rt := reflect.TypeOf(pl)

	enc, err := coder.RowEncoderForStruct(rt)
	if err != nil {
		panic(fmt.Errorf("unable to get row encoder"))
	}
	var buf bytes.Buffer
	if err := enc(pl, &buf); err != nil {
		panic(fmt.Errorf("unable to do row encoding"))
	}
	return buf.Bytes()
}

// Write is a cross-language PTransform which writes Rows to the specified database via JDBC.
// tableName is a required parameter, and by default, the write statement is generated from it.
// The generated write statement can be overridden by passing in a WriteStatement option.
// If an expansion service address is not provided,
// an appropriate expansion service will be automatically started; however
// this is slower than having a persistent expansion service running.
//
// The default write statement is: "INSERT INTO tableName(column1, ...) INTO VALUES(value1, ...)"
// Example:
//   tableName := "roles"
//	 driverClassName := "org.postgresql.Driver"
// 	 username := "root"
// 	 password := "root123"
// 	 jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//	 jdbcio.Write(s, tableName, driverClassName, jdbcurl, username, password, jdbcio.ExpansionAddrWrite("localhost:9000"))
func Write(s beam.Scope, tableName, driverClassName, jdbcUrl, username, password string, col beam.PCollection, opts ...writeOption) {
	s = s.Scope("jdbcio.Write")

	wpl := config{
		DriverClassName: driverClassName,
		JDBCUrl:         jdbcUrl,
		Username:        username,
		Password:        password,
	}
	cfg := jdbcConfig{config: &wpl}
	for _, opt := range opts {
		opt(&cfg)
	}

	expansionAddr := cfg.expansionAddr
	if expansionAddr == "" {
		if len(cfg.classpaths) > 0 {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget, xlangx.AddClasspaths(cfg.classpaths))
		} else {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget)
		}
	}

	jcs := jdbcConfigSchema{
		Location: tableName,
		Config:   toRow(cfg.config),
	}
	pl := beam.CrossLanguagePayload(jcs)
	beam.CrossLanguage(s, writeURN, pl, expansionAddr, beam.UnnamedInput(col), nil)
}

type writeOption func(*jdbcConfig)

func WriteClasspaths(classpaths []string) writeOption {
	return func(jc *jdbcConfig) {
		jc.classpaths = classpaths
	}
}

// WriteStatement option overrides the default write statement of
// "INSERT INTO tableName(column1, ...) INTO VALUES(value1, ...)".
func WriteStatement(statement string) writeOption {
	return func(jc *jdbcConfig) {
		jc.config.WriteStatement = &statement
	}
}

// WriteConnectionProperties properties of the jdbc connection passed as string
// with format [propertyName=property;].
func WriteConnectionProperties(properties string) writeOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionProperties = &properties
	}
}

// ConnectionInitSQLs required only for MySql and MariaDB. passed as list of strings.
func ConnectionInitSQLs(initStatements []string) writeOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionInitSQLs = &initStatements
	}
}

// ExpansionAddrWrite sets the expansion service for JDBC IO.
func ExpansionAddrWrite(expansionAddr string) writeOption {
	return func(jc *jdbcConfig) {
		jc.expansionAddr = expansionAddr
	}
}

// WriteToPostgres is a cross-language PTransform which writes Rows to the postgres database via JDBC.
// tableName is a required parameter, and by default, a write statement is generated from it.
// The generated write statement can be overridden by passing in a WriteStatement option.
// If an expansion service address is not provided,
// an appropriate expansion service will be automatically started; however
// this is slower than having a persistent expansion service running.
// NOTE: This transform uses "org.postgresql.Driver" as the default driver. If you want to use write transform
// with custom postgres driver then use the conventional jdbcio.Write() transform.
//
// The default write statement is: "INSERT INTO tableName(column1, ...) INTO VALUES(value1, ...)"
// Example:
//   tableName := "roles"
// 	 username := "root"
// 	 password := "root123"
// 	 jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//	 jdbcio.WriteToPostgres(s, tableName, jdbcurl, username, password, jdbcio.ExpansionAddrWrite("localhost:9000"))
func WriteToPostgres(s beam.Scope, tableName, jdbcUrl, username, password string, col beam.PCollection, opts ...writeOption) {
	driverClassName := "org.postgresql.Driver"
	Write(s, tableName, driverClassName, jdbcUrl, username, password, col, opts...)
}

// Read is a cross-language PTransform which read Rows from the specified database via JDBC.
// tableName is a required paramater, and by default, the readQuery is generated from it.
// The generated readQuery can be overridden by passing in a readQuery.If an expansion service
// address is not provided, an appropriate expansion service will be automatically started;
// however this is slower than having a persistent expansion service running.
//
// The default read query is "SELECT * FROM tableName;"
//
// Read also accepts optional parameters as readOptions. All optional parameters
// are predefined in this package as functions that return readOption. To set
// an optional parameter, call the function within Read's function signature.
//
// Example:
//   tableName := "roles"
//   driverClassName := "org.postgresql.Driver"
//   username := "root"
//   password := "root123"
//   jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//   outT := reflect.TypeOf((*JdbcTestRow)(nil)).Elem()
//   jdbcio.Read(s, tableName, driverClassName, jdbcurl, username, password, outT, jdbcio.ExpansionAddrRead("localhost:9000"))
func Read(s beam.Scope, tableName, driverClassName, jdbcUrl, username, password string, outT reflect.Type, opts ...readOption) beam.PCollection {
	s = s.Scope("jdbcio.Read")

	rpl := config{
		DriverClassName: driverClassName,
		JDBCUrl:         jdbcUrl,
		Username:        username,
		Password:        password,
	}
	cfg := jdbcConfig{config: &rpl}
	for _, opt := range opts {
		opt(&cfg)
	}

	expansionAddr := cfg.expansionAddr
	if expansionAddr == "" {
		if len(cfg.classpaths) > 0 {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget, xlangx.AddClasspaths(cfg.classpaths))
		} else {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget)
		}
	}

	jcs := jdbcConfigSchema{
		Location: tableName,
		Config:   toRow(cfg.config),
	}
	pl := beam.CrossLanguagePayload(jcs)
	result := beam.CrossLanguage(s, readURN, pl, expansionAddr, nil, beam.UnnamedOutput(typex.New(outT)))
	return result[beam.UnnamedOutputTag()]
}

type readOption func(*jdbcConfig)

func ReadClasspaths(classpaths []string) readOption {
	return func(jc *jdbcConfig) {
		jc.classpaths = classpaths
	}
}

// ReadQuery overrides the default read query "SELECT * FROM tableName;"
func ReadQuery(query string) readOption {
	return func(jc *jdbcConfig) {
		jc.config.ReadQuery = &query
	}
}

// OutputParallelization specifies if output parallelization is on.
func OutputParallelization(status bool) readOption {
	return func(jc *jdbcConfig) {
		jc.config.OutputParallelization = &status
	}
}

// FetchSize specifies how many rows to fetch.
func FetchSize(size int16) readOption {
	return func(jc *jdbcConfig) {
		jc.config.FetchSize = &size
	}
}

// ReadConnectionProperties specifies properties of the jdbc connection passed
// as string with format [propertyName=property;]*
func ReadConnectionProperties(properties string) readOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionProperties = &properties
	}
}

// ReadConnectionInitSQLs required only for MySql and MariaDB.
// passed as list of strings.
func ReadConnectionInitSQLs(initStatements []string) readOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionInitSQLs = &initStatements
	}
}

// ExpansionAddrRead sets the expansion service for JDBC IO.
func ExpansionAddrRead(expansionAddr string) readOption {
	return func(jc *jdbcConfig) {
		jc.expansionAddr = expansionAddr
	}
}

// ReadFromPostgres is a cross-language PTransform which read Rows from the postgres via JDBC.
// tableName is a required parameter, and by default, a read query is generated from it.
// The generated read query can be overridden by passing in a ReadQuery. If an expansion service
// address is not provided, an appropriate expansion service will be automatically started;
// however this is slower than having a persistent expansion service running.
//
// The default read query is "SELECT * FROM tableName;"
//
// Read also accepts optional parameters as readOptions. All optional parameters
// are predefined in this package as functions that return readOption. To set
// an optional parameter, call the function within Read's function signature.
// NOTE: This transform uses "org.postgresql.Driver" as the default driver. If you want to use read transform
// with custom postgres driver then use the conventional jdbcio.Read() transform.
//
// Example:
//   tableName := "roles"
//   username := "root"
//   password := "root123"
//   jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//   outT := reflect.TypeOf((*JdbcTestRow)(nil)).Elem()
//   jdbcio.Read(s, tableName, jdbcurl, username, password, outT, jdbcio.ExpansionAddrRead("localhost:9000"))
func ReadFromPostgres(s beam.Scope, tableName, jdbcUrl, username, password string, outT reflect.Type, opts ...readOption) beam.PCollection {
	driverClassName := "org.postgresql.Driver"
	return Read(s, tableName, driverClassName, jdbcUrl, username, password, outT, opts...)
}
<commit_msg>Add default classpath when not present (#17491)<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 jdbcio contains cross-language functionality for reading and writing data to JDBC.
// These transforms only work on runners that support cross-language transforms.
//
// Setup
//
// Transforms specified here are cross-language transforms implemented in a
// different SDK (listed below). During pipeline construction, the Go SDK will
// need to connect to an expansion service containing information on these
// transforms in their native SDK.
//
// To use an expansion service, it must be run as a separate process accessible
// during pipeline construction. The address of that process must be passed to
// the transforms in this package.
//
// The version of the expansion service should match the version of the Beam SDK
// being used. For numbered releases of Beam, these expansions services are
// released to the Maven repository as modules. For development versions of
// Beam, it is recommended to build and run it from source using Gradle.
//
// Current supported SDKs, including expansion service modules and reference
// documentation:
// * Java
//    - Vendored Module: beam-sdks-java-extensions-schemaio-expansion-service
//    - Run via Gradle: ./gradlew :sdks:java:extensions:schemaio-expansion-service:build
// 						java -jar <location_of_jar_file_generated_from_above> <port>
//    - Reference Class: org.apache.beam.sdk.io.jdbc.JdbcIO
package jdbcio

import (
	"bytes"
	"fmt"
	"reflect"

	"github.com/apache/beam/sdks/v2/go/pkg/beam"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/xlangx"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
)

func init() {
	beam.RegisterType(reflect.TypeOf((*jdbcConfigSchema)(nil)).Elem())
	beam.RegisterType(reflect.TypeOf((*config)(nil)).Elem())
	beam.RegisterType(reflect.TypeOf((*jdbcConfig)(nil)).Elem())
}

const (
	readURN             = "beam:transform:org.apache.beam:schemaio_jdbc_read:v1"
	writeURN            = "beam:transform:org.apache.beam:schemaio_jdbc_write:v1"
	serviceGradleTarget = ":sdks:java:extensions:schemaio-expansion-service:runExpansionService"
)

var defaultClasspaths = map[string][]string{
	"org.postgresql.Driver": []string{"org.postgresql:postgresql:42.3.3"},
	"com.mysql.jdbc.Driver": []string{"mysql:mysql-connector-java:8.0.28"},
}

var autoStartupAddress string = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget)

// jdbcConfigSchema is the config schema as per the expected corss language payload
// for JDBC IO read and write transform.
type jdbcConfigSchema struct {
	Location   string  `beam:"location"`
	Config     []byte  `beam:"config"`
	DataSchema *[]byte `beam:"dataSchema"`
}

// config is used to set the config field of jdbcConfigSchema. It contains the
// details required to make a connection to the JDBC database.
type config struct {
	DriverClassName       string    `beam:"driverClassName"`
	JDBCUrl               string    `beam:"jdbcUrl"`
	Username              string    `beam:"username"`
	Password              string    `beam:"password"`
	ConnectionProperties  *string   `beam:"connectionProperties"`
	ConnectionInitSQLs    *[]string `beam:"connectionInitSqls"`
	ReadQuery             *string   `beam:"readQuery"`
	WriteStatement        *string   `beam:"writeStatement"`
	FetchSize             *int16    `beam:"fetchSize"`
	OutputParallelization *bool     `beam:"outputParallelization"`
}

// jdbcConfig stores the expansion service and configuration for JDBC IO.
type jdbcConfig struct {
	classpaths    []string
	expansionAddr string
	config        *config
}

// TODO(riteshghorse): update the IO to use wrapper created in BigQueryIO.
func toRow(pl interface{}) []byte {
	rt := reflect.TypeOf(pl)

	enc, err := coder.RowEncoderForStruct(rt)
	if err != nil {
		panic(fmt.Errorf("unable to get row encoder"))
	}
	var buf bytes.Buffer
	if err := enc(pl, &buf); err != nil {
		panic(fmt.Errorf("unable to do row encoding"))
	}
	return buf.Bytes()
}

// Write is a cross-language PTransform which writes Rows to the specified database via JDBC.
// tableName is a required parameter, and by default, the write statement is generated from it.
// The generated write statement can be overridden by passing in a WriteStatement option.
// If an expansion service address is not provided,
// an appropriate expansion service will be automatically started; however
// this is slower than having a persistent expansion service running.
//
// If no additional classpaths are provided using jdbcio.WriteClasspaths() then the default classpath
// for that driver would be used. As of now, the default classpaths are present only for PostgreSQL and MySQL.
//
// The default write statement is: "INSERT INTO tableName(column1, ...) INTO VALUES(value1, ...)"
// Example:
//   tableName := "roles"
//	 driverClassName := "org.postgresql.Driver"
// 	 username := "root"
// 	 password := "root123"
// 	 jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//	 jdbcio.Write(s, tableName, driverClassName, jdbcurl, username, password, jdbcio.ExpansionAddrWrite("localhost:9000"))
//
// With Classpath paramater:
//   jdbcio.Write(s, tableName, driverClassName, jdbcurl, username, password, jdbcio.ExpansionAddrWrite("localhost:9000"), jdbcio.WriteClasspaths([]string{"org.postgresql:postgresql:42.3.3"}))
func Write(s beam.Scope, tableName, driverClassName, jdbcUrl, username, password string, col beam.PCollection, opts ...writeOption) {
	s = s.Scope("jdbcio.Write")

	wpl := config{
		DriverClassName: driverClassName,
		JDBCUrl:         jdbcUrl,
		Username:        username,
		Password:        password,
	}
	cfg := jdbcConfig{config: &wpl}
	for _, opt := range opts {
		opt(&cfg)
	}

	if len(cfg.classpaths) == 0 {
		cfg.classpaths = defaultClasspaths[driverClassName]
	}

	expansionAddr := cfg.expansionAddr
	if expansionAddr == "" {
		if len(cfg.classpaths) > 0 {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget, xlangx.AddClasspaths(cfg.classpaths))
		} else {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget)
		}
	}

	jcs := jdbcConfigSchema{
		Location: tableName,
		Config:   toRow(cfg.config),
	}
	pl := beam.CrossLanguagePayload(jcs)
	beam.CrossLanguage(s, writeURN, pl, expansionAddr, beam.UnnamedInput(col), nil)
}

type writeOption func(*jdbcConfig)

func WriteClasspaths(classpaths []string) writeOption {
	return func(jc *jdbcConfig) {
		jc.classpaths = classpaths
	}
}

// WriteStatement option overrides the default write statement of
// "INSERT INTO tableName(column1, ...) INTO VALUES(value1, ...)".
func WriteStatement(statement string) writeOption {
	return func(jc *jdbcConfig) {
		jc.config.WriteStatement = &statement
	}
}

// WriteConnectionProperties properties of the jdbc connection passed as string
// with format [propertyName=property;].
func WriteConnectionProperties(properties string) writeOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionProperties = &properties
	}
}

// ConnectionInitSQLs required only for MySql and MariaDB. passed as list of strings.
func ConnectionInitSQLs(initStatements []string) writeOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionInitSQLs = &initStatements
	}
}

// ExpansionAddrWrite sets the expansion service for JDBC IO.
func ExpansionAddrWrite(expansionAddr string) writeOption {
	return func(jc *jdbcConfig) {
		jc.expansionAddr = expansionAddr
	}
}

// WriteToPostgres is a cross-language PTransform which writes Rows to the postgres database via JDBC.
// tableName is a required parameter, and by default, a write statement is generated from it.
// The generated write statement can be overridden by passing in a WriteStatement option.
// If an expansion service address is not provided,
// an appropriate expansion service will be automatically started; however
// this is slower than having a persistent expansion service running.
// NOTE: This transform uses "org.postgresql.Driver" as the default driver. If you want to use write transform
// with custom postgres driver then use the conventional jdbcio.Write() transform.
//
// The default write statement is: "INSERT INTO tableName(column1, ...) INTO VALUES(value1, ...)"
// Example:
//   tableName := "roles"
// 	 username := "root"
// 	 password := "root123"
// 	 jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//	 jdbcio.WriteToPostgres(s, tableName, jdbcurl, username, password, jdbcio.ExpansionAddrWrite("localhost:9000"))
func WriteToPostgres(s beam.Scope, tableName, jdbcUrl, username, password string, col beam.PCollection, opts ...writeOption) {
	driverClassName := "org.postgresql.Driver"
	Write(s, tableName, driverClassName, jdbcUrl, username, password, col, opts...)
}

// Read is a cross-language PTransform which read Rows from the specified database via JDBC.
// tableName is a required paramater, and by default, the readQuery is generated from it.
// The generated readQuery can be overridden by passing in a readQuery.If an expansion service
// address is not provided, an appropriate expansion service will be automatically started;
// however this is slower than having a persistent expansion service running.
//
// If no additional classpaths are provided using jdbcio.ReadClasspaths() then the default classpath
// for that driver would be used. As of now, the default classpaths are present only for PostgreSQL and MySQL.
//
// The default read query is "SELECT * FROM tableName;"
//
// Read also accepts optional parameters as readOptions. All optional parameters
// are predefined in this package as functions that return readOption. To set
// an optional parameter, call the function within Read's function signature.
//
// Example:
//   tableName := "roles"
//   driverClassName := "org.postgresql.Driver"
//   username := "root"
//   password := "root123"
//   jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//   outT := reflect.TypeOf((*JdbcTestRow)(nil)).Elem()
//   jdbcio.Read(s, tableName, driverClassName, jdbcurl, username, password, outT, jdbcio.ExpansionAddrRead("localhost:9000"))
//
// With Classpath parameter:
//   jdbcio.Read(s, tableName, driverClassName, jdbcurl, username, password, outT, jdbcio.ExpansionAddrRead("localhost:9000"), jdbcio.ReadClasspaths([]string{"org.postgresql:postgresql:42.3.3"})))
func Read(s beam.Scope, tableName, driverClassName, jdbcUrl, username, password string, outT reflect.Type, opts ...readOption) beam.PCollection {
	s = s.Scope("jdbcio.Read")

	rpl := config{
		DriverClassName: driverClassName,
		JDBCUrl:         jdbcUrl,
		Username:        username,
		Password:        password,
	}
	cfg := jdbcConfig{config: &rpl}
	for _, opt := range opts {
		opt(&cfg)
	}

	if len(cfg.classpaths) == 0 {
		cfg.classpaths = defaultClasspaths[driverClassName]
	}

	expansionAddr := cfg.expansionAddr
	if expansionAddr == "" {
		if len(cfg.classpaths) > 0 {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget, xlangx.AddClasspaths(cfg.classpaths))
		} else {
			expansionAddr = xlangx.UseAutomatedJavaExpansionService(serviceGradleTarget)
		}
	}

	jcs := jdbcConfigSchema{
		Location: tableName,
		Config:   toRow(cfg.config),
	}
	pl := beam.CrossLanguagePayload(jcs)
	result := beam.CrossLanguage(s, readURN, pl, expansionAddr, nil, beam.UnnamedOutput(typex.New(outT)))
	return result[beam.UnnamedOutputTag()]
}

type readOption func(*jdbcConfig)

func ReadClasspaths(classpaths []string) readOption {
	return func(jc *jdbcConfig) {
		jc.classpaths = classpaths
	}
}

// ReadQuery overrides the default read query "SELECT * FROM tableName;"
func ReadQuery(query string) readOption {
	return func(jc *jdbcConfig) {
		jc.config.ReadQuery = &query
	}
}

// OutputParallelization specifies if output parallelization is on.
func OutputParallelization(status bool) readOption {
	return func(jc *jdbcConfig) {
		jc.config.OutputParallelization = &status
	}
}

// FetchSize specifies how many rows to fetch.
func FetchSize(size int16) readOption {
	return func(jc *jdbcConfig) {
		jc.config.FetchSize = &size
	}
}

// ReadConnectionProperties specifies properties of the jdbc connection passed
// as string with format [propertyName=property;]*
func ReadConnectionProperties(properties string) readOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionProperties = &properties
	}
}

// ReadConnectionInitSQLs required only for MySql and MariaDB.
// passed as list of strings.
func ReadConnectionInitSQLs(initStatements []string) readOption {
	return func(jc *jdbcConfig) {
		jc.config.ConnectionInitSQLs = &initStatements
	}
}

// ExpansionAddrRead sets the expansion service for JDBC IO.
func ExpansionAddrRead(expansionAddr string) readOption {
	return func(jc *jdbcConfig) {
		jc.expansionAddr = expansionAddr
	}
}

// ReadFromPostgres is a cross-language PTransform which read Rows from the postgres via JDBC.
// tableName is a required parameter, and by default, a read query is generated from it.
// The generated read query can be overridden by passing in a ReadQuery. If an expansion service
// address is not provided, an appropriate expansion service will be automatically started;
// however this is slower than having a persistent expansion service running.
//
// The default read query is "SELECT * FROM tableName;"
//
// Read also accepts optional parameters as readOptions. All optional parameters
// are predefined in this package as functions that return readOption. To set
// an optional parameter, call the function within Read's function signature.
// NOTE: This transform uses "org.postgresql.Driver" as the default driver. If you want to use read transform
// with custom postgres driver then use the conventional jdbcio.Read() transform.
//
// Example:
//   tableName := "roles"
//   username := "root"
//   password := "root123"
//   jdbcUrl := "jdbc:postgresql://localhost:5432/dbname"
//   outT := reflect.TypeOf((*JdbcTestRow)(nil)).Elem()
//   jdbcio.Read(s, tableName, jdbcurl, username, password, outT, jdbcio.ExpansionAddrRead("localhost:9000"))
func ReadFromPostgres(s beam.Scope, tableName, jdbcUrl, username, password string, outT reflect.Type, opts ...readOption) beam.PCollection {
	driverClassName := "org.postgresql.Driver"
	return Read(s, tableName, driverClassName, jdbcUrl, username, password, outT, opts...)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"os"
	"strconv"

	"github.com/golang/glog"

	"github.com/crackcomm/crawl/nsq/consumer"

	imdb "github.com/crackcomm/crawl/examples/imdb/spider"
)

func main() {
	defer glog.Flush()

	// CRAWL_DEBUG environment variable turns on debug mode
	// crawler then can spit out logs using glog.V(3)
	var verbosity string
	if yes, _ := strconv.ParseBool(os.Getenv("CRAWL_DEBUG")); yes {
		verbosity = "-v=3"
	}

	// We are setting glog to log to stderr
	flag.CommandLine.Parse([]string{"-logtostderr", verbosity})

	// Start consumer
	app := consumer.New(
		consumer.WithSpiders(imdb.Register),
	)

	if err := app.Run(os.Args); err != nil {
		glog.Fatal(err)
	}
}
<commit_msg>nsq consumer skeleton: fix<commit_after>package main

import (
	"flag"
	"os"
	"strconv"

	"github.com/golang/glog"

	"github.com/crackcomm/crawl/nsq/consumer"

	imdb "github.com/crackcomm/crawl/examples/imdb/spider"
)

func main() {
	defer glog.Flush()

	// CRAWL_DEBUG environment variable turns on debug mode
	// crawler then can spit out logs using glog.V(3)
	var verbosity string
	if yes, _ := strconv.ParseBool(os.Getenv("CRAWL_DEBUG")); yes {
		verbosity = "-v=3"
	}

	// We are setting glog to log to stderr
	flag.CommandLine.Parse([]string{"-logtostderr", verbosity})

	// Start consumer
	app := consumer.New(
		consumer.WithSpiders(imdb.Spider),
	)

	if err := app.Run(os.Args); err != nil {
		glog.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package binloginfo_test

import (
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/ngaut/log"
	. "github.com/pingcap/check"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/sessionctx/binloginfo"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/codec"
	"github.com/pingcap/tidb/util/testkit"
	"github.com/pingcap/tidb/util/types"
	"github.com/pingcap/tipb/go-binlog"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func TestT(t *testing.T) {
	CustomVerboseFlag = true
	TestingT(t)
}

type mockBinlogPump struct {
	mu struct {
		sync.Mutex
		payloads [][]byte
	}
}

func (p *mockBinlogPump) WriteBinlog(ctx context.Context, req *binlog.WriteBinlogReq) (*binlog.WriteBinlogResp, error) {
	p.mu.Lock()
	p.mu.payloads = append(p.mu.payloads, req.Payload)
	p.mu.Unlock()
	return &binlog.WriteBinlogResp{}, nil
}

// PullBinlogs implements PumpServer interface.
func (p *mockBinlogPump) PullBinlogs(ctx context.Context, req *binlog.PullBinlogReq) (*binlog.PullBinlogResp, error) {
	return &binlog.PullBinlogResp{}, nil
}

var _ = Suite(&testBinlogSuite{})

type testBinlogSuite struct {
	store    kv.Storage
	unixFile string
	serv     *grpc.Server
	pump     *mockBinlogPump
}

func (s *testBinlogSuite) SetUpSuite(c *C) {
	logLevel := os.Getenv("log_level")
	log.SetLevelByString(logLevel)
	store, err := tikv.NewMockTikvStore()
	c.Assert(err, IsNil)
	s.store = store
	tidb.SetSchemaLease(0)
	s.unixFile = "/tmp/mock-binlog-pump"
	l, err := net.Listen("unix", s.unixFile)
	c.Assert(err, IsNil)
	s.serv = grpc.NewServer()
	s.pump = new(mockBinlogPump)
	binlog.RegisterPumpServer(s.serv, s.pump)
	go s.serv.Serve(l)
	opt := grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
		return net.DialTimeout("unix", addr, timeout)
	})
	clientCon, err := grpc.Dial(s.unixFile, opt, grpc.WithInsecure())
	c.Assert(err, IsNil)
	c.Assert(clientCon, NotNil)
	binloginfo.PumpClient = binlog.NewPumpClient(clientCon)
}

func (s *testBinlogSuite) TearDownSuite(c *C) {
	binloginfo.PumpClient = nil
	s.serv.Stop()
	os.Remove(s.unixFile)
	s.store.Close()
}

func (s *testBinlogSuite) TestBinlog(c *C) {
	tk := testkit.NewTestKit(c, s.store)
	tk.MustExec("use test")
	pump := s.pump
	tk.MustExec("drop table if exists local_binlog")
	ddlQuery := "create table local_binlog (id int primary key, name varchar(10))"
	tk.MustExec(ddlQuery)
	var matched bool // got matched pre DDL and commit DDL
	for i := 0; i < 10; i++ {
		preDDL, commitDDL := getLatestDDLBinlog(c, pump, ddlQuery)
		if preDDL.DdlJobId == commitDDL.DdlJobId {
			c.Assert(commitDDL.StartTs, Equals, preDDL.StartTs)
			c.Assert(commitDDL.CommitTs, Greater, commitDDL.StartTs)
			matched = true
			break
		}
		time.Sleep(time.Millisecond * 10)
	}
	c.Assert(matched, IsTrue)

	tk.MustExec("insert local_binlog values (1, 'abc'), (2, 'cde')")
	prewriteVal := getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.SchemaVersion, Greater, int64(0))
	c.Assert(prewriteVal.Mutations[0].TableId, Greater, int64(0))
	expected := [][]types.Datum{
		{types.NewIntDatum(1), types.NewStringDatum("abc")},
		{types.NewIntDatum(2), types.NewStringDatum("cde")},
	}
	gotRows := mutationRowsToRows(c, prewriteVal.Mutations[0].InsertedRows, 0, 2)
	c.Assert(gotRows, DeepEquals, expected)

	tk.MustExec("update local_binlog set name = 'xyz' where id = 2")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	expected = [][]types.Datum{
		{types.NewIntDatum(2), types.NewStringDatum("xyz")},
	}
	gotRows = mutationRowsToRows(c, prewriteVal.Mutations[0].UpdatedRows, 2, 4)
	c.Assert(gotRows, DeepEquals, expected)

	tk.MustExec("delete from local_binlog where id = 1")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].DeletedIds, DeepEquals, []int64{1})

	// Test table primary key is not integer.
	tk.MustExec("create table local_binlog2 (name varchar(64) primary key, age int)")
	tk.MustExec("insert local_binlog2 values ('abc', 16), ('def', 18)")
	tk.MustExec("delete from local_binlog2 where name = 'def'")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].Sequence[0], Equals, binlog.MutationType_DeletePK)
	_, deletedPK, _ := codec.DecodeOne(prewriteVal.Mutations[0].DeletedPks[0])
	c.Assert(deletedPK.GetString(), Equals, "def")

	// Test Table don't have primary key.
	tk.MustExec("create table local_binlog3 (c1 int, c2 int)")
	tk.MustExec("insert local_binlog3 values (1, 2), (1, 3), (2, 3)")
	tk.MustExec("update local_binlog3 set c1 = 3 where c1 = 2")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	gotRows = mutationRowsToRows(c, prewriteVal.Mutations[0].UpdatedRows, 5, 7)
	expected = [][]types.Datum{
		{types.NewIntDatum(3), types.NewIntDatum(3)},
	}
	c.Assert(gotRows, DeepEquals, expected)

	tk.MustExec("delete from local_binlog3 where c1 = 3 and c2 = 3")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].Sequence[0], Equals, binlog.MutationType_DeleteRow)
	gotRows = mutationRowsToRows(c, prewriteVal.Mutations[0].DeletedRows, 1, 3)
	expected = [][]types.Datum{
		{types.NewIntDatum(3), types.NewIntDatum(3)},
	}
	c.Assert(gotRows, DeepEquals, expected)

	// Test Mutation Sequence.
	tk.MustExec("create table local_binlog4 (c1 int primary key, c2 int)")
	tk.MustExec("insert local_binlog4 values (1, 1), (2, 2), (3, 2)")
	tk.MustExec("begin")
	tk.MustExec("delete from local_binlog4 where c1 = 1")
	tk.MustExec("insert local_binlog4 values (1, 1)")
	tk.MustExec("update local_binlog4 set c2 = 3 where c1 = 3")
	tk.MustExec("commit")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].Sequence, DeepEquals, []binlog.MutationType{
		binlog.MutationType_DeleteID,
		binlog.MutationType_Insert,
		binlog.MutationType_Update,
	})

	checkBinlogCount(c, pump)

	pump.mu.Lock()
	originBinlogLen := len(pump.mu.payloads)
	pump.mu.Unlock()
	tk.MustExec("set @@global.autocommit = 0")
	tk.MustExec("set @@global.autocommit = 1")
	pump.mu.Lock()
	newBinlogLen := len(pump.mu.payloads)
	pump.mu.Unlock()
	c.Assert(newBinlogLen, Equals, originBinlogLen)
}

func getLatestBinlogPrewriteValue(c *C, pump *mockBinlogPump) *binlog.PrewriteValue {
	var bin *binlog.Binlog
	pump.mu.Lock()
	for i := len(pump.mu.payloads) - 1; i >= 0; i-- {
		payload := pump.mu.payloads[i]
		bin = new(binlog.Binlog)
		bin.Unmarshal(payload)
		if bin.Tp == binlog.BinlogType_Prewrite {
			break
		}
	}
	pump.mu.Unlock()
	c.Assert(bin, NotNil)
	preVal := new(binlog.PrewriteValue)
	preVal.Unmarshal(bin.PrewriteValue)
	return preVal
}

func getLatestDDLBinlog(c *C, pump *mockBinlogPump, ddlQuery string) (preDDL, commitDDL *binlog.Binlog) {
	pump.mu.Lock()
	for i := len(pump.mu.payloads) - 1; i >= 0; i-- {
		payload := pump.mu.payloads[i]
		bin := new(binlog.Binlog)
		bin.Unmarshal(payload)
		if bin.Tp == binlog.BinlogType_Commit && bin.DdlJobId > 0 {
			commitDDL = bin
		}
		if bin.Tp == binlog.BinlogType_Prewrite && bin.DdlJobId != 0 {
			preDDL = bin
		}
		if preDDL != nil && commitDDL != nil {
			break
		}
	}
	pump.mu.Unlock()
	c.Assert(preDDL.DdlJobId, Greater, int64(0))
	c.Assert(preDDL.StartTs, Greater, int64(0))
	c.Assert(preDDL.CommitTs, Equals, int64(0))
	c.Assert(string(preDDL.DdlQuery), Equals, ddlQuery)
	return
}

func checkBinlogCount(c *C, pump *mockBinlogPump) {
	var bin *binlog.Binlog
	prewriteCount := 0
	ddlCount := 0
	pump.mu.Lock()
	length := len(pump.mu.payloads)
	for i := length - 1; i >= 0; i-- {
		payload := pump.mu.payloads[i]
		bin = new(binlog.Binlog)
		bin.Unmarshal(payload)
		if bin.Tp == binlog.BinlogType_Prewrite {
			if bin.DdlJobId != 0 {
				ddlCount++
			} else {
				prewriteCount++
			}
		}
	}
	pump.mu.Unlock()
	c.Assert(ddlCount, Greater, 0)
	match := false
	for i := 0; i < 10; i++ {
		pump.mu.Lock()
		length = len(pump.mu.payloads)
		pump.mu.Unlock()
		if (prewriteCount+ddlCount)*2 == length {
			match = true
			break
		}
		time.Sleep(time.Millisecond * 10)
	}
	c.Assert(match, IsTrue)
}

func mutationRowsToRows(c *C, mutationRows [][]byte, firstColumn, secondColumn int) [][]types.Datum {
	var rows [][]types.Datum
	for _, mutationRow := range mutationRows {
		datums, err := codec.Decode(mutationRow, 5)
		c.Assert(err, IsNil)
		for i := range datums {
			if i != firstColumn && i != secondColumn {
				// Column ID or handle
				c.Assert(datums[i].GetInt64(), Greater, int64(0))
			}
			if datums[i].Kind() == types.KindBytes {
				datums[i].SetBytesAsString(datums[i].GetBytes())
			}
		}
		row := []types.Datum{datums[firstColumn], datums[secondColumn]}
		rows = append(rows, row)
	}
	return rows
}
<commit_msg>sessionctx/binloginfo: fix race. (#2130)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package binloginfo_test

import (
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/ngaut/log"
	. "github.com/pingcap/check"
	"github.com/pingcap/tidb"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/ddl"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/sessionctx/binloginfo"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/codec"
	"github.com/pingcap/tidb/util/testkit"
	"github.com/pingcap/tidb/util/types"
	"github.com/pingcap/tipb/go-binlog"
	goctx "golang.org/x/net/context"
	"google.golang.org/grpc"
)

func TestT(t *testing.T) {
	CustomVerboseFlag = true
	TestingT(t)
}

type mockBinlogPump struct {
	mu struct {
		sync.Mutex
		payloads [][]byte
	}
}

func (p *mockBinlogPump) WriteBinlog(ctx goctx.Context, req *binlog.WriteBinlogReq) (*binlog.WriteBinlogResp, error) {
	p.mu.Lock()
	p.mu.payloads = append(p.mu.payloads, req.Payload)
	p.mu.Unlock()
	return &binlog.WriteBinlogResp{}, nil
}

// PullBinlogs implements PumpServer interface.
func (p *mockBinlogPump) PullBinlogs(ctx goctx.Context, req *binlog.PullBinlogReq) (*binlog.PullBinlogResp, error) {
	return &binlog.PullBinlogResp{}, nil
}

var _ = Suite(&testBinlogSuite{})

type testBinlogSuite struct {
	store    kv.Storage
	unixFile string
	serv     *grpc.Server
	pump     *mockBinlogPump
	tk       *testkit.TestKit
	ddl      ddl.DDL
}

func (s *testBinlogSuite) SetUpSuite(c *C) {
	logLevel := os.Getenv("log_level")
	log.SetLevelByString(logLevel)
	store, err := tikv.NewMockTikvStore()
	c.Assert(err, IsNil)
	s.store = store
	tidb.SetSchemaLease(0)
	s.unixFile = "/tmp/mock-binlog-pump"
	os.Remove(s.unixFile)
	l, err := net.Listen("unix", s.unixFile)
	c.Assert(err, IsNil)
	s.serv = grpc.NewServer()
	s.pump = new(mockBinlogPump)
	binlog.RegisterPumpServer(s.serv, s.pump)
	go s.serv.Serve(l)
	opt := grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
		return net.DialTimeout("unix", addr, timeout)
	})
	clientCon, err := grpc.Dial(s.unixFile, opt, grpc.WithInsecure())
	c.Assert(err, IsNil)
	c.Assert(clientCon, NotNil)
	binloginfo.PumpClient = binlog.NewPumpClient(clientCon)
	s.tk = testkit.NewTestKit(c, s.store)
	s.tk.MustExec("use test")
	domain := sessionctx.GetDomain(s.tk.Se.(context.Context))
	s.ddl = domain.DDL()
}

func (s *testBinlogSuite) TearDownSuite(c *C) {
	s.ddl.Stop()
	binloginfo.PumpClient = nil
	s.serv.Stop()
	os.Remove(s.unixFile)
	s.store.Close()
}

func (s *testBinlogSuite) TestBinlog(c *C) {
	tk := s.tk
	pump := s.pump
	tk.MustExec("drop table if exists local_binlog")
	ddlQuery := "create table local_binlog (id int primary key, name varchar(10))"
	tk.MustExec(ddlQuery)
	var matched bool // got matched pre DDL and commit DDL
	for i := 0; i < 10; i++ {
		preDDL, commitDDL := getLatestDDLBinlog(c, pump, ddlQuery)
		if preDDL.DdlJobId == commitDDL.DdlJobId {
			c.Assert(commitDDL.StartTs, Equals, preDDL.StartTs)
			c.Assert(commitDDL.CommitTs, Greater, commitDDL.StartTs)
			matched = true
			break
		}
		time.Sleep(time.Millisecond * 10)
	}
	c.Assert(matched, IsTrue)

	tk.MustExec("insert local_binlog values (1, 'abc'), (2, 'cde')")
	prewriteVal := getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.SchemaVersion, Greater, int64(0))
	c.Assert(prewriteVal.Mutations[0].TableId, Greater, int64(0))
	expected := [][]types.Datum{
		{types.NewIntDatum(1), types.NewStringDatum("abc")},
		{types.NewIntDatum(2), types.NewStringDatum("cde")},
	}
	gotRows := mutationRowsToRows(c, prewriteVal.Mutations[0].InsertedRows, 0, 2)
	c.Assert(gotRows, DeepEquals, expected)

	tk.MustExec("update local_binlog set name = 'xyz' where id = 2")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	expected = [][]types.Datum{
		{types.NewIntDatum(2), types.NewStringDatum("xyz")},
	}
	gotRows = mutationRowsToRows(c, prewriteVal.Mutations[0].UpdatedRows, 2, 4)
	c.Assert(gotRows, DeepEquals, expected)

	tk.MustExec("delete from local_binlog where id = 1")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].DeletedIds, DeepEquals, []int64{1})

	// Test table primary key is not integer.
	tk.MustExec("create table local_binlog2 (name varchar(64) primary key, age int)")
	tk.MustExec("insert local_binlog2 values ('abc', 16), ('def', 18)")
	tk.MustExec("delete from local_binlog2 where name = 'def'")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].Sequence[0], Equals, binlog.MutationType_DeletePK)
	_, deletedPK, _ := codec.DecodeOne(prewriteVal.Mutations[0].DeletedPks[0])
	c.Assert(deletedPK.GetString(), Equals, "def")

	// Test Table don't have primary key.
	tk.MustExec("create table local_binlog3 (c1 int, c2 int)")
	tk.MustExec("insert local_binlog3 values (1, 2), (1, 3), (2, 3)")
	tk.MustExec("update local_binlog3 set c1 = 3 where c1 = 2")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	gotRows = mutationRowsToRows(c, prewriteVal.Mutations[0].UpdatedRows, 5, 7)
	expected = [][]types.Datum{
		{types.NewIntDatum(3), types.NewIntDatum(3)},
	}
	c.Assert(gotRows, DeepEquals, expected)

	tk.MustExec("delete from local_binlog3 where c1 = 3 and c2 = 3")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].Sequence[0], Equals, binlog.MutationType_DeleteRow)
	gotRows = mutationRowsToRows(c, prewriteVal.Mutations[0].DeletedRows, 1, 3)
	expected = [][]types.Datum{
		{types.NewIntDatum(3), types.NewIntDatum(3)},
	}
	c.Assert(gotRows, DeepEquals, expected)

	// Test Mutation Sequence.
	tk.MustExec("create table local_binlog4 (c1 int primary key, c2 int)")
	tk.MustExec("insert local_binlog4 values (1, 1), (2, 2), (3, 2)")
	tk.MustExec("begin")
	tk.MustExec("delete from local_binlog4 where c1 = 1")
	tk.MustExec("insert local_binlog4 values (1, 1)")
	tk.MustExec("update local_binlog4 set c2 = 3 where c1 = 3")
	tk.MustExec("commit")
	prewriteVal = getLatestBinlogPrewriteValue(c, pump)
	c.Assert(prewriteVal.Mutations[0].Sequence, DeepEquals, []binlog.MutationType{
		binlog.MutationType_DeleteID,
		binlog.MutationType_Insert,
		binlog.MutationType_Update,
	})

	checkBinlogCount(c, pump)

	pump.mu.Lock()
	originBinlogLen := len(pump.mu.payloads)
	pump.mu.Unlock()
	tk.MustExec("set @@global.autocommit = 0")
	tk.MustExec("set @@global.autocommit = 1")
	pump.mu.Lock()
	newBinlogLen := len(pump.mu.payloads)
	pump.mu.Unlock()
	c.Assert(newBinlogLen, Equals, originBinlogLen)
}

func getLatestBinlogPrewriteValue(c *C, pump *mockBinlogPump) *binlog.PrewriteValue {
	var bin *binlog.Binlog
	pump.mu.Lock()
	for i := len(pump.mu.payloads) - 1; i >= 0; i-- {
		payload := pump.mu.payloads[i]
		bin = new(binlog.Binlog)
		bin.Unmarshal(payload)
		if bin.Tp == binlog.BinlogType_Prewrite {
			break
		}
	}
	pump.mu.Unlock()
	c.Assert(bin, NotNil)
	preVal := new(binlog.PrewriteValue)
	preVal.Unmarshal(bin.PrewriteValue)
	return preVal
}

func getLatestDDLBinlog(c *C, pump *mockBinlogPump, ddlQuery string) (preDDL, commitDDL *binlog.Binlog) {
	pump.mu.Lock()
	for i := len(pump.mu.payloads) - 1; i >= 0; i-- {
		payload := pump.mu.payloads[i]
		bin := new(binlog.Binlog)
		bin.Unmarshal(payload)
		if bin.Tp == binlog.BinlogType_Commit && bin.DdlJobId > 0 {
			commitDDL = bin
		}
		if bin.Tp == binlog.BinlogType_Prewrite && bin.DdlJobId != 0 {
			preDDL = bin
		}
		if preDDL != nil && commitDDL != nil {
			break
		}
	}
	pump.mu.Unlock()
	c.Assert(preDDL.DdlJobId, Greater, int64(0))
	c.Assert(preDDL.StartTs, Greater, int64(0))
	c.Assert(preDDL.CommitTs, Equals, int64(0))
	c.Assert(string(preDDL.DdlQuery), Equals, ddlQuery)
	return
}

func checkBinlogCount(c *C, pump *mockBinlogPump) {
	var bin *binlog.Binlog
	prewriteCount := 0
	ddlCount := 0
	pump.mu.Lock()
	length := len(pump.mu.payloads)
	for i := length - 1; i >= 0; i-- {
		payload := pump.mu.payloads[i]
		bin = new(binlog.Binlog)
		bin.Unmarshal(payload)
		if bin.Tp == binlog.BinlogType_Prewrite {
			if bin.DdlJobId != 0 {
				ddlCount++
			} else {
				prewriteCount++
			}
		}
	}
	pump.mu.Unlock()
	c.Assert(ddlCount, Greater, 0)
	match := false
	for i := 0; i < 10; i++ {
		pump.mu.Lock()
		length = len(pump.mu.payloads)
		pump.mu.Unlock()
		if (prewriteCount+ddlCount)*2 == length {
			match = true
			break
		}
		time.Sleep(time.Millisecond * 10)
	}
	c.Assert(match, IsTrue)
}

func mutationRowsToRows(c *C, mutationRows [][]byte, firstColumn, secondColumn int) [][]types.Datum {
	var rows [][]types.Datum
	for _, mutationRow := range mutationRows {
		datums, err := codec.Decode(mutationRow, 5)
		c.Assert(err, IsNil)
		for i := range datums {
			if i != firstColumn && i != secondColumn {
				// Column ID or handle
				c.Assert(datums[i].GetInt64(), Greater, int64(0))
			}
			if datums[i].Kind() == types.KindBytes {
				datums[i].SetBytesAsString(datums[i].GetBytes())
			}
		}
		row := []types.Datum{datums[firstColumn], datums[secondColumn]}
		rows = append(rows, row)
	}
	return rows
}
<|endoftext|>


<commit_before>package service

import (
	"testing"

	"github.com/builderscon/octav/octav/db"
	"github.com/builderscon/octav/octav/model"
	"github.com/stretchr/testify/assert"
)

func TestSessionPopulateRowForUpdate(t *testing.T) {
	s := Session()

	var vdb db.Session
	var payload model.UpdateSessionRequest

	payload.ID = "abc"
	payload.Status.Set("accepted")
	vdb.EID = payload.ID

	if !assert.NoError(t, s.populateRowForUpdate(&vdb, &payload), "populateRowForUpdate should succeed") {
		return
	}

	if !assert.Equal(t, payload.ID, vdb.EID, "ID should match") {
		return
	}

	if !assert.Equal(t, payload.Status.String, vdb.Status, "Status should match") {
		return
	}

	if !assert.False(t, vdb.StartsOn.Valid, "StartsOn should be invalid") {
		return
	}
}

func TestFormatSessionTweet(t *testing.T) {
	series := model.ConferenceSeries{
		Slug: "builderscon",
	}
	conf := model.Conference{
		Slug: "tokyo/2016",
	}

	session := model.Session{
		ID: "ff8657cb-a751-4415-ad93-374fb9fda2b6",
		Title: "Highly available and scalable Kubernetes on AWS",
	}

	s, err := formatSessionTweet(&session, &conf, &series)
	if !assert.NoError(t, err, "formatSessionTweet should succeed") {
		return
	}

	if !assert.Equal(t, `New submission "Highly available and scalable Kubernetes on AWS" https://builderscon.io/builderscon/tokyo/2016/session/ff8657cb-a751-4415-ad93-374fb9fda2b6`, s, "tweet should match") {
		return
	}
}
<commit_msg>Fix test<commit_after>package service

import (
	"context"
	"testing"

	"github.com/builderscon/octav/octav/db"
	"github.com/builderscon/octav/octav/model"
	"github.com/stretchr/testify/assert"
)

func TestSessionPopulateRowForUpdate(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	s := Session()

	var vdb db.Session
	var payload model.UpdateSessionRequest

	payload.ID = "abc"
	payload.Status.Set("accepted")
	vdb.EID = payload.ID

	if !assert.NoError(t, s.populateRowForUpdate(ctx, &vdb, &payload), "populateRowForUpdate should succeed") {
		return
	}

	if !assert.Equal(t, payload.ID, vdb.EID, "ID should match") {
		return
	}

	if !assert.Equal(t, payload.Status.String, vdb.Status, "Status should match") {
		return
	}

	if !assert.False(t, vdb.StartsOn.Valid, "StartsOn should be invalid") {
		return
	}
}

func TestFormatSessionTweet(t *testing.T) {
	series := model.ConferenceSeries{
		Slug: "builderscon",
	}
	conf := model.Conference{
		Slug: "tokyo/2016",
	}

	session := model.Session{
		ID:    "ff8657cb-a751-4415-ad93-374fb9fda2b6",
		Title: "Highly available and scalable Kubernetes on AWS",
	}

	s, err := formatSessionTweet(&session, &conf, &series)
	if !assert.NoError(t, err, "formatSessionTweet should succeed") {
		return
	}

	if !assert.Equal(t, `New submission "Highly available and scalable Kubernetes on AWS" https://builderscon.io/builderscon/tokyo/2016/session/ff8657cb-a751-4415-ad93-374fb9fda2b6`, s, "tweet should match") {
		return
	}
}
<|endoftext|>


<commit_before>package swarm

import (
	"time"

	types "github.com/docker/docker/api/types"
	context "golang.org/x/net/context"

	"github.com/ovh/cds/sdk/log"
)

func (h *HatcherySwarm) pullImage(img string, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	//Pull the worker image
	opts := types.ImagePullOptions{}
	log.Info("CanSpawn> pulling image %s", img)
	res, err := h.dockerClient.ImagePull(ctx, img, opts)
	if err != nil {
		log.Warning("CanSpawn> Unable to pull image %s : %s", img, err)
		return err
	}
	return res.Close()
}
<commit_msg>fix(hatchery/swarm): pull image (#2388)<commit_after>package swarm

import (
	"io/ioutil"
	"time"

	types "github.com/docker/docker/api/types"
	context "golang.org/x/net/context"

	"github.com/ovh/cds/sdk/log"
)

func (h *HatcherySwarm) pullImage(img string, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	//Pull the worker image
	opts := types.ImagePullOptions{}
	log.Info("pullImage> pulling image %s", img)
	res, err := h.dockerClient.ImagePull(ctx, img, opts)
	if err != nil {
		log.Warning("pullImage> Unable to pull image %s : %s", img, err)
		return err
	}

	btes, _ := ioutil.ReadAll(res)
	log.Debug("pullImage> %s", string(btes))
	if err := res.Close(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>