# frozen_string_literal: true

module Gitlab
  module Database
    class TablesLocker
      GITLAB_SCHEMAS_TO_IGNORE = %i[gitlab_embedding gitlab_geo gitlab_jh].freeze

      def initialize(logger: nil, dry_run: false, include_partitions: true, options: {})
        @logger = logger
        @dry_run = dry_run
        @result = []
        @include_partitions = include_partitions
        @scope_to_database = options[:scope_to_database].to_s

        validate_scopes
      end

      def unlock_writes
        Gitlab::Database::EachDatabase.each_connection do |connection, database_name|
          next if skip_connection?(database_name)

          tables_to_lock(connection) do |table_name, schema_name|
            # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
            next if schema_name.in? GITLAB_SCHEMAS_TO_IGNORE

            unlock_writes_on_table(table_name, connection, database_name)
          end
        end

        @result
      end

      # It locks the tables on the database where they don't belong. Also it unlocks the tables
      # on the database where they belong
      def lock_writes
        Gitlab::Database::EachDatabase.each_connection(include_shared: false) do |connection, database_name|
          next if skip_connection?(database_name)

          schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)

          tables_to_lock(connection) do |table_name, schema_name|
            # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
            next if schema_name.in? GITLAB_SCHEMAS_TO_IGNORE

            if schemas_for_connection.include?(schema_name)
              unlock_writes_on_table(table_name, connection, database_name)
            else
              lock_writes_on_table(table_name, connection, database_name)
            end
          end
        end

        @result
      end

      private

      def skip_connection?(database_name)
        return false if @scope_to_database.empty? || database_name == @scope_to_database

        true
      end

      def validate_scopes
        return if @scope_to_database.blank? ||
          Gitlab::Database.all_database_connections.key?(@scope_to_database)

        raise "#{@scope_to_database} is not a valid database to scope to."
      end

      # Unlocks the writes on the table and its partitions
      def unlock_writes_on_table(table_name, connection, database_name)
        @result << lock_writes_manager(table_name, connection, database_name).unlock_writes
        return unless @include_partitions

        table_attached_partitions(table_name, connection) do |postgres_partition|
          @result << lock_writes_manager(postgres_partition.identifier, connection, database_name).unlock_writes
        end
      end

      # It locks the writes on the table and its partitions
      def lock_writes_on_table(table_name, connection, database_name)
        @result << lock_writes_manager(table_name, connection, database_name).lock_writes
        return unless @include_partitions

        table_attached_partitions(table_name, connection) do |postgres_partition|
          @result << lock_writes_manager(postgres_partition.identifier, connection, database_name).lock_writes
        end
      end

      def tables_to_lock(connection, &block)
        Gitlab::Database::GitlabSchema.tables_to_schema.each(&block)
        return unless @include_partitions

        Gitlab::Database::SharedModel.using_connection(connection) do
          Postgresql::DetachedPartition.find_each do |detached_partition|
            yield detached_partition.fully_qualified_table_name, detached_partition.table_schema
          end
        end
      end

      def table_attached_partitions(table_name, connection, &block)
        Gitlab::Database::SharedModel.using_connection(connection) do
          break unless Gitlab::Database::PostgresPartitionedTable.find_by_name_in_current_schema(table_name)

          Gitlab::Database::PostgresPartitionedTable.each_partition(table_name, &block)
        end
      end

      def lock_writes_manager(table_name, connection, database_name)
        Gitlab::Database::LockWritesManager.new(
          table_name: table_name,
          connection: connection,
          database_name: database_name,
          with_retries: true,
          logger: @logger,
          dry_run: @dry_run,
          force: !@dry_run
        )
      end
    end
  end
end
