paimon-python/pypaimon/tests/system/consumers_table_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.

"""End-to-end tests for the ``$consumers`` system table."""

import os
import shutil
import tempfile
import unittest

from pypaimon import CatalogFactory, Schema
from pypaimon.consumer.consumer import Consumer
from pypaimon.schema.data_types import DataField
from pypaimon.table.system.consumers_table import ConsumersTable


def _read(table):
    rb = table.new_read_builder()
    return rb.new_read().to_arrow(rb.new_scan().plan().splits())


class ConsumersTableTest(unittest.TestCase):

    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix="consumers_sys_")
        warehouse = os.path.join(self.tmp, "warehouse")
        self.catalog = CatalogFactory.create({"warehouse": warehouse})
        self.catalog.create_database("db", False)
        fields = [DataField.from_dict({"id": 0, "name": "v", "type": "INT"})]
        self.catalog.create_table("db.t", Schema(fields=fields), False)

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def _record_consumers(self, mapping):
        manager = self.catalog.get_table("db.t").consumer_manager()
        for consumer_id, next_snapshot in mapping.items():
            manager.reset_consumer(consumer_id, Consumer(next_snapshot))

    def test_consumers_table_is_loaded_via_catalog(self):
        table = self.catalog.get_table("db.t$consumers")
        self.assertIsInstance(table, ConsumersTable)

    def test_schema_column_layout(self):
        table = self.catalog.get_table("db.t$consumers")
        row_type = table.row_type()
        self.assertEqual(["consumer_id", "next_snapshot_id"],
                         [f.name for f in row_type.fields])
        for field in row_type.fields:
            self.assertFalse(field.type.nullable,
                             "{} should be NOT NULL".format(field.name))
        self.assertEqual(["consumer_id"], table.primary_keys())

    def test_read_returns_every_consumer(self):
        self._record_consumers({"c1": 5, "c2": 8})
        arrow_table = _read(self.catalog.get_table("db.t$consumers"))
        self.assertEqual(["consumer_id", "next_snapshot_id"],
                         arrow_table.schema.names)
        rendered = dict(zip(
            arrow_table.column("consumer_id").to_pylist(),
            arrow_table.column("next_snapshot_id").to_pylist(),
        ))
        self.assertEqual({"c1": 5, "c2": 8}, rendered)

    def test_empty_when_no_consumers(self):
        arrow_table = _read(self.catalog.get_table("db.t$consumers"))
        self.assertEqual(["consumer_id", "next_snapshot_id"],
                         arrow_table.schema.names)
        self.assertEqual(0, arrow_table.num_rows)
        # The declared (string, int64) schema survives an empty table.
        self.assertEqual("string", str(arrow_table.schema.field("consumer_id").type))
        self.assertEqual("int64",
                         str(arrow_table.schema.field("next_snapshot_id").type))

    def test_read_with_projection_keeps_requested_column_only(self):
        self._record_consumers({"c1": 5})
        table = self.catalog.get_table("db.t$consumers")
        rb = table.new_read_builder().with_projection(["consumer_id"])
        arrow_table = rb.new_read().to_arrow(rb.new_scan().plan().splits())
        self.assertEqual(["consumer_id"], arrow_table.schema.names)
        self.assertEqual(["c1"], arrow_table.column("consumer_id").to_pylist())


if __name__ == "__main__":
    unittest.main()