clair/database/pgsql/pgsql_test.go

// Copyright 2016 clair 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 pgsql

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"github.com/pborman/uuid"
	log "github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	yaml "gopkg.in/yaml.v2"

	"github.com/coreos/clair/database"
	"github.com/coreos/clair/pkg/pagination"
)

var (
	withFixtureName, withoutFixtureName string
)

var testPaginationKey = pagination.Must(pagination.NewKey())

func genTemplateDatabase(name string, loadFixture bool) (sourceURL string, dbName string) {
	config := generateTestConfig(name, loadFixture, false)
	source := config.Options["source"].(string)
	name, url, err := parseConnectionString(source)
	if err != nil {
		panic(err)
	}

	fixturePath := config.Options["fixturepath"].(string)

	if err := createDatabase(url, name); err != nil {
		panic(err)
	}

	// migration and fixture
	db, err := sql.Open("postgres", source)
	if err != nil {
		panic(err)
	}

	// Verify database state.
	if err := db.Ping(); err != nil {
		panic(err)
	}

	// Run migrations.
	if err := migrateDatabase(db); err != nil {
		panic(err)
	}

	if loadFixture {
		d, err := ioutil.ReadFile(fixturePath)
		if err != nil {
			panic(err)
		}

		_, err = db.Exec(string(d))
		if err != nil {
			panic(err)
		}
	}

	db.Exec("UPDATE pg_database SET datistemplate=True WHERE datname=$1", name)
	db.Close()

	log.Info("Generated Template database ", name)
	return url, name
}

func dropTemplateDatabase(url string, name string) {
	db, err := sql.Open("postgres", url)
	if err != nil {
		panic(err)
	}

	if _, err := db.Exec("UPDATE pg_database SET datistemplate=False WHERE datname=$1", name); err != nil {
		panic(err)
	}

	if err := db.Close(); err != nil {
		panic(err)
	}

	if err := dropDatabase(url, name); err != nil {
		panic(err)
	}

}

func TestMain(m *testing.M) {
	fURL, fName := genTemplateDatabase("fixture", true)
	nfURL, nfName := genTemplateDatabase("nonfixture", false)

	withFixtureName = fName
	withoutFixtureName = nfName

	rc := m.Run()

	dropTemplateDatabase(fURL, fName)
	dropTemplateDatabase(nfURL, nfName)
	os.Exit(rc)
}

func openCopiedDatabase(testConfig database.RegistrableComponentConfig, fixture bool) (database.Datastore, error) {
	var fixtureName string
	if fixture {
		fixtureName = withFixtureName
	} else {
		fixtureName = withoutFixtureName
	}

	// copy the database into new database
	var pg pgSQL
	// Parse configuration.
	pg.config = Config{
		CacheSize: 16384,
	}

	bytes, err := yaml.Marshal(testConfig.Options)
	if err != nil {
		return nil, fmt.Errorf("pgsql: could not load configuration: %v", err)
	}
	err = yaml.Unmarshal(bytes, &pg.config)
	if err != nil {
		return nil, fmt.Errorf("pgsql: could not load configuration: %v", err)
	}

	dbName, pgSourceURL, err := parseConnectionString(pg.config.Source)
	if err != nil {
		return nil, err
	}

	// Create database.
	if pg.config.ManageDatabaseLifecycle {
		if err = copyDatabase(pgSourceURL, dbName, fixtureName); err != nil {
			return nil, err
		}
	}

	// Open database.
	pg.DB, err = sql.Open("postgres", pg.config.Source)
	fmt.Println("database", pg.config.Source)
	if err != nil {
		pg.Close()
		return nil, fmt.Errorf("pgsql: could not open database: %v", err)
	}

	return &pg, nil
}

// copyDatabase creates a new database with
func copyDatabase(url, name string, templateName string) error {
	// Open database.
	db, err := sql.Open("postgres", url)
	if err != nil {
		return fmt.Errorf("pgsql: could not open 'postgres' database for creation: %v", err)
	}
	defer db.Close()

	// Create database with copy
	_, err = db.Exec("CREATE DATABASE " + name + " WITH TEMPLATE " + templateName)
	if err != nil {
		return fmt.Errorf("pgsql: could not create database: %v", err)
	}

	return nil
}

func openDatabaseForTest(testName string, loadFixture bool) (*pgSQL, error) {
	var (
		db         database.Datastore
		err        error
		testConfig = generateTestConfig(testName, loadFixture, true)
	)

	db, err = openCopiedDatabase(testConfig, loadFixture)

	if err != nil {
		return nil, err
	}
	datastore := db.(*pgSQL)
	return datastore, nil
}

func generateTestConfig(testName string, loadFixture bool, manageLife bool) database.RegistrableComponentConfig {
	dbName := "test_" + strings.ToLower(testName) + "_" + strings.Replace(uuid.New(), "-", "_", -1)

	var fixturePath string
	if loadFixture {
		_, filename, _, _ := runtime.Caller(0)
		fixturePath = filepath.Join(filepath.Dir(filename)) + "/testdata/data.sql"
	}

	source := fmt.Sprintf("postgresql://postgres@127.0.0.1:5432/%s?sslmode=disable", dbName)
	if sourceEnv := os.Getenv("CLAIR_TEST_PGSQL"); sourceEnv != "" {
		source = fmt.Sprintf(sourceEnv, dbName)
	}

	return database.RegistrableComponentConfig{
		Options: map[string]interface{}{
			"source":                  source,
			"cachesize":               0,
			"managedatabaselifecycle": manageLife,
			"fixturepath":             fixturePath,
			"paginationkey":           testPaginationKey.String(),
		},
	}
}

func closeTest(t *testing.T, store database.Datastore, session database.Session) {
	err := session.Rollback()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}

	store.Close()
}

func openSessionForTest(t *testing.T, name string, loadFixture bool) (*pgSQL, *pgSession) {
	store, err := openDatabaseForTest(name, loadFixture)
	if err != nil {
		t.Error(err)
		t.FailNow()
	}
	tx, err := store.Begin()
	if err != nil {
		t.Error(err)
		t.FailNow()
	}

	return store, tx.(*pgSession)
}

func restartSession(t *testing.T, datastore *pgSQL, tx *pgSession, commit bool) *pgSession {
	var err error
	if !commit {
		err = tx.Rollback()
	} else {
		err = tx.Commit()
	}

	if assert.Nil(t, err) {
		session, err := datastore.Begin()
		if assert.Nil(t, err) {
			return session.(*pgSession)
		}
	}
	t.FailNow()
	return nil
}