<commit_before>// +build windows
package tap

import (
	"errors"
	"golang.org/x/sys/windows/registry"
	"log"
	"os"
	"syscall"
)

var (
	// Device Control Codes
	tap_win_ioctl_get_mac             = tap_control_code(1, 0)
	tap_win_ioctl_get_version         = tap_control_code(2, 0)
	tap_win_ioctl_get_mtu             = tap_control_code(3, 0)
	tap_win_ioctl_get_info            = tap_control_code(4, 0)
	tap_ioctl_config_point_to_point   = tap_control_code(5, 0)
	tap_ioctl_set_media_status        = tap_control_code(6, 0)
	tap_win_ioctl_config_dhcp_masq    = tap_control_code(7, 0)
	tap_win_ioctl_get_log_line        = tap_control_code(8, 0)
	tap_win_ioctl_config_dhcp_set_opt = tap_control_code(9, 0)
	tap_ioctl_config_tun              = tap_control_code(10, 0)
	// w32 api
	file_device_unknown = uint32(0x00000022)
)

func ctl_code(device_type, function, method, access uint32) uint32 {
	return (device_type << 16) | (access << 14) | (function << 2) | method
}

func tap_control_code(request, method uint32) uint32 {
	return ctl_code(file_device_unknown, request, method, 0)
}

// GetDeviceId finds out a TAP device from registry, it requires privileged right.
func getdeviceid() (string, string, error) {
	// TAP driver key location
	regkey := `SYSTEM\CurrentControlSet\Control\Class\{4D36E972-E325-11CE-BFC1-08002BE10318}`
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, regkey, registry.ALL_ACCESS)
	if err != nil {
		log.Println("Cannot open the reg key:", err, "Please run this program with privileged right.")
		return "", "", err
	}
	defer k.Close()
	// read all subkeys
	keys, err := k.ReadSubKeyNames(-1)
	if err != nil {
		log.Println("Cannot read subkeys:", err, "Please run this program with privileged right.")
		return "", "", err
	}
	// find the one with ComponentId == "tap0901"
	for _, v := range keys {
		key, err := registry.OpenKey(registry.LOCAL_MACHINE, regkey+"\\"+v, registry.ALL_ACCESS)
		if err != nil {
			log.Println("Failed to open subkey:", err)
			continue
		}
		val, _, err := key.GetStringValue("ComponentId")
		if err != nil {
			log.Println("Failed to get subkey value:", err)
			goto next
		}
		if val == "tap0901" {
			val, _, err = key.GetStringValue("NetCfgInstanceId")
			if err != nil {
				log.Println("err read NetCfgInstanceId:", err)
				goto next
			}
			name, _, err := key.GetStringValue("DeviceInstanceID")
			if err != nil {
				log.Println("err read DeviceInstanceID:", err)
				goto next
			}
			key.Close()
			return val, name, nil
		}
	next:
		key.Close()
	}
	return "", "", errors.New("Device not found")
}

// NewTAP find and open a TAP device.
func newTAP() (ifce *Interface, err error) {
	deviceid, name, err := getdeviceid()
	if err != nil {
		return nil, errors.New("Failed to get DeviceId:" + err.Error())
	}
	path := "\\\\.\\Global\\" + deviceid + ".tap"
	pathp, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return nil, errors.New("Invalid Device path:" + err.Error())
	}
	// type Handle uintptr
	file, err := syscall.CreateFile(pathp, syscall.GENERIC_READ|syscall.GENERIC_WRITE, uint32(syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE), nil, syscall.OPEN_EXISTING, syscall.FILE_ATTRIBUTE_SYSTEM, 0)
	defer func() {
		if err := recover(); err != nil {
			syscall.Close(file)
		}
	}()
	if err != nil {
		return nil, errors.New("Failed to open device:" + err.Error())
	}
	var bytesReturned uint32
	fd := os.NewFile(uintptr(file), path)
	mac := make([]byte, 6)
	err = syscall.DeviceIoControl(file, tap_win_ioctl_get_mac, &mac[0], uint32(len(mac)), &mac[0], uint32(len(mac)), &bytesReturned, nil)
	if err != nil {
		log.Println("Failed to get mac address of the interface: ", err)
		fd.Close()
		return nil, err
	}
	rdbbuf := make([]byte, syscall.MAXIMUM_REPARSE_DATA_BUFFER_SIZE)
	code := []byte{0x01, 0x00, 0x00, 0x00}
	err = syscall.DeviceIoControl(file, tap_ioctl_set_media_status, &code[0], uint32(4), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil)
	if err != nil {
		return nil, errors.New("Failed to bring up device:" + err.Error())
	}
	//TUN
	//code2 := []byte{0x0a, 0x03, 0x00, 0x01, 0x0a, 0x03, 0x00, 0x00, 0xff, 0xff, 0xff, 0x00}
	//err = syscall.DeviceIoControl(file, tap_ioctl_config_tun, &code2[0], uint32(12), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil)
	//if err != nil {
	//	log.Fatalln("code2 err:", err)
	//}
	ifce = &Interface{tap: true, file: fd, name: name, mac: mac}
	return
}
<commit_msg>Properly process the error.<commit_after>// +build windows
package tap

import (
	"errors"
	"golang.org/x/sys/windows/registry"
	"os"
	"syscall"
)

var (
	IfceNameNotFound  = errors.New("Failed to find the name of interface.")
	TapDeviceNotFound = errors.New("Failed to find the tap device in registry.")
	// Device Control Codes
	tap_win_ioctl_get_mac             = tap_control_code(1, 0)
	tap_win_ioctl_get_version         = tap_control_code(2, 0)
	tap_win_ioctl_get_mtu             = tap_control_code(3, 0)
	tap_win_ioctl_get_info            = tap_control_code(4, 0)
	tap_ioctl_config_point_to_point   = tap_control_code(5, 0)
	tap_ioctl_set_media_status        = tap_control_code(6, 0)
	tap_win_ioctl_config_dhcp_masq    = tap_control_code(7, 0)
	tap_win_ioctl_get_log_line        = tap_control_code(8, 0)
	tap_win_ioctl_config_dhcp_set_opt = tap_control_code(9, 0)
	tap_ioctl_config_tun              = tap_control_code(10, 0)
	// w32 api
	file_device_unknown = uint32(0x00000022)
)

func ctl_code(device_type, function, method, access uint32) uint32 {
	return (device_type << 16) | (access << 14) | (function << 2) | method
}

func tap_control_code(request, method uint32) uint32 {
	return ctl_code(file_device_unknown, request, method, 0)
}

// GetDeviceId finds out a TAP device from registry, it requires privileged right.
func getdeviceid() (string, string, error) {
	// TAP driver key location
	regkey := `SYSTEM\CurrentControlSet\Control\Class\{4D36E972-E325-11CE-BFC1-08002BE10318}`
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, regkey, registry.ALL_ACCESS)
	if err != nil {
		return "", "", err
	}
	defer k.Close()
	// read all subkeys
	keys, err := k.ReadSubKeyNames(-1)
	if err != nil {
		return "", "", err
	}
	// find the one with ComponentId == "tap0901"
	for _, v := range keys {
		key, err := registry.OpenKey(registry.LOCAL_MACHINE, regkey+"\\"+v, registry.ALL_ACCESS)
		if err != nil {
			continue
		}
		val, _, err := key.GetStringValue("ComponentId")
		if err != nil {
			goto next
		}
		if val == "tap0901" {
			val, _, err = key.GetStringValue("NetCfgInstanceId")
			if err != nil {
				goto next
			}
			name, _, err := key.GetStringValue("DeviceInstanceID")
			if err != nil {
				log.Println("err read DeviceInstanceID:", err)
				goto next
			}
			key.Close()
			return val, name, nil
		}
	next:
		key.Close()
	}
	return "", "", errors.New("Device not found")
}

// NewTAP find and open a TAP device.
func newTAP() (ifce *Interface, err error) {
	deviceid, name, err := getdeviceid()
	if err != nil {
		return nil, err
	}
	path := "\\\\.\\Global\\" + deviceid + ".tap"
	pathp, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return nil, err
	}
	// type Handle uintptr
	file, err := syscall.CreateFile(pathp, syscall.GENERIC_READ|syscall.GENERIC_WRITE, uint32(syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE), nil, syscall.OPEN_EXISTING, syscall.FILE_ATTRIBUTE_SYSTEM, 0)
	// if err hanppens, close the interface.
	defer func() {
		if err := recover(); err != nil {
			syscall.Close(file)
		}
	}()
	if err != nil {
		return nil, err
	}
	var bytesReturned uint32
	// find the mac address of tap device.
	mac := make([]byte, 6)
	err = syscall.DeviceIoControl(file, tap_win_ioctl_get_mac, &mac[0], uint32(len(mac)), &mac[0], uint32(len(mac)), &bytesReturned, nil)
	if err != nil {
		return nil, err
	}
	// bring up device.
	rdbbuf := make([]byte, syscall.MAXIMUM_REPARSE_DATA_BUFFER_SIZE)
	code := []byte{0x01, 0x00, 0x00, 0x00}
	err = syscall.DeviceIoControl(file, tap_ioctl_set_media_status, &code[0], uint32(4), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil)
	if err != nil {
		return nil, err
	}
	//TUN
	//code2 := []byte{0x0a, 0x03, 0x00, 0x01, 0x0a, 0x03, 0x00, 0x00, 0xff, 0xff, 0xff, 0x00}
	//err = syscall.DeviceIoControl(file, tap_ioctl_config_tun, &code2[0], uint32(12), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil)
	//if err != nil {
	//	log.Fatalln("code2 err:", err)
	//}
	fd := os.NewFile(uintptr(file), path)
	ifce = &Interface{tap: true, file: fd}
	copy(ifce.mac[:6], mac[:6])

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package singleton

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strings"
	"testing"
	"time"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/vt/schema"

	"vitess.io/vitess/go/test/endtoend/cluster"
	"vitess.io/vitess/go/test/endtoend/onlineddl"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	vtParams        mysql.ConnParams

	hostname                   = "localhost"
	keyspaceName               = "ks"
	cell                       = "zone1"
	schemaChangeDirectory      = ""
	tableName                  = `onlineddl_test`
	createTableWrapper         = `CREATE TABLE onlineddl_test(%s)`
	ddlStrategy                = "online -declarative -allow-zero-in-date -postpone-completion"
	onlineSingletonDDLStrategy = `online -singleton`
	createStatement            = `
		CREATE TABLE stress_test (
			id bigint(20) not null,
			rand_val varchar(32) null default '',
			hint_col varchar(64) not null default 'just-created',
			created_timestamp timestamp not null default current_timestamp,
			updates int unsigned not null default 0,
			PRIMARY KEY (id),
			key created_idx(created_timestamp),
			key updates_idx(updates)
		) ENGINE=InnoDB
	`
	// We will run this query with "gh-ost --max-load=Threads_running=1"
	alterTableThrottlingStatement = `
		ALTER TABLE stress_test DROP COLUMN created_timestamp
	`
	multiAlterTableThrottlingStatement = `
		ALTER TABLE stress_test ENGINE=InnoDB;
		ALTER TABLE stress_test ENGINE=InnoDB;
		ALTER TABLE stress_test ENGINE=InnoDB;
	`
	// A trivial statement which must succeed and does not change the schema
	alterTableTrivialStatement = `
		ALTER TABLE stress_test ENGINE=InnoDB
	`
	dropStatement = `
		DROP TABLE stress_test
	`
	multiDropStatements = `DROP TABLE IF EXISTS t1; DROP TABLE IF EXISTS t2; DROP TABLE IF EXISTS t3;`
)

type testCase struct {
	fromSchema                  string
	toSchema                    string
	expectProblems              bool
	removedUniqueKeyNames       string
	droppedNoDefaultColumnNames string
	expandedColumnNames         string
}

var testCases = []testCase{
	{
		fromSchema: `id int primary key, i1 int not null default 0`,
		toSchema:   `id int primary key, i2 int not null default 0`,
	},
	{
		fromSchema: `id int primary key, i1 int not null default 0, unique key i1_uidx(i1)`,
		toSchema:   `id int primary key, i1 int not null default 0, i2 int not null default 0, unique key i1_uidx(i1)`,
	},
	{
		fromSchema:            `id int primary key, i1 int not null default 0, unique key i1_uidx(i1)`,
		toSchema:              `id int primary key, i2 int not null default 0`,
		removedUniqueKeyNames: `i1_uidx`,
	},
	{
		fromSchema:                  `id int primary key, i1 int not null`,
		toSchema:                    `id int primary key, i2 int not null default 0`,
		droppedNoDefaultColumnNames: `i1`,
	},
}

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		clusterInstance.VtTabletExtraArgs = []string{
			"-enable-lag-throttler",
			"-throttle_threshold", "1s",
			"-heartbeat_enable",
			"-heartbeat_interval", "250ms",
		}
		clusterInstance.VtGateExtraArgs = []string{}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		// No need for replicas in this stress test
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 0, false); err != nil {
			return 1, err
		}

		vtgateInstance := clusterInstance.NewVtgateInstance()
		// set the gateway we want to use
		vtgateInstance.GatewayImplementation = "tabletgateway"
		// Start vtgate
		if err := vtgateInstance.Setup(); err != nil {
			return 1, err
		}
		// ensure it is torn down during cluster TearDown
		clusterInstance.VtgateProcess = *vtgateInstance
		vtParams = mysql.ConnParams{
			Host: clusterInstance.Hostname,
			Port: clusterInstance.VtgateMySQLPort,
		}

		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestSchemaChange(t *testing.T) {
	defer cluster.PanicHandler(t)
	shards := clusterInstance.Keyspaces[0].Shards
	require.Equal(t, 1, len(shards))

	var uuids []string
	// CREATE
	t.Run("CREATE TABLE", func(t *testing.T) {
		// The table does not exist
		uuid := testOnlineDDLStatement(t, createStatement, onlineSingletonDDLStrategy, "vtgate", "", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		checkTable(t, tableName, true)
	})
	t.Run("revert CREATE TABLE", func(t *testing.T) {
		// The table existed, so it will now be dropped (renamed)
		uuid := testRevertMigration(t, uuids[len(uuids)-1], "vtgate", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		checkTable(t, tableName, false)
	})
	t.Run("revert revert CREATE TABLE", func(t *testing.T) {
		// Table was dropped (renamed) so it will now be restored
		uuid := testRevertMigration(t, uuids[len(uuids)-1], "vtgate", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		checkTable(t, tableName, true)
	})

	var throttledUUID string
	t.Run("throttled migration", func(t *testing.T) {
		throttledUUID = testOnlineDDLStatement(t, alterTableThrottlingStatement, "gh-ost -singleton --max-load=Threads_running=1", "vtgate", "hint_col", "", false)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, throttledUUID, schema.OnlineDDLStatusRunning)
	})
	t.Run("failed singleton migration, vtgate", func(t *testing.T) {
		uuid := testOnlineDDLStatement(t, alterTableThrottlingStatement, "gh-ost -singleton --max-load=Threads_running=1", "vtgate", "hint_col", "rejected", true)
		assert.Empty(t, uuid)
	})
	t.Run("failed singleton migration, vtctl", func(t *testing.T) {
		uuid := testOnlineDDLStatement(t, alterTableThrottlingStatement, "gh-ost -singleton --max-load=Threads_running=1", "vtctl", "hint_col", "rejected", true)
		assert.Empty(t, uuid)
	})
	t.Run("failed revert migration", func(t *testing.T) {
		uuid := testRevertMigration(t, throttledUUID, "vtgate", "rejected", true)
		assert.Empty(t, uuid)
	})
	t.Run("terminate throttled migration", func(t *testing.T) {
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, throttledUUID, schema.OnlineDDLStatusRunning)
		onlineddl.CheckCancelMigration(t, &vtParams, shards, throttledUUID, true)
		time.Sleep(2 * time.Second)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, throttledUUID, schema.OnlineDDLStatusFailed)
	})
	t.Run("successful gh-ost alter, vtctl", func(t *testing.T) {
		uuid := testOnlineDDLStatement(t, alterTableTrivialStatement, "gh-ost -singleton", "vtctl", "hint_col", "", false)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		onlineddl.CheckCancelMigration(t, &vtParams, shards, uuid, false)
		onlineddl.CheckRetryMigration(t, &vtParams, shards, uuid, false)
	})
	t.Run("successful gh-ost alter, vtgate", func(t *testing.T) {
		uuid := testOnlineDDLStatement(t, alterTableTrivialStatement, "gh-ost -singleton", "vtgate", "hint_col", "", false)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		onlineddl.CheckCancelMigration(t, &vtParams, shards, uuid, false)
		onlineddl.CheckRetryMigration(t, &vtParams, shards, uuid, false)
	})

	t.Run("successful online alter, vtgate", func(t *testing.T) {
		uuid := testOnlineDDLStatement(t, alterTableTrivialStatement, onlineSingletonDDLStrategy, "vtgate", "hint_col", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		onlineddl.CheckCancelMigration(t, &vtParams, shards, uuid, false)
		onlineddl.CheckRetryMigration(t, &vtParams, shards, uuid, false)
		checkTable(t, tableName, true)
	})
	t.Run("revert ALTER TABLE, vttablet", func(t *testing.T) {
		// The table existed, so it will now be dropped (renamed)
		uuid := testRevertMigration(t, uuids[len(uuids)-1], "vttablet", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		checkTable(t, tableName, true)
	})

	var throttledUUIDs []string
	// singleton-context
	t.Run("throttled migrations, singleton-context", func(t *testing.T) {
		uuidList := testOnlineDDLStatement(t, multiAlterTableThrottlingStatement, "gh-ost -singleton-context --max-load=Threads_running=1", "vtctl", "hint_col", "", false)
		throttledUUIDs = strings.Split(uuidList, "\n")
		assert.Equal(t, 3, len(throttledUUIDs))
		for _, uuid := range throttledUUIDs {
			onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusRunning, schema.OnlineDDLStatusQueued)
		}
	})
	t.Run("failed migrations, singleton-context", func(t *testing.T) {
		_ = testOnlineDDLStatement(t, multiAlterTableThrottlingStatement, "gh-ost -singleton-context --max-load=Threads_running=1", "vtctl", "hint_col", "rejected", false)
	})
	t.Run("terminate throttled migrations", func(t *testing.T) {
		for _, uuid := range throttledUUIDs {
			onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusRunning, schema.OnlineDDLStatusQueued)
			onlineddl.CheckCancelMigration(t, &vtParams, shards, uuid, true)
		}
		time.Sleep(2 * time.Second)
		for _, uuid := range throttledUUIDs {
			uuid = strings.TrimSpace(uuid)
			onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusFailed, schema.OnlineDDLStatusCancelled)
		}
	})

	//DROP

	t.Run("online DROP TABLE", func(t *testing.T) {
		uuid := testOnlineDDLStatement(t, dropStatement, onlineSingletonDDLStrategy, "vtgate", "", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		checkTable(t, tableName, false)
	})
	t.Run("revert DROP TABLE", func(t *testing.T) {
		// This will recreate the table (well, actually, rename it back into place)
		uuid := testRevertMigration(t, uuids[len(uuids)-1], "vttablet", "", false)
		uuids = append(uuids, uuid)
		onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
		checkTable(t, tableName, true)
	})

	// Last two tests (we run an incomplete migration)
	t.Run("submit successful migration, no wait, vtgate", func(t *testing.T) {
		_ = testOnlineDDLStatement(t, alterTableTrivialStatement, "gh-ost -singleton", "vtgate", "hint_col", "", true)
	})
	t.Run("fail submit migration, no wait, vtgate", func(t *testing.T) {
		_ = testOnlineDDLStatement(t, alterTableTrivialStatement, "gh-ost -singleton", "vtgate", "hint_col", "rejected", true)
	})
}

// testOnlineDDLStatement runs an online DDL, ALTER statement
func testOnlineDDLStatement(t *testing.T, alterStatement string, ddlStrategy string, executeStrategy string, expectHint string, expectError string, skipWait bool) (uuid string) {
	strategySetting, err := schema.ParseDDLStrategy(ddlStrategy)
	require.NoError(t, err)

	if executeStrategy == "vtgate" {
		result := onlineddl.VtgateExecDDL(t, &vtParams, ddlStrategy, alterStatement, expectError)
		if result != nil {
			row := result.Named().Row()
			if row != nil {
				uuid = row.AsString("uuid", "")
			}
		}
	} else {
		output, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, alterStatement, cluster.VtctlClientParams{DDLStrategy: ddlStrategy, SkipPreflight: true})
		if expectError == "" {
			assert.NoError(t, err)
			uuid = output
		} else {
			assert.Error(t, err)
			assert.Contains(t, output, expectError)
		}
	}
	uuid = strings.TrimSpace(uuid)
	fmt.Println("# Generated UUID (for debug purposes):")
	fmt.Printf("<%s>\n", uuid)

	if !strategySetting.Strategy.IsDirect() && !skipWait {
		time.Sleep(time.Second * 20)
	}

	if expectError == "" && expectHint != "" {
		checkMigratedTable(t, tableName, expectHint)
	}
	return uuid
}

// testRevertMigration reverts a given migration
func testRevertMigration(t *testing.T, revertUUID string, executeStrategy string, expectError string, skipWait bool) (uuid string) {
	revertQuery := fmt.Sprintf("revert vitess_migration '%s'", revertUUID)
	if executeStrategy == "vtgate" {
		result := onlineddl.VtgateExecDDL(t, &vtParams, onlineSingletonDDLStrategy, revertQuery, expectError)
		if result != nil {
			row := result.Named().Row()
			if row != nil {
				uuid = row.AsString("uuid", "")
			}
		}
	} else {
		output, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, revertQuery, cluster.VtctlClientParams{DDLStrategy: onlineSingletonDDLStrategy, SkipPreflight: true})
		if expectError == "" {
			assert.NoError(t, err)
			uuid = output
		} else {
			assert.Error(t, err)
			assert.Contains(t, output, expectError)
		}
	}

	if expectError == "" {
		uuid = strings.TrimSpace(uuid)
		fmt.Println("# Generated UUID (for debug purposes):")
		fmt.Printf("<%s>\n", uuid)
	}
	if !skipWait {
		time.Sleep(time.Second * 20)
	}
	return uuid
}

// checkTable checks the number of tables in the first two shards.
func checkTable(t *testing.T, showTableName string, expectExists bool) bool {
	expectCount := 0
	if expectExists {
		expectCount = 1
	}
	for i := range clusterInstance.Keyspaces[0].Shards {
		if !checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], showTableName, expectCount) {
			return false
		}
	}
	return true
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, showTableName string, expectCount int) bool {
	query := fmt.Sprintf(`show tables like '%%%s%%';`, showTableName)
	queryResult, err := tablet.VttabletProcess.QueryTablet(query, keyspaceName, true)
	require.Nil(t, err)
	return assert.Equal(t, expectCount, len(queryResult.Rows))
}

// checkMigratedTables checks the CREATE STATEMENT of a table after migration
func checkMigratedTable(t *testing.T, tableName, expectHint string) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		createStatement := getCreateTableStatement(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], tableName)
		assert.Contains(t, createStatement, expectHint)
	}
}

// getCreateTableStatement returns the CREATE TABLE statement for a given table
func getCreateTableStatement(t *testing.T, tablet *cluster.Vttablet, tableName string) (statement string) {
	queryResult, err := tablet.VttabletProcess.QueryTablet(fmt.Sprintf("show create table %s;", tableName), keyspaceName, true)
	require.Nil(t, err)

	assert.Equal(t, len(queryResult.Rows), 1)
	assert.Equal(t, len(queryResult.Rows[0]), 2) // table name, create statement
	statement = queryResult.Rows[0][1].ToString()
	return statement
}
<commit_msg>more tests<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package singleton

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strings"
	"testing"
	"time"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/vt/schema"

	"vitess.io/vitess/go/test/endtoend/cluster"
	"vitess.io/vitess/go/test/endtoend/onlineddl"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	shards          []cluster.Shard
	vtParams        mysql.ConnParams

	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	tableName             = `onlineddl_test`
	createTableWrapper    = `CREATE TABLE onlineddl_test(%s)`
	dropTableStatement    = `
		DROP TABLE IF EXISTS onlineddl_test
	`
	ddlStrategy = "online -declarative -allow-zero-in-date"
)

type testCase struct {
	name       string
	fromSchema string
	toSchema   string
	// expectProblems              bool
	removedUniqueKeyNames       string
	droppedNoDefaultColumnNames string
	expandedColumnNames         string
}

var testCases = []testCase{
	{
		name:       "identical schemas",
		fromSchema: `id int primary key, i1 int not null default 0`,
		toSchema:   `id int primary key, i2 int not null default 0`,
	},
	{
		name:       "different schemas, nothing to note",
		fromSchema: `id int primary key, i1 int not null default 0, unique key i1_uidx(i1)`,
		toSchema:   `id int primary key, i1 int not null default 0, i2 int not null default 0, unique key i1_uidx(i1)`,
	},
	{
		name:                  "removed non-nullable unique key",
		fromSchema:            `id int primary key, i1 int not null default 0, unique key i1_uidx(i1)`,
		toSchema:              `id int primary key, i2 int not null default 0`,
		removedUniqueKeyNames: `i1_uidx`,
	},
	{
		name:                  "removed nullable unique key",
		fromSchema:            `id int primary key, i1 int default null, unique key i1_uidx(i1)`,
		toSchema:              `id int primary key, i2 int default null`,
		removedUniqueKeyNames: `i1_uidx`,
	},
	{
		name:                  "expanding unique key removes unique constraint",
		fromSchema:            `id int primary key, i1 int default null, unique key i1_uidx(i1)`,
		toSchema:              `id int primary key, i1 int default null, unique key i1_uidx(i1, id)`,
		removedUniqueKeyNames: `i1_uidx`,
	},
	{
		name:                  "reducing unique key does not unique constraint",
		fromSchema:            `id int primary key, i1 int default null, unique key i1_uidx(i1, id)`,
		toSchema:              `id int primary key, i1 int default null, unique key i1_uidx(i1)`,
		removedUniqueKeyNames: ``,
	},
	{
		name:                        "remove column without default",
		fromSchema:                  `id int primary key, i1 int not null`,
		toSchema:                    `id int primary key, i2 int not null default 0`,
		droppedNoDefaultColumnNames: `i1`,
	},
}

func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		clusterInstance.VtTabletExtraArgs = []string{
			"-enable-lag-throttler",
			"-throttle_threshold", "1s",
			"-heartbeat_enable",
			"-heartbeat_interval", "250ms",
		}
		clusterInstance.VtGateExtraArgs = []string{}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		// No need for replicas in this stress test
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 0, false); err != nil {
			return 1, err
		}

		vtgateInstance := clusterInstance.NewVtgateInstance()
		// set the gateway we want to use
		vtgateInstance.GatewayImplementation = "tabletgateway"
		// Start vtgate
		if err := vtgateInstance.Setup(); err != nil {
			return 1, err
		}
		// ensure it is torn down during cluster TearDown
		clusterInstance.VtgateProcess = *vtgateInstance
		vtParams = mysql.ConnParams{
			Host: clusterInstance.Hostname,
			Port: clusterInstance.VtgateMySQLPort,
		}

		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestSchemaChange(t *testing.T) {
	defer cluster.PanicHandler(t)
	shards = clusterInstance.Keyspaces[0].Shards
	require.Equal(t, 1, len(shards))

	for _, testcase := range testCases {
		t.Run(testcase.name, func(t *testing.T) {

			t.Run("ensure table dropped", func(t *testing.T) {
				uuid := testOnlineDDLStatement(t, dropTableStatement, ddlStrategy, "vtgate", "", "", false)
				onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
				checkTable(t, tableName, false)
			})

			t.Run("create from-table", func(t *testing.T) {
				fromStatement := fmt.Sprintf(createTableWrapper, testcase.fromSchema)
				uuid := testOnlineDDLStatement(t, fromStatement, ddlStrategy, "vtgate", "", "", false)
				onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
				checkTable(t, tableName, true)
			})
			var uuid string
			t.Run("run migration", func(t *testing.T) {
				toStatement := fmt.Sprintf(createTableWrapper, testcase.toSchema)
				uuid = testOnlineDDLStatement(t, toStatement, ddlStrategy, "vtgate", "", "", false)
				onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
				checkTable(t, tableName, true)
			})
			t.Run("check migration", func(t *testing.T) {
				rs := onlineddl.ReadMigrations(t, &vtParams, uuid)
				require.NotNil(t, rs)
				for _, row := range rs.Named().Rows {
					removedUniqueKeyNames := row.AsString("removed_unique_key_names", "")
					droppedNoDefaultColumnNames := row.AsString("dropped_no_default_column_names", "")
					expandedColumnNames := row.AsString("expanded_column_names", "")

					assert.Equal(t, removedUniqueKeyNames, testcase.removedUniqueKeyNames)
					assert.Equal(t, droppedNoDefaultColumnNames, testcase.droppedNoDefaultColumnNames)
					assert.Equal(t, expandedColumnNames, testcase.expandedColumnNames)
				}
			})
		})
	}
}

// testOnlineDDLStatement runs an online DDL, ALTER statement
func testOnlineDDLStatement(t *testing.T, alterStatement string, ddlStrategy string, executeStrategy string, expectHint string, expectError string, skipWait bool) (uuid string) {
	strategySetting, err := schema.ParseDDLStrategy(ddlStrategy)
	require.NoError(t, err)

	if executeStrategy == "vtgate" {
		result := onlineddl.VtgateExecDDL(t, &vtParams, ddlStrategy, alterStatement, expectError)
		if result != nil {
			row := result.Named().Row()
			if row != nil {
				uuid = row.AsString("uuid", "")
			}
		}
	} else {
		output, err := clusterInstance.VtctlclientProcess.ApplySchemaWithOutput(keyspaceName, alterStatement, cluster.VtctlClientParams{DDLStrategy: ddlStrategy, SkipPreflight: true})
		if expectError == "" {
			assert.NoError(t, err)
			uuid = output
		} else {
			assert.Error(t, err)
			assert.Contains(t, output, expectError)
		}
	}
	uuid = strings.TrimSpace(uuid)
	fmt.Println("# Generated UUID (for debug purposes):")
	fmt.Printf("<%s>\n", uuid)

	if !strategySetting.Strategy.IsDirect() && !skipWait {
		status := onlineddl.WaitForMigrationStatus(t, &vtParams, shards, uuid, 20*time.Second, schema.OnlineDDLStatusComplete, schema.OnlineDDLStatusFailed)
		fmt.Printf("# Migration status (for debug purposes): <%s>\n", status)
	}

	if expectError == "" && expectHint != "" {
		checkMigratedTable(t, tableName, expectHint)
	}
	return uuid
}

// checkTable checks the number of tables in the first two shards.
func checkTable(t *testing.T, showTableName string, expectExists bool) bool {
	expectCount := 0
	if expectExists {
		expectCount = 1
	}
	for i := range clusterInstance.Keyspaces[0].Shards {
		if !checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], showTableName, expectCount) {
			return false
		}
	}
	return true
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, showTableName string, expectCount int) bool {
	query := fmt.Sprintf(`show tables like '%%%s%%';`, showTableName)
	queryResult, err := tablet.VttabletProcess.QueryTablet(query, keyspaceName, true)
	require.Nil(t, err)
	return assert.Equal(t, expectCount, len(queryResult.Rows))
}

// checkMigratedTables checks the CREATE STATEMENT of a table after migration
func checkMigratedTable(t *testing.T, tableName, expectHint string) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		createStatement := getCreateTableStatement(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], tableName)
		assert.Contains(t, createStatement, expectHint)
	}
}

// getCreateTableStatement returns the CREATE TABLE statement for a given table
func getCreateTableStatement(t *testing.T, tablet *cluster.Vttablet, tableName string) (statement string) {
	queryResult, err := tablet.VttabletProcess.QueryTablet(fmt.Sprintf("show create table %s;", tableName), keyspaceName, true)
	require.Nil(t, err)

	assert.Equal(t, len(queryResult.Rows), 1)
	assert.Equal(t, len(queryResult.Rows[0]), 2) // table name, create statement
	statement = queryResult.Rows[0][1].ToString()
	return statement
}
<|endoftext|>


<commit_before>package awslogs

import "github.com/aws/aws-sdk-go/service/cloudwatchlogs"

type mockcwlogsclient struct {
	createLogStreamArgument chan *cloudwatchlogs.CreateLogStreamInput
	createLogStreamResult   chan *createLogStreamResult
	putLogEventsArgument    chan *cloudwatchlogs.PutLogEventsInput
	putLogEventsResult      chan *putLogEventsResult
}

type createLogStreamResult struct {
	successResult *cloudwatchlogs.CreateLogStreamOutput
	errorResult   error
}

type putLogEventsResult struct {
	successResult *cloudwatchlogs.PutLogEventsOutput
	errorResult   error
}

func newMockClient() *mockcwlogsclient {
	return &mockcwlogsclient{
		createLogStreamArgument: make(chan *cloudwatchlogs.CreateLogStreamInput, 1),
		createLogStreamResult:   make(chan *createLogStreamResult, 1),
		putLogEventsArgument:    make(chan *cloudwatchlogs.PutLogEventsInput, 1),
		putLogEventsResult:      make(chan *putLogEventsResult, 1),
	}
}

func newMockClientBuffered(buflen int) *mockcwlogsclient {
	return &mockcwlogsclient{
		createLogStreamArgument: make(chan *cloudwatchlogs.CreateLogStreamInput, buflen),
		createLogStreamResult:   make(chan *createLogStreamResult, buflen),
		putLogEventsArgument:    make(chan *cloudwatchlogs.PutLogEventsInput, buflen),
		putLogEventsResult:      make(chan *putLogEventsResult, buflen),
	}
}

func (m *mockcwlogsclient) CreateLogStream(input *cloudwatchlogs.CreateLogStreamInput) (*cloudwatchlogs.CreateLogStreamOutput, error) {
	m.createLogStreamArgument <- input
	output := <-m.createLogStreamResult
	return output.successResult, output.errorResult
}

func (m *mockcwlogsclient) PutLogEvents(input *cloudwatchlogs.PutLogEventsInput) (*cloudwatchlogs.PutLogEventsOutput, error) {
	m.putLogEventsArgument <- input
	output := <-m.putLogEventsResult
	return output.successResult, output.errorResult
}

type mockmetadataclient struct {
	regionResult chan *regionResult
}

type regionResult struct {
	successResult string
	errorResult   error
}

func newMockMetadataClient() *mockmetadataclient {
	return &mockmetadataclient{
		regionResult: make(chan *regionResult, 1),
	}
}

func (m *mockmetadataclient) Region() (string, error) {
	output := <-m.regionResult
	return output.successResult, output.errorResult
}
<commit_msg>awslogs: Fix a race in mockcwlogsclient<commit_after>package awslogs

import "github.com/aws/aws-sdk-go/service/cloudwatchlogs"

type mockcwlogsclient struct {
	createLogStreamArgument chan *cloudwatchlogs.CreateLogStreamInput
	createLogStreamResult   chan *createLogStreamResult
	putLogEventsArgument    chan *cloudwatchlogs.PutLogEventsInput
	putLogEventsResult      chan *putLogEventsResult
}

type createLogStreamResult struct {
	successResult *cloudwatchlogs.CreateLogStreamOutput
	errorResult   error
}

type putLogEventsResult struct {
	successResult *cloudwatchlogs.PutLogEventsOutput
	errorResult   error
}

func newMockClient() *mockcwlogsclient {
	return &mockcwlogsclient{
		createLogStreamArgument: make(chan *cloudwatchlogs.CreateLogStreamInput, 1),
		createLogStreamResult:   make(chan *createLogStreamResult, 1),
		putLogEventsArgument:    make(chan *cloudwatchlogs.PutLogEventsInput, 1),
		putLogEventsResult:      make(chan *putLogEventsResult, 1),
	}
}

func newMockClientBuffered(buflen int) *mockcwlogsclient {
	return &mockcwlogsclient{
		createLogStreamArgument: make(chan *cloudwatchlogs.CreateLogStreamInput, buflen),
		createLogStreamResult:   make(chan *createLogStreamResult, buflen),
		putLogEventsArgument:    make(chan *cloudwatchlogs.PutLogEventsInput, buflen),
		putLogEventsResult:      make(chan *putLogEventsResult, buflen),
	}
}

func (m *mockcwlogsclient) CreateLogStream(input *cloudwatchlogs.CreateLogStreamInput) (*cloudwatchlogs.CreateLogStreamOutput, error) {
	m.createLogStreamArgument <- input
	output := <-m.createLogStreamResult
	return output.successResult, output.errorResult
}

func (m *mockcwlogsclient) PutLogEvents(input *cloudwatchlogs.PutLogEventsInput) (*cloudwatchlogs.PutLogEventsOutput, error) {
	events := make([]*cloudwatchlogs.InputLogEvent, len(input.LogEvents))
	copy(events, input.LogEvents)
	m.putLogEventsArgument <- &cloudwatchlogs.PutLogEventsInput{
		LogEvents:     events,
		SequenceToken: input.SequenceToken,
		LogGroupName:  input.LogGroupName,
		LogStreamName: input.LogStreamName,
	}
	output := <-m.putLogEventsResult
	return output.successResult, output.errorResult
}

type mockmetadataclient struct {
	regionResult chan *regionResult
}

type regionResult struct {
	successResult string
	errorResult   error
}

func newMockMetadataClient() *mockmetadataclient {
	return &mockmetadataclient{
		regionResult: make(chan *regionResult, 1),
	}
}

func (m *mockmetadataclient) Region() (string, error) {
	output := <-m.regionResult
	return output.successResult, output.errorResult
}
<|endoftext|>


<commit_before>package dnsr

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	lru "github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

//go:generate sh generate.sh

var (
	Root        *Resolver
	DebugLogger io.Writer
)

func init() {
	Root = New(strings.Count(root, "\n"))
	for t := range dns.ParseZone(strings.NewReader(root), "", "") {
		if t.Error == nil {
			Root.saveDNSRR(t.RR)
		}
	}
}

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache  *lru.Cache
	client *dns.Client
}

// New initializes a Resolver with the specified cache size. Cache size defaults to 10,000 if size <= 0.
func New(size int) *Resolver {
	if size <= 0 {
		size = 10000
	}
	cache, _ := lru.New(size)
	r := &Resolver{
		client: &dns.Client{},
		cache:  cache,
	}
	return r
}

// Resolve finds DNS records of type qtype for the domain qname. It returns a channel of *RR.
// The implementation guarantees that the output channel will close, so it is safe to range over.
// For nonexistent domains (where a DNS server will return NXDOMAIN), it will simply close the output channel.
// Specify an empty string in qtype to receive any DNS records found (currently A, AAAA, NS, CNAME, and TXT).
func (r *Resolver) Resolve(qname string, qtype string) <-chan *RR {
	return r.resolve(qname, qtype, 0)
}

func (r *Resolver) resolve(qname string, qtype string, depth int) <-chan *RR {
	c := make(chan *RR, 20)
	go func() {
		if DebugLogger != nil {
			start := time.Now()
			defer func() {
				dur := time.Since(start)
				if dur >= 1*time.Millisecond {
					fmt.Fprintf(DebugLogger, "%s└─── %dms: resolve(\"%s\", \"%s\", %d)\n",
						strings.Repeat("│   ", depth), dur/time.Millisecond, qname, qtype, depth)
				}
			}()
		}
		qname = toLowerFQDN(qname)
		defer close(c)
		if rrs := r.cacheGet(qname, qtype); rrs != nil {
			inject(c, rrs...)
			return
		}
		if DebugLogger != nil {
			fmt.Fprintf(DebugLogger, "%s┌─── resolve(\"%s\", \"%s\")\n",
				strings.Repeat("│   ", depth), qname, qtype)
		}
		pname, ok := qname, true
		if qtype == "NS" {
			pname, ok = parent(qname)
			if !ok {
				return
			}
		}
	outer:
		for ; ok; pname, ok = parent(pname) {
			for nrr := range r.resolve(pname, "NS", depth+1) {
				if qtype != "" {
					if rrs := r.cacheGet(qname, qtype); rrs != nil {
						inject(c, rrs...)
						return
					}
				}
				if nrr.Type != "NS" {
					continue
				}
				for arr := range r.resolve(nrr.Value, "A", depth+1) {
					if arr.Type != "A" { // FIXME: support AAAA records?
						continue
					}
					addr := arr.Value + ":53"
					dtype, ok := dns.StringToType[qtype]
					if !ok {
						dtype = dns.TypeA
					}
					qmsg := &dns.Msg{}
					qmsg.SetQuestion(qname, dtype)
					qmsg.MsgHdr.RecursionDesired = false
					// fmt.Printf(";; dig +norecurse @%s %s %s\n", a.A.String(), qname, dns.TypeToString[qtype])
					rmsg, dur, err := r.client.Exchange(qmsg, addr)
					if err != nil {
						continue // FIXME: handle errors better from flaky/failing NS servers
					}
					if DebugLogger != nil {
						fmt.Fprintf(DebugLogger, "%s│    %dms: dig @%s %s %s\n", strings.Repeat("│   ", depth), dur/time.Millisecond, arr.Value, qname, dns.TypeToString[dtype])
					}
					r.saveDNSRR(rmsg.Answer...)
					r.saveDNSRR(rmsg.Ns...)
					r.saveDNSRR(rmsg.Extra...)
					if rmsg.Rcode == dns.RcodeNameError {
						r.cacheAdd(qname, nil) // FIXME: cache NXDOMAIN responses responsibly
						return
					}
					break outer
				}
			}
		}

		if rrs := r.cacheGet(qname, ""); rrs != nil {
			if !inject(c, rrs...) {
				return
			}

			for _, crr := range rrs {
				if crr.Type != "CNAME" {
					continue
				}
				if DebugLogger != nil {
					fmt.Fprintf(DebugLogger, "%s│    CNAME: %s\n", strings.Repeat("│   ", depth), crr.String())
				}
				for rr := range r.resolve(crr.Value, qtype, depth+1) {
					r.cacheAdd(qname, rr)
					if !inject(c, rr) {
						return
					}
				}
			}
		}
	}()
	return c
}

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

func convertRR(drr dns.RR) *RR {
	switch t := drr.(type) {
	case *dns.NS:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Ns}
	case *dns.CNAME:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Target}
	case *dns.A:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.A.String()}
	case *dns.AAAA:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.AAAA.String()}
	case *dns.TXT:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], strings.Join(t.Txt, "\t")}
	default:
		// fmt.Printf("%s\n", drr.String())
	}
	return nil
}

func inject(c chan<- *RR, rrs ...*RR) bool {
	for _, rr := range rrs {
		select {
		case c <- rr:
		default:
			// return false
		}
	}
	return true
}

func parent(name string) (string, bool) {
	labels := dns.SplitDomainName(name)
	if labels == nil {
		return "", false
	}
	return toLowerFQDN(strings.Join(labels[1:], ".")), true
}

func toLowerFQDN(name string) string {
	return dns.Fqdn(strings.ToLower(name))
}

type key struct {
	Name string
	Type string
}

type entry struct {
	m   sync.RWMutex
	rrs map[RR]struct{}
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(drrs ...dns.RR) {
	for _, drr := range drrs {
		if rr := convertRR(drr); rr != nil {
			r.cacheAdd(rr.Name, rr)
		}
	}
}

// cacheAdd adds 0 or more DNS records to the resolver cache for a specific
// domain name and record type. This ensures the cache entry exists, even
// if empty, for NXDOMAIN responses.
func (r *Resolver) cacheAdd(qname string, rr *RR) {
	qname = toLowerFQDN(qname)
	e := r.getEntry(qname)
	if e == nil {
		e = &entry{rrs: make(map[RR]struct{}, 0)}
		e.m.Lock()
		r.cache.Add(qname, e)
	} else {
		e.m.Lock()
	}
	defer e.m.Unlock()
	if rr != nil {
		e.rrs[*rr] = struct{}{}
	}
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(qname string, qtype string) []*RR {
	e := r.getEntry(qname)
	if e == nil && r != Root {
		e = Root.getEntry(qname)
	}
	if e == nil {
		return nil
	}
	e.m.RLock()
	defer e.m.RUnlock()
	if len(e.rrs) == 0 {
		return []*RR{}
	}
	rrs := make([]*RR, 0, len(e.rrs))
	for rr, _ := range e.rrs {
		// fmt.Printf("%s\n", rr.String())
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, &RR{rr.Name, rr.Type, rr.Value})
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil
	}
	return rrs
}

// getEntry returns a single cache entry or nil if an entry does not exist in the cache.
func (r *Resolver) getEntry(qname string) *entry {
	c, ok := r.cache.Get(qname)
	if !ok {
		return nil
	}
	e, ok := c.(*entry)
	if !ok {
		return nil
	}
	return e
}
<commit_msg>timeout at 500ms<commit_after>package dnsr

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	lru "github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

//go:generate sh generate.sh

var (
	Root        *Resolver
	DebugLogger io.Writer
	Timeout     = 500 * time.Millisecond
)

func init() {
	Root = New(strings.Count(root, "\n"))
	for t := range dns.ParseZone(strings.NewReader(root), "", "") {
		if t.Error == nil {
			Root.saveDNSRR(t.RR)
		}
	}
}

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache  *lru.Cache
	client *dns.Client
}

// New initializes a Resolver with the specified cache size. Cache size defaults to 10,000 if size <= 0.
func New(size int) *Resolver {
	if size <= 0 {
		size = 10000
	}
	cache, _ := lru.New(size)
	r := &Resolver{
		client: &dns.Client{
			DialTimeout:  Timeout,
			ReadTimeout:  Timeout,
			WriteTimeout: Timeout,
		},
		cache: cache,
	}
	return r
}

// Resolve finds DNS records of type qtype for the domain qname. It returns a channel of *RR.
// The implementation guarantees that the output channel will close, so it is safe to range over.
// For nonexistent domains (where a DNS server will return NXDOMAIN), it will simply close the output channel.
// Specify an empty string in qtype to receive any DNS records found (currently A, AAAA, NS, CNAME, and TXT).
func (r *Resolver) Resolve(qname string, qtype string) <-chan *RR {
	return r.resolve(qname, qtype, 0)
}

func (r *Resolver) resolve(qname string, qtype string, depth int) <-chan *RR {
	c := make(chan *RR, 20)
	go func() {
		if DebugLogger != nil {
			fmt.Fprintf(DebugLogger, "%s┌─── resolve(\"%s\", \"%s\", %d)\n",
				strings.Repeat("│   ", depth), qname, qtype, depth)

			start := time.Now()
			defer func() {
				dur := time.Since(start)
				if true || dur >= 1*time.Millisecond {
					fmt.Fprintf(DebugLogger, "%s└─── %dms: resolve(\"%s\", \"%s\", %d)\n",
						strings.Repeat("│   ", depth), dur/time.Millisecond, qname, qtype, depth)
				}
			}()
		}
		qname = toLowerFQDN(qname)
		defer close(c)
		if rrs := r.cacheGet(qname, qtype); rrs != nil {
			inject(c, rrs...)
			return
		}
		pname, ok := qname, true
		if qtype == "NS" {
			pname, ok = parent(qname)
			if !ok {
				return
			}
		}
	outer:
		for ; ok; pname, ok = parent(pname) {
			for nrr := range r.resolve(pname, "NS", depth+1) {
				if qtype != "" {
					if rrs := r.cacheGet(qname, qtype); rrs != nil {
						inject(c, rrs...)
						return
					}
				}
				if nrr.Type != "NS" {
					continue
				}
				for arr := range r.resolve(nrr.Value, "A", depth+1) {
					if arr.Type != "A" { // FIXME: support AAAA records?
						continue
					}
					addr := arr.Value + ":53"
					dtype, ok := dns.StringToType[qtype]
					if !ok {
						dtype = dns.TypeA
					}
					qmsg := &dns.Msg{}
					qmsg.SetQuestion(qname, dtype)
					qmsg.MsgHdr.RecursionDesired = false
					// fmt.Printf(";; dig +norecurse @%s %s %s\n", a.A.String(), qname, dns.TypeToString[qtype])
					start := time.Now()
					rmsg, _, err := r.client.Exchange(qmsg, addr)
					dur := time.Since(start)
					if DebugLogger != nil {
						fmt.Fprintf(DebugLogger, "%s│    %dms: dig @%s %s %s\n",
							strings.Repeat("│   ", depth), dur/time.Millisecond, arr.Value, qname, dns.TypeToString[dtype])
						if err != nil {
							fmt.Fprintf(DebugLogger, "%s│    %dms: ERROR: %s\n",
								strings.Repeat("│   ", depth), dur/time.Millisecond, err.Error())
						}
					}
					if err != nil {
						continue // FIXME: handle errors better from flaky/failing NS servers
					}
					r.saveDNSRR(rmsg.Answer...)
					r.saveDNSRR(rmsg.Ns...)
					r.saveDNSRR(rmsg.Extra...)
					if rmsg.Rcode == dns.RcodeNameError {
						r.cacheAdd(qname, nil) // FIXME: cache NXDOMAIN responses responsibly
						return
					}
					break outer
				}
			}
		}

		if rrs := r.cacheGet(qname, ""); rrs != nil {
			if !inject(c, rrs...) {
				return
			}

			for _, crr := range rrs {
				if crr.Type != "CNAME" {
					continue
				}
				if DebugLogger != nil {
					fmt.Fprintf(DebugLogger, "%s│    CNAME: %s\n", strings.Repeat("│   ", depth), crr.String())
				}
				for rr := range r.resolve(crr.Value, qtype, depth+1) {
					r.cacheAdd(qname, rr)
					if !inject(c, rr) {
						return
					}
				}
			}
		}
	}()
	return c
}

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

func convertRR(drr dns.RR) *RR {
	switch t := drr.(type) {
	case *dns.NS:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Ns}
	case *dns.CNAME:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Target}
	case *dns.A:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.A.String()}
	case *dns.AAAA:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.AAAA.String()}
	case *dns.TXT:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], strings.Join(t.Txt, "\t")}
	default:
		// fmt.Printf("%s\n", drr.String())
	}
	return nil
}

func inject(c chan<- *RR, rrs ...*RR) bool {
	for _, rr := range rrs {
		select {
		case c <- rr:
		default:
			// return false
		}
	}
	return true
}

func parent(name string) (string, bool) {
	labels := dns.SplitDomainName(name)
	if labels == nil {
		return "", false
	}
	return toLowerFQDN(strings.Join(labels[1:], ".")), true
}

func toLowerFQDN(name string) string {
	return dns.Fqdn(strings.ToLower(name))
}

type key struct {
	Name string
	Type string
}

type entry struct {
	m   sync.RWMutex
	rrs map[RR]struct{}
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(drrs ...dns.RR) {
	for _, drr := range drrs {
		if rr := convertRR(drr); rr != nil {
			r.cacheAdd(rr.Name, rr)
		}
	}
}

// cacheAdd adds 0 or more DNS records to the resolver cache for a specific
// domain name and record type. This ensures the cache entry exists, even
// if empty, for NXDOMAIN responses.
func (r *Resolver) cacheAdd(qname string, rr *RR) {
	qname = toLowerFQDN(qname)
	e := r.getEntry(qname)
	if e == nil {
		e = &entry{rrs: make(map[RR]struct{}, 0)}
		e.m.Lock()
		r.cache.Add(qname, e)
	} else {
		e.m.Lock()
	}
	defer e.m.Unlock()
	if rr != nil {
		e.rrs[*rr] = struct{}{}
	}
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(qname string, qtype string) []*RR {
	e := r.getEntry(qname)
	if e == nil && r != Root {
		e = Root.getEntry(qname)
	}
	if e == nil {
		return nil
	}
	e.m.RLock()
	defer e.m.RUnlock()
	if len(e.rrs) == 0 {
		return []*RR{}
	}
	rrs := make([]*RR, 0, len(e.rrs))
	for rr, _ := range e.rrs {
		// fmt.Printf("%s\n", rr.String())
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, &RR{rr.Name, rr.Type, rr.Value})
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil
	}
	return rrs
}

// getEntry returns a single cache entry or nil if an entry does not exist in the cache.
func (r *Resolver) getEntry(qname string) *entry {
	c, ok := r.cache.Get(qname)
	if !ok {
		return nil
	}
	e, ok := c.(*entry)
	if !ok {
		return nil
	}
	return e
}
<|endoftext|>


<commit_before>package clock

import (
	"time"
)

type mockTicker struct {
	c    chan time.Time
	stop chan bool

	clock    Clock
	interval time.Duration
	start    time.Time
}

var _ Ticker = new(mockTicker)

// note: this probably does not function the same way as the time.Timer
// in the event that the clock skips more than the timer interval. I've
// not yet dug deep into the runtimeTimer to see how that works.
// PRs are appreciated!
func (m *mockTicker) wait() {
	for i := time.Duration(1); true; i++ {
		delta := m.start.Add(m.interval * i).Sub(m.clock.Now())

		select {
		case <-m.stop:
			return
		case <-m.clock.After(delta):
			select {
			case m.c <- m.clock.Now():
			case <-m.stop:
				return
			}
		}
	}
}

func (m *mockTicker) Chan() <-chan time.Time {
	return m.c
}

func (m *mockTicker) Stop() {
	m.stop <- true
}

// Creates a new Ticker using the provided Clock. You should not use this
// directly outside of unit tests; use Clock.NewTicker().
func NewMockTicker(c Clock, interval time.Duration) Ticker {
	t := &mockTicker{
		c:        make(chan time.Time),
		stop:     make(chan bool),
		interval: interval,
		start:    c.Now(),
		clock:    c,
	}
	go t.wait()

	return t
}
<commit_msg>fix :use ready chan to fix ticker race condition (#14)<commit_after>package clock

import (
	"time"
)

type mockTicker struct {
	c    chan time.Time
	stop chan bool

	clock    Clock
	interval time.Duration
	start    time.Time
}

var _ Ticker = new(mockTicker)

// note: this probably does not function the same way as the time.Timer
// in the event that the clock skips more than the timer interval. I've
// not yet dug deep into the runtimeTimer to see how that works.
// PRs are appreciated!
func (m *mockTicker) wait(ready chan<- struct{}) {
	for i := time.Duration(1); true; i++ {
		delta := m.start.Add(m.interval * i).Sub(m.clock.Now())
		afterChan := m.clock.After(delta)

		if i == time.Duration(1) {
			ready <- struct{}{}
		}

		select {
		case <-m.stop:
			return
		case <-afterChan:
			select {
			case m.c <- m.clock.Now():
			case <-m.stop:
				return
			}
		}
	}
}

func (m *mockTicker) Chan() <-chan time.Time {
	return m.c
}

func (m *mockTicker) Stop() {
	m.stop <- true
}

// Creates a new Ticker using the provided Clock. You should not use this
// directly outside of unit tests; use Clock.NewTicker().
func NewMockTicker(c Clock, interval time.Duration) Ticker {
	t := &mockTicker{
		c:        make(chan time.Time),
		stop:     make(chan bool),
		interval: interval,
		start:    c.Now(),
		clock:    c,
	}

	ready := make(chan struct{})
	go t.wait(ready)
	<-ready
	return t
}

<|endoftext|>


<commit_before>// Provides support for parsing magnet links.

package rain

import (
	"encoding/base32"
	"encoding/hex"
	"errors"
	"net/url"
	"strings"
)

type magnet struct {
	InfoHash [20]byte
	Name     string
	Trackers []string
}

func parseMagnet(s string) (*magnet, error) {
	u, err := url.Parse(s)
	if err != nil {
		return nil, err
	}

	params := u.Query()

	xts, ok := params["xt"]
	if !ok {
		return nil, errors.New("missing xt param")
	}
	if len(xts) == 0 {
		return nil, errors.New("empty xt param")
	}

	xt := xts[0]
	if !strings.HasPrefix(xt, "urn:btih:") {
		return nil, errors.New("invalid xt param: must start with \"urn:btih:\"")
	}
	xt = xt[9:]

	var magnet magnet

	magnet.InfoHash, err = infoHashString(xt)
	if err != nil {
		return nil, err
	}

	names := params["dn"]
	if len(names) != 0 {
		magnet.Name = names[0]
	}

	magnet.Trackers = params["tr"]

	return &magnet, nil
}

// infoHashString returns a new info hash value from a string.
// s must be 40 (hex encoded) or 32 (base32 encoded) characters, otherwise it returns error.
func infoHashString(s string) ([20]byte, error) {
	var ih [20]byte
	var b []byte
	var err error
	if len(s) == 40 {
		b, err = hex.DecodeString(s)
	} else if len(s) == 32 {
		b, err = base32.StdEncoding.DecodeString(s)
	} else {
		return ih, errors.New("info hash must be 32 or 40 characters")
	}
	if err != nil {
		return ih, err
	}
	copy(ih[:], b)
	return ih, nil
}
<commit_msg>check scheme while parsing magnet<commit_after>// Provides support for parsing magnet links.

package rain

import (
	"encoding/base32"
	"encoding/hex"
	"errors"
	"net/url"
	"strings"
)

type magnet struct {
	InfoHash [20]byte
	Name     string
	Trackers []string
}

func parseMagnet(s string) (*magnet, error) {
	u, err := url.Parse(s)
	if err != nil {
		return nil, err
	}

	if u.Scheme != "magnet" {
		return nil, errors.New("not a magnet link")
	}

	params := u.Query()

	xts, ok := params["xt"]
	if !ok {
		return nil, errors.New("missing xt param")
	}
	if len(xts) == 0 {
		return nil, errors.New("empty xt param")
	}

	xt := xts[0]
	if !strings.HasPrefix(xt, "urn:btih:") {
		return nil, errors.New("invalid xt param: must start with \"urn:btih:\"")
	}
	xt = xt[9:]

	var magnet magnet

	magnet.InfoHash, err = infoHashString(xt)
	if err != nil {
		return nil, err
	}

	names := params["dn"]
	if len(names) != 0 {
		magnet.Name = names[0]
	}

	magnet.Trackers = params["tr"]

	return &magnet, nil
}

// infoHashString returns a new info hash value from a string.
// s must be 40 (hex encoded) or 32 (base32 encoded) characters, otherwise it returns error.
func infoHashString(s string) ([20]byte, error) {
	var ih [20]byte
	var b []byte
	var err error
	if len(s) == 40 {
		b, err = hex.DecodeString(s)
	} else if len(s) == 32 {
		b, err = base32.StdEncoding.DecodeString(s)
	} else {
		return ih, errors.New("info hash must be 32 or 40 characters")
	}
	if err != nil {
		return ih, err
	}
	copy(ih[:], b)
	return ih, nil
}
<|endoftext|>


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

package main

import (
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/jroimartin/gocui"
)

const delta = 0.1

type HelpWidget struct {
	name string
	x, y int
	w, h int
	body string
}

func NewHelpWidget(name string, x, y int, body string) *HelpWidget {
	lines := strings.Split(body, "\n")

	w := 0
	for _, l := range lines {
		if len(l) > w {
			w = len(l)
		}
	}
	h := len(lines) + 1
	w = w + 1

	return &HelpWidget{name: name, x: x, y: y, w: w, h: h, body: body}
}

func (w *HelpWidget) Layout(g *gocui.Gui) error {
	v, err := g.SetView(w.name, w.x, w.y, w.x+w.w, w.y+w.h)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		fmt.Fprint(v, w.body)
	}
	return nil
}

type StatusbarWidget struct {
	name string
	x, y int
	w    int
	val  float32
}

func NewStatusbarWidget(name string, x, y, w int) *StatusbarWidget {
	return &StatusbarWidget{name: name, x: x, y: y, w: w}
}

func (w *StatusbarWidget) SetVal(val float32) error {
	if val < 0 || val > 1+delta/2 {
		return errors.New("invalid value")
	}
	w.val = val
	return nil
}

func (w *StatusbarWidget) Val() float32 {
	return w.val
}

func (w *StatusbarWidget) Layout(g *gocui.Gui) error {
	v, err := g.SetView(w.name, w.x, w.y, w.x+w.w, w.y+2)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
	}
	v.Clear()
	val := int(w.val * float32(w.w-1))
	fmt.Fprint(v, strings.Repeat("▒", val))
	return nil
}

type ButtonWidget struct {
	name    string
	x, y    int
	w       int
	label   string
	handler func(g *gocui.Gui, v *gocui.View) error
}

func NewButtonWidget(name string, x, y int, label string, handler func(g *gocui.Gui, v *gocui.View) error) *ButtonWidget {
	return &ButtonWidget{name: name, x: x, y: y, w: len(label) + 1, label: label, handler: handler}
}

func (w *ButtonWidget) Layout(g *gocui.Gui) error {
	v, err := g.SetView(w.name, w.x, w.y, w.x+w.w, w.y+2)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		if _, err := g.SetCurrentView(w.name); err != nil {
			return err
		}
		if err := g.SetKeybinding(w.name, gocui.KeyEnter, gocui.ModNone, w.handler); err != nil {
			return err
		}
		fmt.Fprint(v, w.label)
	}
	return nil
}

func main() {
	g, err := gocui.NewGui()
	if err != nil {
		log.Panicln(err)
	}
	defer g.Close()

	g.Highlight = true
	g.SelFgColor = gocui.ColorRed

	help := NewHelpWidget("help", 1, 1, helpText)
	status := NewStatusbarWidget("status", 1, 6, 50)
	butdown := NewButtonWidget("butdown", 52, 6, "DOWN", statusDown(status))
	butup := NewButtonWidget("butup", 58, 6, "UP", statusUp(status))
	g.SetManager(help, butdown, butup, status)

	if err := g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, quit); err != nil {
		log.Panicln(err)
	}
	if err := g.SetKeybinding("", gocui.KeyTab, gocui.ModNone, toggleButton); err != nil {
		log.Panicln(err)
	}

	if err := g.MainLoop(); err != nil && err != gocui.ErrQuit {
		log.Panicln(err)
	}
}

func quit(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func toggleButton(g *gocui.Gui, v *gocui.View) error {
	nextview := "butdown"
	if v == nil || v.Name() == "butdown" {
		nextview = "butup"
	}
	_, err := g.SetCurrentView(nextview)
	return err
}

func statusUp(status *StatusbarWidget) func(g *gocui.Gui, v *gocui.View) error {
	return func(g *gocui.Gui, v *gocui.View) error {
		return statusSet(status, delta)
	}
}

func statusDown(status *StatusbarWidget) func(g *gocui.Gui, v *gocui.View) error {
	return func(g *gocui.Gui, v *gocui.View) error {
		return statusSet(status, -delta)
	}
}

func statusSet(sw *StatusbarWidget, inc float32) error {
	val := sw.Val() + inc
	if val < 0 || val > 1+delta/2 {
		return nil
	}
	return sw.SetVal(val)
}

const helpText = `KEYBINDINGS
Tab: Move between buttons
^C: Exit`
<commit_msg>Minor change in _examples/widgets.go<commit_after>// Copyright 2014 The gocui Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/jroimartin/gocui"
)

const delta = 0.1

type HelpWidget struct {
	name string
	x, y int
	w, h int
	body string
}

func NewHelpWidget(name string, x, y int, body string) *HelpWidget {
	lines := strings.Split(body, "\n")

	w := 0
	for _, l := range lines {
		if len(l) > w {
			w = len(l)
		}
	}
	h := len(lines) + 1
	w = w + 1

	return &HelpWidget{name: name, x: x, y: y, w: w, h: h, body: body}
}

func (w *HelpWidget) Layout(g *gocui.Gui) error {
	v, err := g.SetView(w.name, w.x, w.y, w.x+w.w, w.y+w.h)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		fmt.Fprint(v, w.body)
	}
	return nil
}

type StatusbarWidget struct {
	name string
	x, y int
	w    int
	val  float32
}

func NewStatusbarWidget(name string, x, y, w int) *StatusbarWidget {
	return &StatusbarWidget{name: name, x: x, y: y, w: w}
}

func (w *StatusbarWidget) SetVal(val float32) error {
	if val < 0 || val > 1+delta/2 {
		return errors.New("invalid value")
	}
	w.val = val
	return nil
}

func (w *StatusbarWidget) Val() float32 {
	return w.val
}

func (w *StatusbarWidget) Layout(g *gocui.Gui) error {
	v, err := g.SetView(w.name, w.x, w.y, w.x+w.w, w.y+2)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
	}
	v.Clear()
	val := int(w.val * float32(w.w-1))
	fmt.Fprint(v, strings.Repeat("▒", val))
	return nil
}

type ButtonWidget struct {
	name    string
	x, y    int
	w       int
	label   string
	handler func(g *gocui.Gui, v *gocui.View) error
}

func NewButtonWidget(name string, x, y int, label string, handler func(g *gocui.Gui, v *gocui.View) error) *ButtonWidget {
	return &ButtonWidget{name: name, x: x, y: y, w: len(label) + 1, label: label, handler: handler}
}

func (w *ButtonWidget) Layout(g *gocui.Gui) error {
	v, err := g.SetView(w.name, w.x, w.y, w.x+w.w, w.y+2)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		if _, err := g.SetCurrentView(w.name); err != nil {
			return err
		}
		if err := g.SetKeybinding(w.name, gocui.KeyEnter, gocui.ModNone, w.handler); err != nil {
			return err
		}
		fmt.Fprint(v, w.label)
	}
	return nil
}

func main() {
	g, err := gocui.NewGui()
	if err != nil {
		log.Panicln(err)
	}
	defer g.Close()

	g.Highlight = true
	g.SelFgColor = gocui.ColorRed

	help := NewHelpWidget("help", 1, 1, helpText)
	status := NewStatusbarWidget("status", 1, 6, 50)
	butdown := NewButtonWidget("butdown", 52, 6, "DOWN", statusDown(status))
	butup := NewButtonWidget("butup", 58, 6, "UP", statusUp(status))
	g.SetManager(help, status, butdown, butup)

	if err := g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, quit); err != nil {
		log.Panicln(err)
	}
	if err := g.SetKeybinding("", gocui.KeyTab, gocui.ModNone, toggleButton); err != nil {
		log.Panicln(err)
	}

	if err := g.MainLoop(); err != nil && err != gocui.ErrQuit {
		log.Panicln(err)
	}
}

func quit(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func toggleButton(g *gocui.Gui, v *gocui.View) error {
	nextview := "butdown"
	if v == nil || v.Name() == "butdown" {
		nextview = "butup"
	}
	_, err := g.SetCurrentView(nextview)
	return err
}

func statusUp(status *StatusbarWidget) func(g *gocui.Gui, v *gocui.View) error {
	return func(g *gocui.Gui, v *gocui.View) error {
		return statusSet(status, delta)
	}
}

func statusDown(status *StatusbarWidget) func(g *gocui.Gui, v *gocui.View) error {
	return func(g *gocui.Gui, v *gocui.View) error {
		return statusSet(status, -delta)
	}
}

func statusSet(sw *StatusbarWidget, inc float32) error {
	val := sw.Val() + inc
	if val < 0 || val > 1+delta/2 {
		return nil
	}
	return sw.SetVal(val)
}

const helpText = `KEYBINDINGS
Tab: Move between buttons
^C: Exit`
<|endoftext|>


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

package gitlab

import (
	"fmt"
	"net/http"
)

// TopicsService handles communication with the topics related methods
// of the GitLab API.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/topics.html
type TopicsService struct {
	client *Client
}

// Topic represents a GitLab project topic.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/topics.html
type Topic struct {
	ID                 int    `json:"id"`
	Name               string `json:"name"`
	Description        string `json:"description"`
	TotalProjectsCount uint64 `json:"total_projects_count"`
	AvatarURL          string `json:"avatar_url"`
}

func (t Topic) String() string {
	return Stringify(t)
}

// ListTopicsOptions represents the available ListTopics() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#list-topics
type ListTopicsOptions struct {
	ListOptions
	Search *string `url:"search,omitempty" json:"search,omitempty"`
}

// ListTopics Returns a list of project topics in the GitLab instance ordered by
// number of associated projects.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#list-topics
func (s *TopicsService) ListTopics(opt *ListTopicsOptions, options ...RequestOptionFunc) ([]*Topic, *Response, error) {

	req, err := s.client.NewRequest(http.MethodGet, "topics", opt, options)
	if err != nil {
		return nil, nil, err
	}

	var t []*Topic
	resp, err := s.client.Do(req, &t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// GetTopic Get a project topic by ID. It returns 200 together
// with the topic information if the topic exists. It returns 404 if the tag does not exist.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#get-a-topic
func (s *TopicsService) GetTopic(tid int, options ...RequestOptionFunc) (*Topic, *Response, error) {

	u := fmt.Sprintf("topics/%d", tid)

	req, err := s.client.NewRequest(http.MethodGet, u, nil, options)
	if err != nil {
		return nil, nil, err
	}

	var t *Topic
	resp, err := s.client.Do(req, &t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// CreateTopicOptions represents the available CreateTopic() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ce/api/tags.html#create-a-new-tag
type CreateTopicOptions struct {
	Name        *string `url:"name,omitempty" json:"name,omitempty"`
	Description *string `url:"description,omitempty" json:"description,omitempty"`
	//	Avatar      *string `url:"avatar,omitempty" json:"avatar,omitempty"`
}

// CreateTopic creates a new project topic. Only available to administrators.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#create-a-project-topic
func (s *TopicsService) CreateTopic(opt *CreateTopicOptions, options ...RequestOptionFunc) (*Topic, *Response, error) {

	req, err := s.client.NewRequest(http.MethodPost, "topics", opt, options)
	if err != nil {
		return nil, nil, err
	}

	t := new(Topic)
	resp, err := s.client.Do(req, t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// UpdateTopicOptions represents the available CreateTopic() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ce/api/tags.html#create-a-new-tag
type UpdateTopicOptions struct {
	Name        *string `url:"name,omitempty" json:"name,omitempty"`
	Description *string `url:"description,omitempty" json:"description,omitempty"`
	// Avatar      *string `url:"avatar,omitempty" json:"avatar,omitempty"`
}

// UpdateTopic updates a project topic. Only available to administrators.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#update-a-project-topic
func (s *TopicsService) UpdateTopic(tid int, opt *UpdateTopicOptions, options ...RequestOptionFunc) (*Topic, *Response, error) {

	u := fmt.Sprintf("topics/%d", tid)

	req, err := s.client.NewRequest(http.MethodPut, u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	topic := new(Topic)
	resp, err := s.client.Do(req, topic)
	if err != nil {
		return nil, resp, err
	}

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

package gitlab

import (
	"fmt"
	"net/http"
)

// TopicsService handles communication with the topics related methods
// of the GitLab API.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/topics.html
type TopicsService struct {
	client *Client
}

// Topic represents a GitLab project topic.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/topics.html
type Topic struct {
	ID                 int    `json:"id"`
	Name               string `json:"name"`
	Description        string `json:"description"`
	TotalProjectsCount uint64 `json:"total_projects_count"`
	AvatarURL          string `json:"avatar_url"`
}

func (t Topic) String() string {
	return Stringify(t)
}

// ListTopicsOptions represents the available ListTopics() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#list-topics
type ListTopicsOptions struct {
	ListOptions
	Search *string `url:"search,omitempty" json:"search,omitempty"`
}

// ListTopics Returns a list of project topics in the GitLab instance ordered by
// number of associated projects.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#list-topics
func (s *TopicsService) ListTopics(opt *ListTopicsOptions, options ...RequestOptionFunc) ([]*Topic, *Response, error) {

	req, err := s.client.NewRequest(http.MethodGet, "topics", opt, options)
	if err != nil {
		return nil, nil, err
	}

	var t []*Topic
	resp, err := s.client.Do(req, &t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// GetTopic Get a project topic by ID. It returns 200 together
// with the topic information if the topic exists. It returns 404 if the tag does not exist.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#get-a-topic
func (s *TopicsService) GetTopic(tid interface{}, options ...RequestOptionFunc) (*Topic, *Response, error) {

	group, err := parseID(tid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("topics/%s", pathEscape(group))

	req, err := s.client.NewRequest(http.MethodGet, u, nil, options)
	if err != nil {
		return nil, nil, err
	}

	var t *Topic
	resp, err := s.client.Do(req, &t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// CreateTopicOptions represents the available CreateTopic() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ce/api/tags.html#create-a-new-tag
type CreateTopicOptions struct {
	Name        *string `url:"name,omitempty" json:"name,omitempty"`
	Description *string `url:"description,omitempty" json:"description,omitempty"`
	//	Avatar      *string `url:"avatar,omitempty" json:"avatar,omitempty"`
}

// CreateTopic creates a new project topic. Only available to administrators.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#create-a-project-topic
func (s *TopicsService) CreateTopic(opt *CreateTopicOptions, options ...RequestOptionFunc) (*Topic, *Response, error) {

	req, err := s.client.NewRequest(http.MethodPost, "topics", opt, options)
	if err != nil {
		return nil, nil, err
	}

	t := new(Topic)
	resp, err := s.client.Do(req, t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// UpdateTopicOptions represents the available CreateTopic() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ce/api/tags.html#create-a-new-tag
type UpdateTopicOptions struct {
	Name        *string `url:"name,omitempty" json:"name,omitempty"`
	Description *string `url:"description,omitempty" json:"description,omitempty"`
	// Avatar      *string `url:"avatar,omitempty" json:"avatar,omitempty"`
}

// UpdateTopic updates a project topic. Only available to administrators.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#update-a-project-topic
func (s *TopicsService) UpdateTopic(tid interface{}, opt *UpdateTopicOptions, options ...RequestOptionFunc) (*Topic, *Response, error) {

	topic, err := parseID(tid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("topics/%s", pathEscape(topic))

	req, err := s.client.NewRequest(http.MethodPut, u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	t := new(Topic)
	resp, err := s.client.Do(req, t)
	if err != nil {
		return nil, resp, err
	}

	return t, resp, err
}

// DeleteTopic deletes a project topic. Only available to administrators.
//
// GitLab API docs:
// https://docs.gitlab.com/ee/api/topics.html#update-a-project-topic
func (s *TopicsService) DeleteTopic(tid interface{}, options ...RequestOptionFunc) (*Response, error) {

	topic, err := parseID(tid)
	if err != nil {
		return nil, err
	}
	u := fmt.Sprintf("topics/%s", pathEscape(topic))

	req, err := s.client.NewRequest(http.MethodDelete, u, nil, options)
	if err != nil {
		return nil, err
	}

	return s.client.Do(req, nil)
}
<|endoftext|>


<commit_before>package tasks

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/MyHomeworkSpace/api-server/mit"
	"github.com/MyHomeworkSpace/api-server/util"

	"github.com/MyHomeworkSpace/api-server/config"
)

// some classes have weird times and aren't on the catalog, so we just give up on them
var skipClasses = []string{"15.284", "15.386"}

type catalogListing struct {
	ID         string `json:"id"`
	ShortTitle string `json:"short"`
	Title      string `json:"title"`

	OfferedFall   bool `bool:"fall"`
	OfferedIAP    bool `bool:"iap"`
	OfferedSpring bool `bool:"spring"`

	FallInstructors   string `json:"fallI"`
	SpringInstructors string `json:"springI"`
}

type subjectOffering struct {
	ID      string `json:"id"`
	Title   string `json:"title"`
	Section string `json:"section"`
	Term    string `json:"term"`

	Time  string `json:"time"`
	Place string `json:"place"`

	FacultyID   string `json:"facultyID"`
	FacultyName string `json:"facultyName"`

	IsFake   bool `json:"fake"`
	IsMaster bool `json:"master"`

	IsDesign     bool `json:"design"`
	IsLab        bool `json:"lab"`
	IsLecture    bool `json:"lecture"`
	IsRecitation bool `json:"recitation"`
}

// StartImportFromMIT begins an import of the given data from the MIT Data Warehouse.
func StartImportFromMIT(source string, db *sql.DB) error {
	if source != "catalog" && source != "coursews" && source != "offerings" {
		return errors.New("tasks: invalid parameter")
	}

	taskID := "mit_" + source
	taskName := "MIT Import - " + source

	go taskWatcher(taskID, taskName, importFromMIT, source, db)
	return nil
}

func importFromMIT(lastCompletion *time.Time, source string, db *sql.DB) (taskResponse, error) {
	mitConfig := config.GetCurrent().MIT
	params := url.Values{}

	params.Add("source", source)

	currentTerm := mit.GetCurrentTerm()
	params.Add("termCode", currentTerm.Code)
	params.Add("academicYear", currentTerm.Code[:4])

	// TODO: remove
	params.Add("lastUpdateDate", "2019-01-01")

	requestURL := mitConfig.DataProxyURL + "fetch?" + params.Encode()
	if source == "coursews" {
		// actually use the secret coursews API
		requestURL = "https://coursews.mit.edu/coursews/"
	}

	client := &http.Client{}
	request, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		return taskResponse{}, err
	}

	if source != "coursews" {
		request.Header.Add("X-MHS-Auth", mitConfig.ProxyToken)
	}

	response, err := client.Do(request)
	if err != nil {
		return taskResponse{}, err
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		bodyBytes, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return taskResponse{}, err
		}

		return taskResponse{}, fmt.Errorf(
			"tasks: MIT data server returned status code %d, body: '%s'",
			response.StatusCode,
			string(bodyBytes),
		)
	}

	tx, err := db.Begin()
	if err != nil {
		return taskResponse{}, err
	}

	rowsAffected := int64(0)

	if source == "catalog" {
		listings := []catalogListing{}
		err = json.NewDecoder(response.Body).Decode(&listings)
		if err != nil {
			return taskResponse{}, err
		}

		for _, listing := range listings {
			result, err := tx.Exec(
				`INSERT INTO
					mit_listings(id, shortTitle, title, offeredFall, offeredIAP, offeredSpring, fallInstructors, springInstructors)
					VALUES(?, ?, ?, ?, ?, ?, ?, ?)
				ON DUPLICATE KEY UPDATE
					id = VALUES(id),
					shortTitle = VALUES(shortTitle),
					title = VALUES(title),
					offeredFall = VALUES(offeredFall),
					offeredIAP = VALUES(offeredIAP),
					offeredSpring = VALUES(offeredSpring),
					fallInstructors = VALUES(fallInstructors),
					springInstructors = VALUES(springInstructors)
				`,
				listing.ID,
				listing.ShortTitle,
				listing.Title,
				listing.OfferedFall,
				listing.OfferedIAP,
				listing.OfferedSpring,
				listing.FallInstructors,
				listing.SpringInstructors,
			)
			if err != nil {
				return taskResponse{}, err
			}

			rowAffected, err := result.RowsAffected()
			if err != nil {
				return taskResponse{}, err
			}
			if rowAffected > 0 {
				rowsAffected++
			}
		}
	} else if source == "coursews" {
		// unfortunately, this API gives us data in a rather annoying format
		// so we cannot just parse them into a go struct
		// instead we have to use a hacky series of typecasts :(
		wsData := map[string]interface{}{}
		err = json.NewDecoder(response.Body).Decode(&wsData)
		if err != nil {
			return taskResponse{}, err
		}

		items := wsData["items"].([]interface{})

		// assume it's all the same term
		termInfo, err := mit.GetTermByCode(currentTerm.Code)
		if err != nil {
			return taskResponse{}, err
		}

		// first, clear out any data from a previous term
		_, err = tx.Exec("DELETE FROM mit_offerings WHERE term <> ?", currentTerm.Code)
		if err != nil {
			return taskResponse{}, err
		}

		for _, itemInterface := range items {
			item := itemInterface.(map[string]interface{})

			itemType := item["type"].(string)

			if itemType == "Class" {
				// we actually don't care about classes
				continue
			}

			if itemType != "LectureSession" && itemType != "LabSession" && itemType != "RecitationSession" {
				return taskResponse{}, fmt.Errorf("tasks: unknown coursews item type '%s'", itemType)
			}

			itemLabel := item["label"].(string)
			itemSectionOf := item["section-of"].(string)
			itemTimeAndPlace := item["timeAndPlace"].(string)

			if itemTimeAndPlace == "null null" {
				// this record tells us absolutely nothing
				// ignore it
				continue
			}

			// for some reason, the labels are the section + class number joined together
			// for example, section L01 of 6.003 has a label of "L016.003"
			// parse out the section ID
			sectionID := strings.Replace(itemLabel, itemSectionOf, "", -1)

			// in order to maximize suffering, the coursews api also combines the time and place fields
			// examples of this include "MW9.30-11 4-251" (easy)
			// or "TR9-11 (MEETS 4/7 TO 5/14) MEC-209" and "M EVE (6-8 PM) BOSTON PRE-REL" (why??)
			// the high quality algorithm to parse this is to break the string into spaces, and keep removing words until it works
			timeAndPlaceParts := strings.Split(itemTimeAndPlace, " ")
			currentTimeString := ""
			parsed := false
			for i := len(timeAndPlaceParts); i > 0; i-- {
				currentTimeString = ""
				for j := 0; j < i; j++ {
					if j != 0 {
						currentTimeString += " "
					}
					currentTimeString += timeAndPlaceParts[j]
				}

				// attempt
				_, err = mit.ParseTimeInfo(currentTimeString, termInfo)
				if err != nil {
					// oofie
					continue
				}

				// we survived!
				parsed = true
				break
			}

			if !parsed {
				return taskResponse{}, fmt.Errorf("tasks: failed to parse timeAndPlace '%s'", itemTimeAndPlace)
			}

			time := currentTimeString
			place := strings.TrimSpace(strings.Replace(itemTimeAndPlace, currentTimeString, "", -1))

			if place == "" {
				return taskResponse{}, fmt.Errorf("tasks: didn't get place from timeAndPlace '%s'", itemTimeAndPlace)
			}

			if place == "null" {
				place = ""
			}

			isDesign := false // design sections seem to not be included?
			isLab := (sectionID[0] == 'B')
			isLecture := (sectionID[0] == 'L')
			isRecitation := (sectionID[0] == 'R')

			// now, try to insert this new record
			// since we have very little info, we do NOT overwrite existing faculty/extra data if we have some
			result, err := tx.Exec(
				`INSERT INTO
					mit_offerings(id, title, section, term, time, place, facultyID, facultyName, isFake, isMaster, isDesign, isLab, isLecture, isRecitation)
					VALUES(?, '', ?, ?, ?, ?, '', '', 0, 0, ?, ?, ?, ?)
				ON DUPLICATE KEY UPDATE
					id = VALUES(id),
					section = VALUES(section),
					term = VALUES(term),
					time = VALUES(time),
					place = VALUES(place),
					isDesign = VALUES(isDesign),
					isLab = VALUES(isLab),
					isLecture = VALUES(isLecture),
					isRecitation = VALUES(isRecitation)`,
				itemSectionOf,
				sectionID,
				currentTerm.Code,
				time,
				place,
				isDesign,
				isLab,
				isLecture,
				isRecitation,
			)
			if err != nil {
				return taskResponse{}, err
			}

			rowAffected, err := result.RowsAffected()
			if err != nil {
				return taskResponse{}, err
			}
			if rowAffected > 0 {
				rowsAffected++
			}
		}
	} else if source == "offerings" {
		offerings := []subjectOffering{}
		err = json.NewDecoder(response.Body).Decode(&offerings)
		if err != nil {
			return taskResponse{}, err
		}

		// first, clear out any data from a previous term
		_, err = tx.Exec("DELETE FROM mit_offerings WHERE term <> ?", currentTerm.Code)
		if err != nil {
			return taskResponse{}, err
		}

		for _, offering := range offerings {
			if offering.Time == "" {
				continue
			}

			if util.StringSliceContains(skipClasses, offering.ID) {
				continue
			}

			termInfo, err := mit.GetTermByCode(offering.Term)
			if err != nil {
				return taskResponse{}, err
			}

			// check that we can parse the time info
			_, err = mit.ParseTimeInfo(offering.Time, termInfo)
			if err != nil {
				return taskResponse{}, fmt.Errorf("mit: failed to parse time of offering of %s (%s): %s", offering.ID, offering.Section, err.Error())
			}

			result, err := tx.Exec(
				`INSERT INTO
					mit_offerings(id, title, section, term, time, place, facultyID, facultyName, isFake, isMaster, isDesign, isLab, isLecture, isRecitation)
					VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
				ON DUPLICATE KEY UPDATE
					id = VALUES(id),
					title = VALUES(title),
					section = VALUES(section),
					term = VALUES(term),
					time = VALUES(time),
					place = VALUES(place),
					facultyID = VALUES(facultyID),
					facultyName = VALUES(facultyName),
					isFake = VALUES(isFake),
					isMaster = VALUES(isMaster),
					isDesign = VALUES(isDesign),
					isLab = VALUES(isLab),
					isLecture = VALUES(isLecture),
					isRecitation = VALUES(isRecitation)`,
				offering.ID,
				offering.Title,
				offering.Section,
				offering.Term,
				offering.Time,
				offering.Place,
				offering.FacultyID,
				offering.FacultyName,
				offering.IsFake,
				offering.IsMaster,
				offering.IsDesign,
				offering.IsLab,
				offering.IsLecture,
				offering.IsRecitation,
			)
			if err != nil {
				return taskResponse{}, err
			}

			rowAffected, err := result.RowsAffected()
			if err != nil {
				return taskResponse{}, err
			}
			if rowAffected > 0 {
				rowsAffected++
			}
		}
	}

	return taskResponse{
		RowsAffected: rowsAffected,
	}, tx.Commit()
}
<commit_msg>mit: add additional classes to import ignore list<commit_after>package tasks

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/MyHomeworkSpace/api-server/mit"
	"github.com/MyHomeworkSpace/api-server/util"

	"github.com/MyHomeworkSpace/api-server/config"
)

// some classes have weird times and aren't on the catalog, so we just give up on them
var skipClasses = []string{"15.003", "15.S64", "15.284", "15.386"}

type catalogListing struct {
	ID         string `json:"id"`
	ShortTitle string `json:"short"`
	Title      string `json:"title"`

	OfferedFall   bool `bool:"fall"`
	OfferedIAP    bool `bool:"iap"`
	OfferedSpring bool `bool:"spring"`

	FallInstructors   string `json:"fallI"`
	SpringInstructors string `json:"springI"`
}

type subjectOffering struct {
	ID      string `json:"id"`
	Title   string `json:"title"`
	Section string `json:"section"`
	Term    string `json:"term"`

	Time  string `json:"time"`
	Place string `json:"place"`

	FacultyID   string `json:"facultyID"`
	FacultyName string `json:"facultyName"`

	IsFake   bool `json:"fake"`
	IsMaster bool `json:"master"`

	IsDesign     bool `json:"design"`
	IsLab        bool `json:"lab"`
	IsLecture    bool `json:"lecture"`
	IsRecitation bool `json:"recitation"`
}

// StartImportFromMIT begins an import of the given data from the MIT Data Warehouse.
func StartImportFromMIT(source string, db *sql.DB) error {
	if source != "catalog" && source != "coursews" && source != "offerings" {
		return errors.New("tasks: invalid parameter")
	}

	taskID := "mit_" + source
	taskName := "MIT Import - " + source

	go taskWatcher(taskID, taskName, importFromMIT, source, db)
	return nil
}

func importFromMIT(lastCompletion *time.Time, source string, db *sql.DB) (taskResponse, error) {
	mitConfig := config.GetCurrent().MIT
	params := url.Values{}

	params.Add("source", source)

	currentTerm := mit.GetCurrentTerm()
	params.Add("termCode", currentTerm.Code)
	params.Add("academicYear", currentTerm.Code[:4])

	// TODO: remove
	params.Add("lastUpdateDate", "2019-01-01")

	requestURL := mitConfig.DataProxyURL + "fetch?" + params.Encode()
	if source == "coursews" {
		// actually use the secret coursews API
		requestURL = "https://coursews.mit.edu/coursews/"
	}

	client := &http.Client{}
	request, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		return taskResponse{}, err
	}

	if source != "coursews" {
		request.Header.Add("X-MHS-Auth", mitConfig.ProxyToken)
	}

	response, err := client.Do(request)
	if err != nil {
		return taskResponse{}, err
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		bodyBytes, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return taskResponse{}, err
		}

		return taskResponse{}, fmt.Errorf(
			"tasks: MIT data server returned status code %d, body: '%s'",
			response.StatusCode,
			string(bodyBytes),
		)
	}

	tx, err := db.Begin()
	if err != nil {
		return taskResponse{}, err
	}

	rowsAffected := int64(0)

	if source == "catalog" {
		listings := []catalogListing{}
		err = json.NewDecoder(response.Body).Decode(&listings)
		if err != nil {
			return taskResponse{}, err
		}

		for _, listing := range listings {
			result, err := tx.Exec(
				`INSERT INTO
					mit_listings(id, shortTitle, title, offeredFall, offeredIAP, offeredSpring, fallInstructors, springInstructors)
					VALUES(?, ?, ?, ?, ?, ?, ?, ?)
				ON DUPLICATE KEY UPDATE
					id = VALUES(id),
					shortTitle = VALUES(shortTitle),
					title = VALUES(title),
					offeredFall = VALUES(offeredFall),
					offeredIAP = VALUES(offeredIAP),
					offeredSpring = VALUES(offeredSpring),
					fallInstructors = VALUES(fallInstructors),
					springInstructors = VALUES(springInstructors)
				`,
				listing.ID,
				listing.ShortTitle,
				listing.Title,
				listing.OfferedFall,
				listing.OfferedIAP,
				listing.OfferedSpring,
				listing.FallInstructors,
				listing.SpringInstructors,
			)
			if err != nil {
				return taskResponse{}, err
			}

			rowAffected, err := result.RowsAffected()
			if err != nil {
				return taskResponse{}, err
			}
			if rowAffected > 0 {
				rowsAffected++
			}
		}
	} else if source == "coursews" {
		// unfortunately, this API gives us data in a rather annoying format
		// so we cannot just parse them into a go struct
		// instead we have to use a hacky series of typecasts :(
		wsData := map[string]interface{}{}
		err = json.NewDecoder(response.Body).Decode(&wsData)
		if err != nil {
			return taskResponse{}, err
		}

		items := wsData["items"].([]interface{})

		// assume it's all the same term
		termInfo, err := mit.GetTermByCode(currentTerm.Code)
		if err != nil {
			return taskResponse{}, err
		}

		// first, clear out any data from a previous term
		_, err = tx.Exec("DELETE FROM mit_offerings WHERE term <> ?", currentTerm.Code)
		if err != nil {
			return taskResponse{}, err
		}

		for _, itemInterface := range items {
			item := itemInterface.(map[string]interface{})

			itemType := item["type"].(string)

			if itemType == "Class" {
				// we actually don't care about classes
				continue
			}

			if itemType != "LectureSession" && itemType != "LabSession" && itemType != "RecitationSession" {
				return taskResponse{}, fmt.Errorf("tasks: unknown coursews item type '%s'", itemType)
			}

			itemLabel := item["label"].(string)
			itemSectionOf := item["section-of"].(string)
			itemTimeAndPlace := item["timeAndPlace"].(string)

			if itemTimeAndPlace == "null null" {
				// this record tells us absolutely nothing
				// ignore it
				continue
			}

			// for some reason, the labels are the section + class number joined together
			// for example, section L01 of 6.003 has a label of "L016.003"
			// parse out the section ID
			sectionID := strings.Replace(itemLabel, itemSectionOf, "", -1)

			// in order to maximize suffering, the coursews api also combines the time and place fields
			// examples of this include "MW9.30-11 4-251" (easy)
			// or "TR9-11 (MEETS 4/7 TO 5/14) MEC-209" and "M EVE (6-8 PM) BOSTON PRE-REL" (why??)
			// the high quality algorithm to parse this is to break the string into spaces, and keep removing words until it works
			timeAndPlaceParts := strings.Split(itemTimeAndPlace, " ")
			currentTimeString := ""
			parsed := false
			for i := len(timeAndPlaceParts); i > 0; i-- {
				currentTimeString = ""
				for j := 0; j < i; j++ {
					if j != 0 {
						currentTimeString += " "
					}
					currentTimeString += timeAndPlaceParts[j]
				}

				// attempt
				_, err = mit.ParseTimeInfo(currentTimeString, termInfo)
				if err != nil {
					// oofie
					continue
				}

				// we survived!
				parsed = true
				break
			}

			if !parsed {
				return taskResponse{}, fmt.Errorf("tasks: failed to parse timeAndPlace '%s'", itemTimeAndPlace)
			}

			time := currentTimeString
			place := strings.TrimSpace(strings.Replace(itemTimeAndPlace, currentTimeString, "", -1))

			if place == "" {
				return taskResponse{}, fmt.Errorf("tasks: didn't get place from timeAndPlace '%s'", itemTimeAndPlace)
			}

			if place == "null" {
				place = ""
			}

			isDesign := false // design sections seem to not be included?
			isLab := (sectionID[0] == 'B')
			isLecture := (sectionID[0] == 'L')
			isRecitation := (sectionID[0] == 'R')

			// now, try to insert this new record
			// since we have very little info, we do NOT overwrite existing faculty/extra data if we have some
			result, err := tx.Exec(
				`INSERT INTO
					mit_offerings(id, title, section, term, time, place, facultyID, facultyName, isFake, isMaster, isDesign, isLab, isLecture, isRecitation)
					VALUES(?, '', ?, ?, ?, ?, '', '', 0, 0, ?, ?, ?, ?)
				ON DUPLICATE KEY UPDATE
					id = VALUES(id),
					section = VALUES(section),
					term = VALUES(term),
					time = VALUES(time),
					place = VALUES(place),
					isDesign = VALUES(isDesign),
					isLab = VALUES(isLab),
					isLecture = VALUES(isLecture),
					isRecitation = VALUES(isRecitation)`,
				itemSectionOf,
				sectionID,
				currentTerm.Code,
				time,
				place,
				isDesign,
				isLab,
				isLecture,
				isRecitation,
			)
			if err != nil {
				return taskResponse{}, err
			}

			rowAffected, err := result.RowsAffected()
			if err != nil {
				return taskResponse{}, err
			}
			if rowAffected > 0 {
				rowsAffected++
			}
		}
	} else if source == "offerings" {
		offerings := []subjectOffering{}
		err = json.NewDecoder(response.Body).Decode(&offerings)
		if err != nil {
			return taskResponse{}, err
		}

		// first, clear out any data from a previous term
		_, err = tx.Exec("DELETE FROM mit_offerings WHERE term <> ?", currentTerm.Code)
		if err != nil {
			return taskResponse{}, err
		}

		for _, offering := range offerings {
			if offering.Time == "" {
				continue
			}

			if util.StringSliceContains(skipClasses, offering.ID) {
				continue
			}

			termInfo, err := mit.GetTermByCode(offering.Term)
			if err != nil {
				return taskResponse{}, err
			}

			// check that we can parse the time info
			_, err = mit.ParseTimeInfo(offering.Time, termInfo)
			if err != nil {
				return taskResponse{}, fmt.Errorf("mit: failed to parse time of offering of %s (%s): %s", offering.ID, offering.Section, err.Error())
			}

			result, err := tx.Exec(
				`INSERT INTO
					mit_offerings(id, title, section, term, time, place, facultyID, facultyName, isFake, isMaster, isDesign, isLab, isLecture, isRecitation)
					VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
				ON DUPLICATE KEY UPDATE
					id = VALUES(id),
					title = VALUES(title),
					section = VALUES(section),
					term = VALUES(term),
					time = VALUES(time),
					place = VALUES(place),
					facultyID = VALUES(facultyID),
					facultyName = VALUES(facultyName),
					isFake = VALUES(isFake),
					isMaster = VALUES(isMaster),
					isDesign = VALUES(isDesign),
					isLab = VALUES(isLab),
					isLecture = VALUES(isLecture),
					isRecitation = VALUES(isRecitation)`,
				offering.ID,
				offering.Title,
				offering.Section,
				offering.Term,
				offering.Time,
				offering.Place,
				offering.FacultyID,
				offering.FacultyName,
				offering.IsFake,
				offering.IsMaster,
				offering.IsDesign,
				offering.IsLab,
				offering.IsLecture,
				offering.IsRecitation,
			)
			if err != nil {
				return taskResponse{}, err
			}

			rowAffected, err := result.RowsAffected()
			if err != nil {
				return taskResponse{}, err
			}
			if rowAffected > 0 {
				rowsAffected++
			}
		}
	}

	return taskResponse{
		RowsAffected: rowsAffected,
	}, tx.Commit()
}
<|endoftext|>