paimon-python/pypaimon/tests/native_commit_test.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

from unittest.mock import Mock, patch

import pyarrow as pa
import pytest

from pypaimon import CatalogFactory, Schema
from pypaimon.catalog.catalog_environment import CatalogEnvironment
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.options.options import Options
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.file_store_table import FileStoreTable
from pypaimon.write.commit_message import CommitMessage
from pypaimon.execution.native.write.native_commit import (
    create_native_commit,
    native_runtime_available,
    native_messages_supported,
)
from pypaimon.write.table_write import StreamTableWrite


pytestmark = pytest.mark.python_write


def requires_native(test):
    # Exercise the companion Rust release in the Native CI job.
    test = pytest.mark.native_plan(test)
    return pytest.mark.skipif(
        not native_runtime_available(), reason='pypaimon-rust not installed')(test)


def _table(tmp_path, mode='append', backend='filesystem', catalog=None):
    options = {'warehouse': str(tmp_path / 'warehouse')}
    if backend == 'jdbc':
        options.update({'metastore': 'jdbc', 'uri': 'jdbc:sqlite:' + str(tmp_path / 'catalog.db')})
    if catalog is None:
        catalog = CatalogFactory.create(options)
        catalog.create_database('default', True)
    options = {'file.format': 'parquet', 'commit.native.enabled': 'true'}
    if mode == 'pk':
        options['bucket'] = '1'
    elif mode == 'de':
        options.update({'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true'})
    catalog.create_table('default.t', Schema.from_pyarrow_schema(
        pa.schema([('id', pa.int64()), ('pt', pa.string())]),
        options=options, primary_keys=['id'] if mode == 'pk' else [],
        partition_keys=[] if mode in ('pk', 'unpartitioned') else ['pt']), False)
    table = catalog.get_table('default.t')
    if backend == 'path':
        table = FileStoreTable.from_path(table.table_path)
    if backend == 'jdbc':
        catalog.close()
    return table


def _prepare(builder, rows, identifier=None):
    writer = builder.new_write()
    try:
        writer.write_arrow(pa.Table.from_pylist(
            rows, schema=pa.schema([('id', pa.int64()), ('pt', pa.string())])))
        return writer.prepare_commit() if identifier is None else writer.prepare_commit(identifier)
    finally:
        writer.close()


def _rows(table):
    builder = table.new_read_builder()
    return sorted(builder.new_read().to_arrow(
        builder.new_scan().plan().splits()).to_pylist(), key=lambda row: row['id'])


def _must_not_fallback(commit, method='commit'):
    return patch.object(commit.file_store_commit, method, side_effect=AssertionError('Python fallback'))


def _seed(table):
    builder = table.copy({'commit.native.enabled': 'false'}).new_batch_write_builder()
    commit = builder.new_commit()
    try:
        commit.commit(_prepare(builder, [
            {'id': 1, 'pt': 'a'}, {'id': 2, 'pt': 'b'}, {'id': 3, 'pt': None}]))
    finally:
        commit.close()


@pytest.mark.python_commit
def test_native_commit_is_opt_in():
    assert not CoreOptions(Options({})).native_commit_enabled()
    assert CoreOptions(Options({'commit.native.enabled': 'true'})).native_commit_enabled()


@pytest.mark.parametrize('backend', ['filesystem', 'path', 'jdbc'])
def test_non_rest_catalog_keeps_python_commit(tmp_path, backend):
    table = _table(tmp_path, backend=backend)
    assert create_native_commit(table, 'job') is None
    builder = table.new_batch_write_builder()
    commit = builder.new_commit()
    try:
        commit.commit(_prepare(builder, [{'id': 1, 'pt': 'a'}]))
        assert _rows(table) == [{'id': 1, 'pt': 'a'}]
    finally:
        commit.close()


def test_overwrite_builder_api_is_batch_only(tmp_path):
    table = _table(tmp_path)
    batch = table.new_batch_write_builder()
    assert batch.overwrite({'pt': 'a'}) is batch
    assert batch.static_partition == {'pt': 'a'}
    stream = table.new_stream_write_builder()
    with pytest.raises(AttributeError):
        stream.overwrite({'pt': 'a'})
    with pytest.raises(TypeError):
        StreamTableWrite(table, stream.commit_user, {'pt': 'a'})


@requires_native
@pytest.mark.parametrize('mode', ['append', 'pk'])
def test_native_batch_roundtrip_preserves_identity(tmp_path, native_rest_catalog, mode):
    table = _table(tmp_path, mode, catalog=native_rest_catalog)
    builder = table.new_batch_write_builder()
    rows = [{'id': 1, 'pt': 'a'}, {'id': 2, 'pt': None}]
    messages = _prepare(builder, rows)
    commit = builder.new_commit()
    try:
        with _must_not_fallback(commit):
            commit.commit(messages)
        snapshot = table.snapshot_manager().get_latest_snapshot()
        assert snapshot.commit_user == builder.commit_user
        assert snapshot.commit_identifier == BATCH_COMMIT_IDENTIFIER
        assert _rows(table) == rows
        with pytest.raises(RuntimeError, match='one-time'):
            commit.commit(messages)
    finally:
        commit.close()


@requires_native
def test_replaced_rest_table_rejects_stale_commit(tmp_path, native_rest_catalog):
    table = _table(tmp_path, catalog=native_rest_catalog)
    table.catalog_environment.uuid = 'stale-table-id'
    assert create_native_commit(table, 'job') is None


@requires_native
def test_native_stream_reuses_commit_user_and_identifiers(tmp_path, native_rest_catalog):
    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_stream_write_builder()
    commit = builder.new_commit()
    try:
        with _must_not_fallback(commit):
            for identifier in (7, 8):
                commit.commit(_prepare(builder, [{'id': identifier, 'pt': None}], identifier), identifier)
                snapshot = table.snapshot_manager().get_latest_snapshot()
                assert snapshot.commit_user == builder.commit_user
                assert snapshot.commit_identifier == identifier
        assert _rows(table) == [{'id': 7, 'pt': None}, {'id': 8, 'pt': None}]
    finally:
        commit.close()


@requires_native
@pytest.mark.parametrize('mode,dynamic,spec', [
    ('append', True, {'pt': 'ignored'}),
    ('append', False, {'pt': 'a'}),
    ('pk', True, {}),
    ('unpartitioned', True, {}),
])
def test_native_overwrite_replaces_only_target_rows(
        tmp_path, native_rest_catalog, mode, dynamic, spec):
    table = _table(tmp_path, mode, catalog=native_rest_catalog).copy({
        'dynamic-partition-overwrite': str(dynamic).lower(), 'commit.user-prefix': 'python'})
    _seed(table)
    builder = table.new_batch_write_builder().overwrite(spec)
    commit = builder.new_commit()
    try:
        messages = _prepare(builder, [{'id': 4, 'pt': 'a'}])
        with _must_not_fallback(commit, 'overwrite'):
            commit.commit(messages)
        expected = [{'id': 4, 'pt': 'a'}]
        if table.partition_keys:
            expected = [{'id': 2, 'pt': 'b'}, {'id': 3, 'pt': None}] + expected
        assert _rows(table) == expected
        snapshot = table.snapshot_manager().get_latest_snapshot()
        assert snapshot.commit_user == builder.commit_user
        assert snapshot.commit_identifier == BATCH_COMMIT_IDENTIFIER
        assert snapshot.commit_kind == 'OVERWRITE'
        with pytest.raises(RuntimeError, match='one-time'):
            commit.commit(messages)
    finally:
        commit.close()


@requires_native
@pytest.mark.parametrize('mode,dynamic,spec,remaining', [
    ('append', True, {}, [1, 2, 3]),
    ('append', False, {'pt': 'a'}, [2, 3]),
    ('append', False, {'pt': None}, [1, 2]),
    ('append', False, {'pt': '__DEFAULT_PARTITION__'}, [1, 2]),
    ('append', False, {}, []),
    ('unpartitioned', True, {}, []),
    ('pk', True, {}, []),
])
def test_native_empty_overwrite_semantics(
        tmp_path, native_rest_catalog, mode, dynamic, spec, remaining):
    table = _table(tmp_path, mode, catalog=native_rest_catalog).copy({
        'dynamic-partition-overwrite': str(dynamic).lower()})
    _seed(table)
    builder = table.new_batch_write_builder().overwrite(spec)
    commit = builder.new_commit()
    try:
        with _must_not_fallback(commit, 'overwrite'):
            commit.commit([])
        assert [row['id'] for row in _rows(table)] == remaining
        snapshot = table.snapshot_manager().get_latest_snapshot()
        if remaining == [1, 2, 3]:
            assert snapshot.id == 1
        else:
            assert snapshot.id == 2
            assert snapshot.commit_user == builder.commit_user
            assert snapshot.commit_kind == 'OVERWRITE'
    finally:
        commit.close()


@requires_native
@pytest.mark.parametrize('value', ['off', '0', ' false '])
def test_native_overwrite_normalizes_python_boolean_option(tmp_path, native_rest_catalog, value):
    table = _table(tmp_path, catalog=native_rest_catalog).copy({
        'dynamic-partition-overwrite': value, 'snapshot.ignore-empty-commit': value})
    _seed(table)
    commit = table.new_batch_write_builder().overwrite({'pt': 'a'}).new_commit()
    try:
        with _must_not_fallback(commit, 'overwrite'):
            commit.commit([])
        assert [row['id'] for row in _rows(table)] == [2, 3]
    finally:
        commit.close()


@pytest.mark.parametrize('native', [False, pytest.param(True, marks=pytest.mark.native_plan)])
@pytest.mark.parametrize('case', ['unpartitioned', 'static-empty', 'static-missing', 'dynamic-empty'])
def test_empty_overwrite_records_java_snapshot(tmp_path, native_rest_catalog, native, case):
    if native and not native_runtime_available():
        pytest.skip('pypaimon-rust not installed')
    table = _table(tmp_path, 'unpartitioned' if case == 'unpartitioned' else 'append',
                   catalog=native_rest_catalog).copy({
                       'commit.native.enabled': str(native).lower(),
                       'dynamic-partition-overwrite': str(case == 'dynamic-empty').lower(),
                   })
    if case == 'static-missing':
        _seed(table)
    builder = table.new_batch_write_builder().overwrite(
        {'pt': 'missing'} if case.startswith('static') else {})
    commit = builder.new_commit()
    try:
        if native:
            with _must_not_fallback(commit, 'overwrite'):
                commit.commit([])
        else:
            commit.commit([])
        snapshot = table.snapshot_manager().get_latest_snapshot()
        if case == 'dynamic-empty':
            assert snapshot is None
        else:
            assert snapshot.id == (2 if case == 'static-missing' else 1)
            assert snapshot.commit_user == builder.commit_user
            assert snapshot.commit_kind == 'OVERWRITE'
            assert snapshot.total_record_count == (3 if case == 'static-missing' else 0)
        assert [row['id'] for row in _rows(table)] == ([1, 2, 3] if case == 'static-missing' else [])
    finally:
        commit.close()


@requires_native
@pytest.mark.parametrize('mode', ['batch', 'stream'])
@pytest.mark.parametrize('ignore', [True, False])
def test_native_empty_commit_preserves_python_option(
        tmp_path, native_rest_catalog, mode, ignore):
    table = _table(tmp_path, catalog=native_rest_catalog).copy({
        'snapshot.ignore-empty-commit': str(ignore).lower()})
    commit = getattr(table, 'new_' + mode + '_write_builder')().new_commit()
    try:
        with _must_not_fallback(commit):
            if mode == 'batch':
                commit.commit([])
            else:
                commit.commit([], 7)
        snapshot = table.snapshot_manager().get_latest_snapshot()
        if ignore:
            assert snapshot is None
        else:
            assert snapshot.commit_identifier == (BATCH_COMMIT_IDENTIFIER if mode == 'batch' else 7)
            assert snapshot.total_record_count == 0
    finally:
        commit.close()


@requires_native
@pytest.mark.parametrize('overwrite', [False, True])
def test_native_abort_removes_uncommitted_files(tmp_path, native_rest_catalog, overwrite):
    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_batch_write_builder()
    if overwrite:
        builder.overwrite()
    messages = _prepare(builder, [{'id': 1, 'pt': None}])
    files = list(tmp_path.rglob('data-*.parquet'))
    assert files
    commit = builder.new_commit()
    try:
        with patch.object(commit.file_store_commit, 'abort', side_effect=AssertionError('Python fallback')):
            commit.abort(messages)
        assert not any(path.exists() for path in files)
        assert table.snapshot_manager().get_latest_snapshot() is None
    finally:
        commit.close()


@requires_native
@pytest.mark.parametrize('serialized', [False, True])
def test_python_written_escaped_partition_supports_native_abort(
        tmp_path, native_rest_catalog, serialized):
    from pypaimon.write.commit_message_serializer import deserialize_commit_message, serialize_commit_message

    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_batch_write_builder()
    messages = _prepare(builder, [{'id': 1, 'pt': 'a/b'}])
    file = messages[0].new_files[0]
    assert table.file_io.exists(file.file_path)
    assert '/pt=a%2Fb/' in file.file_path
    if serialized:
        messages = [deserialize_commit_message(serialize_commit_message(message, table.partition_keys_fields),
                                               table.partition_keys_fields) for message in messages]
    assert native_messages_supported(table, messages)

    commit = builder.new_commit()
    try:
        with patch.object(commit.file_store_commit, 'abort',
                          side_effect=AssertionError('Python fallback')):
            commit.abort(messages)
        assert not table.file_io.exists(file.file_path)
        assert table.snapshot_manager().get_latest_snapshot() is None
    finally:
        commit.close()


@pytest.mark.parametrize('failure', ['missing', 'construction', 'conversion'])
def test_preflight_failure_uses_python(tmp_path, failure):
    table = _table(tmp_path)
    builder = table.new_batch_write_builder()
    messages = _prepare(builder, [{'id': 1, 'pt': 'a'}])
    commit = builder.new_commit()
    native = Mock()
    with patch('pypaimon.execution.native.write.native_commit.create_native_commit', return_value=native) as create:
        if failure == 'missing':
            create.return_value = None
        elif failure == 'construction':
            create.side_effect = RuntimeError('unsupported FileIO')
        with patch('pypaimon.execution.native.write.native_commit.to_native_commit_messages',
                   side_effect=ValueError('unsupported payload')):
            commit.commit(messages)
    native.commit.assert_not_called()
    assert _rows(table) == [{'id': 1, 'pt': 'a'}]
    commit.close()


@pytest.mark.parametrize('method', ['commit', 'abort'])
@pytest.mark.parametrize('overwrite', [False, True])
def test_native_mutation_failure_never_falls_back_or_aborts(tmp_path, method, overwrite):
    table = _table(tmp_path)
    builder = table.new_batch_write_builder()
    if overwrite:
        builder.overwrite()
    messages = _prepare(builder, [{'id': 1, 'pt': 'a'}])
    commit = builder.new_commit()
    native = Mock()
    getattr(native, method).side_effect = OSError('outcome unknown')
    with patch('pypaimon.execution.native.write.native_commit.create_native_commit', return_value=native), \
            patch('pypaimon.execution.native.write.native_commit.to_native_commit_messages', return_value=['wire']), \
            patch.object(commit.file_store_commit, 'commit') as fallback, \
            patch.object(commit.file_store_commit, 'overwrite') as fallback_overwrite, \
            patch.object(commit.file_store_commit, 'abort') as abort:
        with pytest.raises(OSError, match='outcome unknown'):
            getattr(commit, method)(messages)
        fallback.assert_not_called()
        fallback_overwrite.assert_not_called()
        abort.assert_not_called()
        if method == 'commit':
            native.abort.assert_not_called()
            with pytest.raises(RuntimeError, match='one-time'):
                commit.commit(messages)
    assert list(tmp_path.rglob('data-*.parquet'))
    commit.close()


@requires_native
@pytest.mark.parametrize('overwrite', [False, True])
def test_publication_response_loss_does_not_duplicate_or_delete_files(
        tmp_path, native_rest_catalog, overwrite):
    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_batch_write_builder()
    if overwrite:
        builder.overwrite()
    messages = _prepare(builder, [{'id': 1, 'pt': None}])
    native = create_native_commit(table, builder.commit_user, builder.static_partition)
    proxy = Mock(wraps=native)

    def publish_then_fail(*args):
        native.commit(*args)
        raise OSError('response lost')

    proxy.commit.side_effect = publish_then_fail
    commit = builder.new_commit()
    try:
        with patch('pypaimon.execution.native.write.native_commit.create_native_commit', return_value=proxy), \
                _must_not_fallback(commit, 'overwrite' if overwrite else 'commit'), \
                patch.object(commit.file_store_commit, 'abort', side_effect=AssertionError('abort')):
            with pytest.raises(OSError, match='response lost'):
                commit.commit(messages)
        assert table.snapshot_manager().get_latest_snapshot().id == 1
        assert _rows(table) == [{'id': 1, 'pt': None}]
        proxy.abort.assert_not_called()
    finally:
        commit.close()


@pytest.mark.parametrize('properties', [{}, {'source': 'python'}])
@pytest.mark.parametrize('overwrite', [False, True])
def test_non_rest_snapshot_properties_keep_python_commit(tmp_path, properties, overwrite):
    table = _table(tmp_path)
    builder = table.new_batch_write_builder()
    if overwrite:
        builder.overwrite()
    commit = builder.new_commit()
    with patch('pypaimon.execution.native.write.native_commit.create_native_commit',
               wraps=create_native_commit) as create:
        commit.commit(_prepare(builder, [{'id': 1, 'pt': None}]), snapshot_properties=properties)
        create.assert_called_once()
        assert commit._native_commit is None
    assert table.snapshot_manager().get_latest_snapshot().properties == (properties or None)
    commit.close()


@requires_native
@pytest.mark.parametrize('warmup', [False, True])
def test_callbacks_added_after_construction_use_native(
        tmp_path, native_rest_catalog, warmup):
    if warmup and not native_runtime_available():
        pytest.skip('native warmup requires the commit bindings')
    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_stream_write_builder()
    commit = builder.new_commit()
    if warmup:
        with _must_not_fallback(commit):
            commit.commit(_prepare(builder, [{'id': 6, 'pt': None}], 6), 6)
    callback = Mock()
    commit.add_commit_callback(callback)
    with _must_not_fallback(commit):
        commit.commit(_prepare(builder, [{'id': 1, 'pt': None}], 7), 7)
    callback.call.assert_called_once()
    context = callback.call.call_args[0][0]
    assert context.identifier == 7
    assert context.snapshot.id == (2 if warmup else 1)
    assert sum(entry.file.row_count for entry in context.commit_entries) == 1
    commit.close()
    callback.close.assert_called_once()


def test_truncate_keeps_python_when_native_publication_is_unavailable(tmp_path):
    table = _table(tmp_path)
    _seed(table)
    with patch('pypaimon.execution.native.write.native_commit.create_native_commit', return_value=None) as create:
        for method, args in [('truncate_partitions', ([{'pt': 'a'}],)), ('truncate_table', ())]:
            commit = table.new_batch_write_builder().new_commit()
            try:
                getattr(commit, method)(*args)
            finally:
                commit.close()
        assert create.call_count == 2
    assert _rows(table) == []


def test_overwrite_conversion_failure_preserves_partition_scope(tmp_path):
    table = _table(tmp_path).copy({'dynamic-partition-overwrite': 'false'})
    _seed(table)
    builder = table.new_batch_write_builder().overwrite({'pt': 'a'})
    commit = builder.new_commit()
    native = Mock()
    try:
        messages = _prepare(builder, [{'id': 4, 'pt': 'a'}])
        with patch('pypaimon.execution.native.write.native_commit.create_native_commit', return_value=native), \
                patch('pypaimon.execution.native.write.native_commit.to_native_commit_messages',
                      side_effect=ValueError('unsupported payload')):
            commit.commit(messages)
        native.commit.assert_not_called()
        native.abort.assert_not_called()
        assert _rows(table) == [{'id': 2, 'pt': 'b'}, {'id': 3, 'pt': None}, {'id': 4, 'pt': 'a'}]
        assert table.snapshot_manager().get_latest_snapshot().commit_user == builder.commit_user
    finally:
        commit.close()


def test_overwrite_callbacks_select_python(tmp_path):
    table = _table(tmp_path)
    builder = table.new_batch_write_builder().overwrite()
    commit = builder.new_commit()
    callback = Mock()
    commit.add_commit_callback(callback)
    with patch(
            'pypaimon.execution.native.write.native_commit.create_native_commit',
            wraps=create_native_commit) as create:
        commit.commit(_prepare(builder, [{'id': 1, 'pt': 'a'}]))
        create.assert_called_once()
        assert commit._native_commit is None
    assert _rows(table) == [{'id': 1, 'pt': 'a'}]
    callback.call.assert_called_once()
    commit.close()
    callback.close.assert_called_once()


def test_disabled_option_never_initializes_native(tmp_path):
    table = _table(tmp_path).copy({'commit.native.enabled': 'false'})
    builder = table.new_batch_write_builder()
    commit = builder.new_commit()
    with patch('pypaimon.execution.native.write.native_commit.create_native_commit') as create:
        commit.commit(_prepare(builder, [{'id': 1, 'pt': None}]))
        create.assert_not_called()
    assert _rows(table) == [{'id': 1, 'pt': None}]
    commit.close()


@pytest.mark.parametrize('kind', ['custom-env', 'custom-io'])
def test_incompatible_publication_environment_is_not_reconstructed(
        tmp_path, native_rest_catalog, kind):
    table = _table(tmp_path, catalog=native_rest_catalog)
    if kind == 'custom-env':
        class CustomEnvironment(CatalogEnvironment):
            pass
        table.catalog_environment = CustomEnvironment()
    else:
        table.file_io = Mock()
    with patch('pypaimon.execution.native.write.native_commit.native_runtime_available', return_value=True), \
            patch('pypaimon.execution.native.write.native_commit._native_rest_table') as resolve:
        assert create_native_commit(table, 'job') is None
        resolve.assert_not_called()


@pytest.mark.parametrize('option', ['hadoop_conf', 'prefer_io_loader', 'fallback_io_loader'])
def test_inactive_context_loader_fields_do_not_disable_native_commit(tmp_path, native_rest_catalog, option):
    from pypaimon.execution.native.write.native_commit import _rest_catalog_supported
    table = _table(tmp_path, catalog=native_rest_catalog)
    unused = Mock(side_effect=AssertionError('Inactive IO loader was invoked'))
    setattr(table.catalog_environment.catalog_loader.context(), option, unused)
    assert _rest_catalog_supported(table)
    unused.assert_not_called()


def test_missing_runtime_falls_back_without_reconstructing_table(
        tmp_path, native_rest_catalog):
    table = _table(tmp_path, catalog=native_rest_catalog)
    with patch('pypaimon.execution.native.write.native_commit.native_runtime_available', return_value=False), \
            patch('pypaimon.execution.native.write.native_commit._native_rest_table') as resolve:
        assert create_native_commit(table, 'job') is None
        resolve.assert_not_called()


def test_preassigned_row_id_messages_use_native_validation(
        tmp_path, native_rest_catalog):
    table = _table(tmp_path, 'de', catalog=native_rest_catalog)
    assert native_messages_supported(table, [CommitMessage((), 0, [])])
    assert native_messages_supported(table, [CommitMessage((), 0, [], check_from_snapshot=7)])
    messages = _prepare(table.new_batch_write_builder(), [{'id': 1, 'pt': 'a'}])
    messages[0].new_files[0].first_row_id = 1
    messages[0].check_from_snapshot = 7
    assert native_messages_supported(table, messages)
    assert not native_messages_supported(table, [CommitMessage((), 0, [], compact_after=[Mock()])])


@requires_native
@pytest.mark.parametrize('mode', ['append', 'de'])
def test_custom_manifest_target_uses_native_commit(tmp_path, native_rest_catalog, mode):
    from pypaimon.manifest.manifest_list_manager import ManifestListManager

    table = _table(tmp_path, mode, catalog=native_rest_catalog).copy({
        'manifest.target-file-size': '1 B'})
    builder = table.new_batch_write_builder()
    commit = builder.new_commit()
    try:
        messages = _prepare(builder, [{'id': 4, 'pt': 'a'}, {'id': 5, 'pt': 'b'}])
        with _must_not_fallback(commit):
            commit.commit(messages)
        assert _rows(table) == [{'id': 4, 'pt': 'a'}, {'id': 5, 'pt': 'b'}]
        snapshot = table.snapshot_manager().get_latest_snapshot()
        metas = ManifestListManager(table).read_delta(snapshot)
        assert len(metas) == 1  # Java checks only at multiples of 1000 records.
        assert metas[0].num_added_files == 2
    finally:
        commit.close()


def test_close_releases_python_resources_even_if_native_close_fails(tmp_path):
    commit = _table(tmp_path).new_batch_write_builder().new_commit()
    commit._native_commit = Mock()
    commit._native_commit.close.side_effect = OSError('close failure')
    with patch.object(commit.file_store_commit, 'close') as close:
        with pytest.raises(OSError, match='close failure'):
            commit.close()
        close.assert_called_once()


@requires_native
@pytest.mark.parametrize('mode', ['append', 'pk', 'de'])
@pytest.mark.parametrize('overwrite', [False, True])
def test_native_callbacks_keep_python_context_and_allow_native_read(
        tmp_path, native_rest_catalog, mode, overwrite):
    from pypaimon.manifest.schema.manifest_entry import ManifestEntry
    from pypaimon.table.row.generic_row import GenericRow
    from pypaimon.write.commit_callback import CommitCallback, CommitCallbackContext

    table = _table(tmp_path, mode, catalog=native_rest_catalog)
    table = table.copy({
        'scan.native-plan.enabled': 'true', 'read.native.enabled': 'true',
        'dynamic-partition-overwrite': 'false'})
    builder = table.new_batch_write_builder()
    if overwrite:
        _seed(table)
        builder.overwrite()
    calls = []

    class ReadingCallback(CommitCallback):
        def call(self, context):
            assert isinstance(context, CommitCallbackContext)
            assert all(isinstance(entry, ManifestEntry) for entry in context.commit_entries)
            assert all(isinstance(entry.partition, GenericRow) for entry in context.commit_entries)
            # Re-enter Rust's runtime from the Python callback after publication.
            with patch(
                'pypaimon.execution.python.read.table_read.PythonTableRead._create_split_read',
                side_effect=AssertionError('callback read fell back')
            ):
                calls.append((context, _rows(table)))

    commit = builder.new_commit()
    callback = ReadingCallback()
    follower = Mock()
    commit.add_commit_callback(callback)
    commit.add_commit_callback(follower)
    rows = [{'id': 9, 'pt': 'a/b%'}]
    messages = _prepare(builder, rows)
    paths = [file.file_path for message in messages for file in message.new_files]
    try:
        with _must_not_fallback(commit, 'overwrite' if overwrite else 'commit'):
            commit.commit(messages, snapshot_properties={'source': 'callback-test'})
        assert len(calls) == 1
        context, actual_rows = calls[0]
        assert actual_rows == rows
        assert context.identifier == BATCH_COMMIT_IDENTIFIER
        assert context.snapshot.properties == {'source': 'callback-test'}
        assert context.snapshot == table.snapshot_manager().get_latest_snapshot()
        adds = [entry for entry in context.commit_entries if entry.kind == 0]
        assert [entry.file.file_path for entry in adds] == paths
        assert all(table.file_io.exists(file_path) for file_path in paths)
        if mode == 'de':
            assert all(entry.file.first_row_id is not None for entry in adds)
            assert all(entry.file.min_sequence_number == context.snapshot.id for entry in adds)
        if overwrite:
            deletes = [entry for entry in context.commit_entries if entry.kind == 1]
            assert deletes
            assert all(entry.file.file_path for entry in deletes)
        follower.call.assert_called_once_with(context)
    finally:
        commit.close()
    follower.close.assert_called_once()


@requires_native
@pytest.mark.parametrize('lose_response', [False, True])
@pytest.mark.parametrize('callback_fails', [False, True])
def test_native_callback_recovery_and_error_preserve_files(
        tmp_path, native_rest_catalog, lose_response, callback_fails):
    from pypaimon.tests.rest.rest_server import RESTCatalogServer

    table = _table(tmp_path, 'de', catalog=native_rest_catalog)
    table = table.copy({'commit.min-retry-wait': '0 ms', 'commit.max-retry-wait': '0 ms'})
    builder = table.new_batch_write_builder()
    messages = _prepare(builder, [{'id': 1, 'pt': 'a/b%'}])
    paths = [file.file_path for message in messages for file in message.new_files]
    commit = builder.new_commit()
    callback = Mock()
    if callback_fails:
        callback.call.side_effect = LookupError('notification failed')
    commit.add_commit_callback(callback)
    follower = Mock()
    commit.add_commit_callback(follower)
    publish = RESTCatalogServer._table_commit_handle
    attempts = []

    def publish_then_lose_response(server, *args, **kwargs):
        result = publish(server, *args, **kwargs)
        attempts.append(result[1])
        assert result[1] == 200
        # A malformed successful response leaves publication outcome uncertain.
        return ('not-json', 200) if lose_response else result

    try:
        with patch.object(RESTCatalogServer, '_table_commit_handle', publish_then_lose_response), \
                _must_not_fallback(commit), \
                patch.object(commit.file_store_commit, 'abort', side_effect=AssertionError('abort')):
            if callback_fails:
                with pytest.raises(LookupError, match='notification failed'):
                    commit.commit(messages)
            else:
                commit.commit(messages)
        assert attempts == [200], 'notification failure must not republish the snapshot'
        context = callback.call.call_args[0][0]
        assert context.snapshot.id == 1
        assert context.snapshot.next_row_id == 1
        assert context.commit_entries[0].file.first_row_id == 0
        assert context.commit_entries[0].file.file_path == paths[0]
        assert tuple(context.commit_entries[0].partition.values) == ('a/b%',)
        assert _rows(table) == [{'id': 1, 'pt': 'a/b%'}]
        assert all(table.file_io.exists(file_path) for file_path in paths)
        callback.call.assert_called_once()
        if callback_fails:
            follower.call.assert_not_called()
        else:
            follower.call.assert_called_once_with(context)
    finally:
        commit.close()
    callback.close.assert_called_once()
    follower.close.assert_called_once()


@requires_native
def test_native_callback_empty_commit_and_prepublication_failure(tmp_path, native_rest_catalog):
    table = _table(tmp_path, catalog=native_rest_catalog).copy({'snapshot.ignore-empty-commit': 'false'})
    builder = table.new_stream_write_builder()
    commit = builder.new_commit()
    callback = Mock()
    commit.add_commit_callback(callback)
    try:
        with _must_not_fallback(commit):
            commit.commit([], 1)
        context = callback.call.call_args[0][0]
        assert context.snapshot.id == 1
        assert context.identifier == 1
        assert context.commit_entries == []
        bad = _prepare(builder, [{'id': 1, 'pt': None}], 2)
        # A repeated file entry must fail before publication and notification.
        with _must_not_fallback(commit), pytest.raises(ValueError, match='Duplicate'):
            commit.commit(bad + bad, 2)
        callback.call.assert_called_once()
        assert table.snapshot_manager().get_latest_snapshot().id == 1
    finally:
        commit.close()


@requires_native
def test_native_callbacks_added_between_commits_share_context_and_close_once(
        tmp_path, native_rest_catalog):
    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_stream_write_builder()
    commit = builder.new_commit()
    first, second = Mock(), Mock()
    events = []

    def first_call(context):
        events.append(('first', context.identifier, context))
        # Preserve Python's shared Context contract between registered callbacks.
        context.observed = True

    def second_call(context):
        assert context.observed
        events.append(('second', context.identifier, context))

    first.call.side_effect = first_call
    second.call.side_effect = second_call
    first.close.side_effect = OSError('close failed')
    commit.add_commit_callback(first)
    try:
        with _must_not_fallback(commit):
            commit.commit(_prepare(builder, [{'id': 1, 'pt': None}], 1), 1)
            commit.add_commit_callback(second)
            commit.commit(_prepare(builder, [{'id': 2, 'pt': None}], 2), 2)
            commit.commit(_prepare(builder, [{'id': 3, 'pt': None}], 3), 3)
        assert [(name, identifier) for name, identifier, _ in events] == [
            ('first', 1), ('first', 2), ('second', 2), ('first', 3), ('second', 3)]
        assert events[1][2] is events[2][2]
        assert events[3][2] is events[4][2]
    finally:
        commit.close()
    first.close.assert_called_once()
    second.close.assert_called_once()


@requires_native
@pytest.mark.parametrize('close', [False, True])
def test_native_commit_callback_cycle_is_collected(tmp_path, native_rest_catalog, close):
    import gc
    import weakref

    table = _table(tmp_path, catalog=native_rest_catalog)
    builder = table.new_stream_write_builder()
    commit = builder.new_commit()
    callback = Mock()
    callback.owner = commit
    reference = weakref.ref(callback)
    commit.add_commit_callback(callback)
    with _must_not_fallback(commit):
        commit.commit(_prepare(builder, [{'id': 1, 'pt': None}], 1), 1)
    if close:
        commit.close()
    del callback, commit
    gc.collect()
    assert reference() is None