<commit_before>package fixtures

import (
	"database/sql"
	"fmt"
	"strings"

	"gopkg.in/yaml.v2"
)

// Load processes a YAML fixture and inserts/updates the database accordingly
func Load(data []byte, db *sql.DB, driver string) error {
	// Unmarshal the YAML data into a []Row slice
	var rows []Row
	if err := yaml.Unmarshal(data, &rows); err != nil {
		return err
	}

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

	// Iterate over rows define in the fixture
	for _, row := range rows {
		// Load internat struct variables
		row.Init()

		// Run a SELECT query to find out if we need to insert or UPDATE
		selectQuery := fmt.Sprintf(
			"SELECT COUNT(*) FROM %s WHERE %s",
			row.Table,
			row.GetWhere(driver, 0),
		)
		var count int
		err = tx.QueryRow(selectQuery, row.GetPKValues()...).Scan(&count)
		if err != nil {
			tx.Rollback() // rollback the transaction
			return err
		}

		if count == 0 {
			// Primary key not found, let's run an INSERT query
			insertQuery := fmt.Sprintf(
				"INSERT INTO %s(%s) VALUES(%s)",
				row.Table,
				strings.Join(row.GetInsertColumns(), ", "),
				strings.Join(row.GetInsertPlaceholders(driver), ", "),
			)
			_, err := tx.Exec(insertQuery, row.GetInsertValues()...)
			if err != nil {
				tx.Rollback() // rollback the transaction
				return err
			}
		} else {
			// Primary key found, let's run UPDATE query
			updateQuery := fmt.Sprintf(
				"UPDATE %s SET %s WHERE %s",
				row.Table,
				strings.Join(row.GetUpdatePlaceholders(driver), ", "),
				row.GetWhere(driver, row.GetUpdateColumnsLength()),
			)
			values := append(row.GetUpdateValues(), row.GetPKValues()...)
			_, err := tx.Exec(updateQuery, values...)
			if err != nil {
				tx.Rollback() // rollback the transaction
				return err
			}<commit_msg>Added fix for Postgres primary key sequence.<commit_after>package fixtures

import (
	"database/sql"
	"fmt"
	"strings"

	"gopkg.in/yaml.v2"
)

// Load processes a YAML fixture and inserts/updates the database accordingly
func Load(data []byte, db *sql.DB, driver string) error {
	// Unmarshal the YAML data into a []Row slice
	var rows []Row
	if err := yaml.Unmarshal(data, &rows); err != nil {
		return err
	}

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

	// Iterate over rows define in the fixture
	for _, row := range rows {
		// Load internat struct variables
		row.Init()

		// Run a SELECT query to find out if we need to insert or UPDATE
		selectQuery := fmt.Sprintf(
			"SELECT COUNT(*) FROM %s WHERE %s",
			row.Table,
			row.GetWhere(driver, 0),
		)
		var count int
		err = tx.QueryRow(selectQuery, row.GetPKValues()...).Scan(&count)
		if err != nil {
			tx.Rollback() // rollback the transaction
			return err
		}

		if count == 0 {
			// Primary key not found, let's run an INSERT query
			insertQuery := fmt.Sprintf(
				"INSERT INTO %s(%s) VALUES(%s)",
				row.Table,
				strings.Join(row.GetInsertColumns(), ", "),
				strings.Join(row.GetInsertPlaceholders(driver), ", "),
			)
			_, err := tx.Exec(insertQuery, row.GetInsertValues()...)
			if err != nil {
				tx.Rollback() // rollback the transaction
				return err
			}
			if driver == postgresDriver && row.GetInsertColumns()[0] == "id" {
				// Fixed the primary ID sequence for Postgres
				query := fmt.Sprintf(
					"SELECT pg_catalog.setval(pg_get_serial_sequence('%s', 'id'), "+
						"(SELECT MAX(id) FROM %s));",
					row.Table,
					row.Table,
				)
				_, err := tx.Exec(query)
				if err != nil {
					tx.Rollback() // rollback the transaction
					return err
				}
			}
		} else {
			// Primary key found, let's run UPDATE query
			updateQuery := fmt.Sprintf(
				"UPDATE %s SET %s WHERE %s",
				row.Table,
				strings.Join(row.GetUpdatePlaceholders(driver), ", "),
				row.GetWhere(driver, row.GetUpdateColumnsLength()),
			)
			values := append(row.GetUpdateValues(), row.GetPKValues()...)
			_, err := tx.Exec(updateQuery, values...)
			if err != nil {
				tx.Rollback() // rollback the transaction
				return err
			}
{'fraction_non_alphanumeric': 0.1426448736998514, 'fraction_numerical': 0.001733531451213472, 'mean_word_length': 2.909970958373669, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5050478', 'n_tokens_mistral': 1572, 'n_tokens_neox': 1303, 'n_words': 398}