From 2fccab172ffd8cc91dcd7c4f87cd58ea93e930aa Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 14 Aug 2026 12:10:33 +0800 Subject: [PATCH] test: add enterprise first-class function e2e --- ...est_first_class_function_enterprise_e2e.py | 181 ++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 python/python/tests/test_first_class_function_enterprise_e2e.py diff --git a/python/python/tests/test_first_class_function_enterprise_e2e.py b/python/python/tests/test_first_class_function_enterprise_e2e.py new file mode 100644 index 000000000..1be22b499 --- /dev/null +++ b/python/python/tests/test_first_class_function_enterprise_e2e.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from __future__ import annotations + +import pyarrow as pa + +from lancedb import udf + + +@udf( + inputs={"value": pa.int64()}, + output=pa.int64(), + python="3.12", + packages=["pyarrow==24.0.0"], + output_nullable=True, +) +def double_nullable(value): + if value is None: + return None + return value * 2 + + +def test_first_class_function_enterprise_lifecycle(): + import json + import os + import uuid + from datetime import timedelta + + import pytest + + import lancedb + from lancedb.exceptions import FunctionError + from lancedb.expr import col + + host = os.environ.get("LANCEDB_FCF_E2E_HOST") + if not host: + pytest.skip("LANCEDB_FCF_E2E_HOST is required for the live enterprise test") + + database_uri = os.environ.get("LANCEDB_FCF_E2E_DB_URI", "db://fcf-e2e-local") + api_key = os.environ.get("LANCEDB_FCF_E2E_API_KEY", "fake") + run_suffix = uuid.uuid4().hex[:12] + table_name = f"fcf_e2e_{run_suffix}" + function_name = f"fcf_e2e.double_{run_suffix}" + job_timeout = timedelta(minutes=5) + query_timeout = timedelta(seconds=30) + + def connect(): + return lancedb.connect( + database_uri, + api_key=api_key, + host_override=host, + ) + + setup_db = connect() + setup_db.create_table( + table_name, + data=pa.Table.from_pylist( + [ + {"row_id": 1, "value": 2}, + {"row_id": 2, "value": 5}, + {"row_id": 3, "value": None}, + ], + schema=pa.schema( + [ + pa.field("row_id", pa.int64(), nullable=False), + pa.field("value", pa.int64(), nullable=True), + ] + ), + ), + ) + + registration_job = setup_db.functions.register(function_name, double_nullable) + registration_job_id = registration_job.id + assert isinstance(registration_job_id, str) and registration_job_id + registered_function = registration_job.wait(timeout=job_timeout) + assert type(registered_function) is lancedb.Function + assert isinstance(registered_function.id, str) and registered_function.id + with pytest.raises(AttributeError): + registered_function.id = "mutated" + + catalog_reader = connect() + function_by_name = catalog_reader.functions.get(function_name) + function_by_id = catalog_reader.functions.get_by_id(registered_function.id) + expected_signature = ((("value", pa.int64()),), pa.int64(), True) + expected_identity = ( + registered_function.id, + *expected_signature, + ) + for function in (registered_function, function_by_name, function_by_id): + assert type(function) is lancedb.Function + assert ( + function.id, + function.parameters, + function.output_type, + function.output_nullable, + ) == expected_identity + + generated_column_table = catalog_reader.open_table(table_name) + generated_column_job = generated_column_table.add_generated_column( + "derived", + registered_function(value=col("value")), + ) + generated_column_job_id = generated_column_job.id + assert isinstance(generated_column_job_id, str) and generated_column_job_id + assert generated_column_job.wait(timeout=job_timeout) is None + + complete_reader = connect().open_table(table_name) + complete_status = complete_reader.generated_column_status("derived") + assert complete_status == "complete" + initial_rows = sorted( + complete_reader.search() + .select(["row_id", "value", "derived"]) + .limit(3) + .to_list(timeout=query_timeout), + key=lambda row: row["row_id"], + ) + assert initial_rows == [ + {"row_id": 1, "value": 2, "derived": 4}, + {"row_id": 2, "value": 5, "derived": 10}, + {"row_id": 3, "value": None, "derived": None}, + ] + + update_result = complete_reader.update( + where="row_id = 2", + values={"value": 7}, + ) + assert update_result.rows_updated == 1 + + incomplete_reader = connect().open_table(table_name) + incomplete_status = incomplete_reader.generated_column_status("derived") + assert incomplete_status == "incomplete" + with pytest.raises(FunctionError) as raised: + ( + incomplete_reader.search() + .select(["row_id", "derived"]) + .limit(3) + .to_list(timeout=query_timeout) + ) + assert raised.value.code == "generated_column_incomplete" + + refresh_job = incomplete_reader.refresh_generated_column("derived") + refresh_job_id = refresh_job.id + assert isinstance(refresh_job_id, str) and refresh_job_id + assert refresh_job.wait(timeout=job_timeout) is None + + refreshed_reader = connect().open_table(table_name) + refreshed_status = refreshed_reader.generated_column_status("derived") + assert refreshed_status == "complete" + final_rows = sorted( + refreshed_reader.search() + .select(["row_id", "value", "derived"]) + .limit(3) + .to_list(timeout=query_timeout), + key=lambda row: row["row_id"], + ) + assert final_rows == [ + {"row_id": 1, "value": 2, "derived": 4}, + {"row_id": 2, "value": 7, "derived": 14}, + {"row_id": 3, "value": None, "derived": None}, + ] + + evidence = { + "run_suffix": run_suffix, + "database": database_uri.removeprefix("db://"), + "table": table_name, + "function": function_name, + "function_id": registered_function.id, + "job_ids": { + "register": registration_job_id, + "add_generated_column": generated_column_job_id, + "refresh_generated_column": refresh_job_id, + }, + "status_transitions": [ + complete_status, + incomplete_status, + refreshed_status, + ], + "final_rows": final_rows, + } + print(json.dumps(evidence, sort_keys=True, separators=(",", ":")))