im_server/server/db/mysql/snowflake_primary_keys.go

//go:build mysql
// +build mysql

package mysql

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

	"github.com/jmoiron/sqlx"
	"github.com/tinode/chat/server/store"
)

type snowflakePrimaryKeySpec struct {
	table       string
	uniqueIndex string
	businessKey []string
}

var snowflakePrimaryKeySpecs = []snowflakePrimaryKeySpec{
	{table: "im_user_security", uniqueIndex: "uk_im_user_security_user", businessKey: []string{"user_id"}},
	{table: "im_kvmeta", uniqueIndex: "uk_im_kvmeta_key", businessKey: []string{"key"}},
	{table: "im_user_profiles", uniqueIndex: "uk_im_user_profiles_user", businessKey: []string{"user_id"}},
	{table: "im_user_privacy_settings", uniqueIndex: "uk_im_user_privacy_settings_user", businessKey: []string{"user_id"}},
	{table: "im_device_push_tokens", uniqueIndex: "uk_im_device_push_tokens_device", businessKey: []string{"device_id"}},
	{table: "im_friendships", uniqueIndex: "uk_im_friendships_pair", businessKey: []string{"user_id", "peer_id"}},
	{table: "im_user_blocks", uniqueIndex: "uk_im_user_blocks_pair", businessKey: []string{"blocker_id", "blocked_id"}},
	{table: "im_youth_mode_settings", uniqueIndex: "uk_im_youth_mode_settings_user_purpose", businessKey: []string{"user_id", "purpose"}},
	{table: "im_group_settings", uniqueIndex: "uk_im_group_settings_topic", businessKey: []string{"topic"}},
	{table: "im_group_announcement_reads", uniqueIndex: "uk_im_group_announcement_reads_announcement_user", businessKey: []string{"announcement_id", "user_id"}},
	{table: "im_group_mutes", uniqueIndex: "uk_im_group_mutes_topic_user", businessKey: []string{"topic", "user_id"}},
}

const snowflakeMigrationLock = "CONCAT('im_snowflake_pk:', LEFT(SHA2(DATABASE(), 256), 40))"

func nextSnowflakeID() (int64, error) {
	uid := store.Store.GetUid()
	if uid.IsZero() {
		return 0, errors.New("mysql: snowflake generator is not initialized")
	}
	id := store.DecodeUid(uid)
	if id <= 0 {
		return 0, errors.New("mysql: snowflake generator returned an invalid ID")
	}
	return id, nil
}

// ensureSnowflakePrimaryKeys upgrades legacy business-key primary keys in the
// selected tenant database. The former primary-key columns remain unique.
func ensureSnowflakePrimaryKeys(db *sqlx.DB) error {
	ctx := context.Background()
	conn, err := db.Connx(ctx)
	if err != nil {
		return err
	}
	defer conn.Close()

	var locked sql.NullInt64
	if err = conn.QueryRowContext(ctx,
		"SELECT GET_LOCK("+snowflakeMigrationLock+", 60)").Scan(&locked); err != nil {
		return err
	}
	if !locked.Valid || locked.Int64 != 1 {
		return errors.New("mysql: timed out waiting for snowflake primary-key migration lock")
	}
	defer conn.ExecContext(ctx, "SELECT RELEASE_LOCK("+snowflakeMigrationLock+")")

	for _, spec := range snowflakePrimaryKeySpecs {
		if err = migrateSnowflakePrimaryKey(ctx, conn, spec); err != nil {
			return fmt.Errorf("mysql: migrate %s primary key: %w", spec.table, err)
		}
	}
	return nil
}

func migrateSnowflakePrimaryKey(ctx context.Context, conn *sqlx.Conn, spec snowflakePrimaryKeySpec) error {
	var exists int
	if err := conn.QueryRowContext(ctx, `SELECT EXISTS(
		SELECT 1 FROM information_schema.tables WHERE table_schema=DATABASE() AND table_name=?)`, spec.table).Scan(&exists); err != nil {
		return err
	}
	if exists == 0 {
		return nil
	}

	if err := conn.QueryRowContext(ctx, `SELECT EXISTS(
		SELECT 1 FROM information_schema.columns WHERE table_schema=DATABASE() AND table_name=? AND column_name='id')`, spec.table).Scan(&exists); err != nil {
		return err
	}
	if exists == 0 {
		if _, err := conn.ExecContext(ctx, "ALTER TABLE "+quoteIdent(spec.table)+" ADD COLUMN id BIGINT NULL FIRST"); err != nil {
			return err
		}
	}

	for {
		if err := conn.QueryRowContext(ctx, "SELECT EXISTS(SELECT 1 FROM "+quoteIdent(spec.table)+" WHERE id IS NULL OR id=0 LIMIT 1)").Scan(&exists); err != nil {
			return err
		}
		if exists == 0 {
			break
		}
		id, err := nextSnowflakeID()
		if err != nil {
			return err
		}
		result, err := conn.ExecContext(ctx, "UPDATE "+quoteIdent(spec.table)+" SET id=? WHERE id IS NULL OR id=0 LIMIT 1", id)
		if err != nil {
			return err
		}
		if affected, err := result.RowsAffected(); err != nil || affected != 1 {
			if err != nil {
				return err
			}
			return errors.New("could not backfill exactly one row")
		}
	}

	if err := conn.QueryRowContext(ctx, `SELECT EXISTS(
		SELECT 1 FROM information_schema.statistics WHERE table_schema=DATABASE() AND table_name=? AND index_name=?)`,
		spec.table, spec.uniqueIndex).Scan(&exists); err != nil {
		return err
	}
	if exists == 0 {
		businessKey := spec.businessKey
		var hasTenantID int
		if err := conn.QueryRowContext(ctx, `SELECT EXISTS(
			SELECT 1 FROM information_schema.columns WHERE table_schema=DATABASE() AND table_name=? AND column_name='tenant_id')`,
			spec.table).Scan(&hasTenantID); err != nil {
			return err
		}
		if hasTenantID == 1 {
			businessKey = append([]string{"tenant_id"}, businessKey...)
		}
		columns := make([]string, len(businessKey))
		for i, column := range businessKey {
			columns[i] = quoteIdent(column)
		}
		statement := "CREATE UNIQUE INDEX " + quoteIdent(spec.uniqueIndex) + " ON " +
			quoteIdent(spec.table) + "(" + strings.Join(columns, ",") + ")"
		if _, err := conn.ExecContext(ctx, statement); err != nil {
			return err
		}
	}

	var primaryColumns sql.NullString
	if err := conn.QueryRowContext(ctx, `SELECT GROUP_CONCAT(column_name ORDER BY seq_in_index SEPARATOR ',')
		FROM information_schema.statistics
		WHERE table_schema=DATABASE() AND table_name=? AND index_name='PRIMARY'`, spec.table).Scan(&primaryColumns); err != nil {
		return err
	}
	if primaryColumns.Valid && primaryColumns.String == "id" {
		var dataType, nullable, extra string
		if err := conn.QueryRowContext(ctx, `SELECT data_type,is_nullable,extra
			FROM information_schema.columns
			WHERE table_schema=DATABASE() AND table_name=? AND column_name='id'`, spec.table).
			Scan(&dataType, &nullable, &extra); err != nil {
			return err
		}
		if strings.EqualFold(dataType, "bigint") && strings.EqualFold(nullable, "NO") &&
			!strings.Contains(strings.ToLower(extra), "auto_increment") {
			return nil
		}
		_, err := conn.ExecContext(ctx, "ALTER TABLE "+quoteIdent(spec.table)+" MODIFY COLUMN id BIGINT NOT NULL")
		return err
	}

	statement := "ALTER TABLE " + quoteIdent(spec.table) + " MODIFY COLUMN id BIGINT NOT NULL"
	if primaryColumns.Valid && primaryColumns.String != "" {
		statement += ", DROP PRIMARY KEY"
	}
	statement += ", ADD PRIMARY KEY(id)"
	_, err := conn.ExecContext(ctx, statement)
	return err
}