diff --git a/.github/workflows/spark-search.yml b/.github/workflows/spark-search.yml index 7d416c700..63dc7b5ee 100644 --- a/.github/workflows/spark-search.yml +++ b/.github/workflows/spark-search.yml @@ -24,7 +24,10 @@ on: - "Makefile" - "docker/**" - "integration-tests/**" + - "lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java" - "lance-spark-base_2.12/src/main/java/org/lance/spark/search/**" + - "lance-spark-base_2.12/src/main/java/org/lance/spark/write/*ColumnsBackfill*.java" + - "lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/*ColumnsBackfillExec.scala" - "lance-spark-base_2.12/src/main/scala/org/lance/spark/search/**" - "lance-spark-base_2.12/src/test/java/org/lance/spark/search/**" - "lance-spark-*/src/main/scala/org/lance/spark/extensions/**" @@ -71,6 +74,11 @@ env: SEARCH_PYTEST_CMD: >- pytest /home/lance/tests/test_lance_spark.py::TestDQLSearchTableFunctions -v --timeout=180 + MANAGED_BACKFILL_PYTEST_CMD: >- + pytest + /home/lance/tests/test_lance_spark.py::TestDMLAddColumn::test_add_column_from_view_on_rest + /home/lance/tests/test_lance_spark.py::TestDMLUpdateColumn::test_update_column_from_view_on_rest + -v --timeout=180 jobs: search-docker-test: @@ -153,3 +161,16 @@ jobs: LANCE_SPARK_REST_DIR_PORT="${LANCE_SPARK_REST_DIR_PORT}" \ DOCKER_RUN_ARGS="${DOCKER_RUN_ARGS}" \ PYTEST_CMD="${SEARCH_PYTEST_CMD}" + # Keep managed versioning isolated from the existing REST search test server. + - name: Run managed REST directory namespace backfill tests + if: ${{ steps.rest.outputs.start_rest_dir == 'true' }} + run: | + make docker-test \ + SPARK_VERSION=${SPARK_VERSION} \ + SCALA_VERSION=${SCALA_VERSION} \ + TEST_BACKENDS=rest-dir \ + LANCE_SPARK_START_REST_DIR=1 \ + LANCE_SPARK_REST_DIR_ROOT=/home/lance/rest-managed-data \ + LANCE_SPARK_REST_DIR_PORT=10025 \ + LANCE_SPARK_REST_DIR_MANAGED_VERSIONING=1 \ + PYTEST_CMD="${MANAGED_BACKFILL_PYTEST_CMD}" diff --git a/Makefile b/Makefile index d321cce5d..fd4b9ed6a 100644 --- a/Makefile +++ b/Makefile @@ -197,6 +197,7 @@ docker-test: $(if $(LANCE_SPARK_START_REST_DIR),-e LANCE_SPARK_START_REST_DIR=$(LANCE_SPARK_START_REST_DIR)) \ $(if $(LANCE_SPARK_REST_DIR_ROOT),-e LANCE_SPARK_REST_DIR_ROOT=$(LANCE_SPARK_REST_DIR_ROOT)) \ $(if $(LANCE_SPARK_REST_DIR_PORT),-e LANCE_SPARK_REST_DIR_PORT=$(LANCE_SPARK_REST_DIR_PORT)) \ + $(if $(LANCE_SPARK_REST_DIR_MANAGED_VERSIONING),-e LANCE_SPARK_REST_DIR_MANAGED_VERSIONING=$(LANCE_SPARK_REST_DIR_MANAGED_VERSIONING)) \ $(if $(TEST_BACKENDS),-e TEST_BACKENDS=$(TEST_BACKENDS)) \ $(if $(LANCE_FTS_FORMAT_VERSION),-e LANCE_FTS_FORMAT_VERSION=$(LANCE_FTS_FORMAT_VERSION)) \ $(if $(AWS_REGION),-e AWS_REGION=$(AWS_REGION)) \ diff --git a/integration-tests/LanceRestDirNamespaceServer.java b/integration-tests/LanceRestDirNamespaceServer.java index 8178b9aa7..d3fe8cc91 100644 --- a/integration-tests/LanceRestDirNamespaceServer.java +++ b/integration-tests/LanceRestDirNamespaceServer.java @@ -12,29 +12,195 @@ * limitations under the License. */ -import org.lance.namespace.RestAdapter; - +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import java.io.IOException; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; import java.util.HashMap; import java.util.Map; +import java.util.Set; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicLong; +import org.lance.namespace.RestAdapter; public final class LanceRestDirNamespaceServer { + private static final String CREATE_TABLE_VERSION_COUNT_PATH = + "/__lance_test/create_table_version_count"; + private LanceRestDirNamespaceServer() {} public static void main(String[] args) throws Exception { String root = args.length > 0 ? args[0] : "/home/lance/rest-data"; String host = args.length > 1 ? args[1] : "127.0.0.1"; int port = args.length > 2 ? Integer.parseInt(args[2]) : 10024; + boolean managedVersioning = args.length > 3 && Boolean.parseBoolean(args[3]); Map backendConfig = new HashMap<>(); backendConfig.put("root", root); + if (managedVersioning) { + // DirectoryNamespace uses manifest storage to track versions managed by the namespace. + backendConfig.put("manifest_enabled", "true"); + backendConfig.put("table_version_tracking_enabled", "true"); + } - RestAdapter adapter = new RestAdapter("dir", backendConfig, host, port); - Runtime.getRuntime().addShutdownHook(new Thread(adapter::close)); + RestAdapter adapter = new RestAdapter("dir", backendConfig, host, managedVersioning ? 0 : port); adapter.start(); + + ManagedVersioningProxy proxy = + managedVersioning ? new ManagedVersioningProxy(host, port, adapter.getPort()) : null; + if (proxy != null) { + proxy.start(); + } + + Runtime.getRuntime() + .addShutdownHook( + new Thread( + () -> { + if (proxy != null) { + proxy.close(); + } + adapter.close(); + })); System.out.printf( - "Lance REST directory namespace listening on http://%s:%d with root %s%n", - host, adapter.getPort(), root); + "Lance REST directory namespace listening on http://%s:%d with root %s " + + "(managed versioning: %s)%n", + host, managedVersioning ? port : adapter.getPort(), root, managedVersioning); new CountDownLatch(1).await(); } + + /** + * Proxies the managed REST test server so tests can observe successful namespace-owned version + * commits. Direct Dataset commits bypass this proxy endpoint and therefore do not increment the + * counter. + */ + private static final class ManagedVersioningProxy implements AutoCloseable { + private static final Set HOP_BY_HOP_HEADERS = + Set.of("connection", "content-length", "expect", "host", "transfer-encoding", "upgrade"); + + private final String backendUri; + private final HttpClient client = HttpClient.newHttpClient(); + private final AtomicLong createTableVersionCount = new AtomicLong(); + private final ExecutorService executor = Executors.newCachedThreadPool(); + private final HttpServer server; + + private ManagedVersioningProxy(String host, int port, int backendPort) throws IOException { + this.backendUri = "http://" + host + ":" + backendPort; + this.server = HttpServer.create(new InetSocketAddress(host, port), 0); + this.server.createContext("/", this::handle); + this.server.setExecutor(executor); + } + + private void start() { + server.start(); + } + + private void handle(HttpExchange exchange) throws IOException { + try { + if (CREATE_TABLE_VERSION_COUNT_PATH.equals(exchange.getRequestURI().getPath())) { + handleCreateTableVersionCount(exchange); + return; + } + + byte[] requestBody = exchange.getRequestBody().readAllBytes(); + URI target = URI.create(backendUri + exchange.getRequestURI().toASCIIString()); + HttpRequest.Builder request = + HttpRequest.newBuilder(target) + .method( + exchange.getRequestMethod(), + HttpRequest.BodyPublishers.ofByteArray(requestBody)); + exchange + .getRequestHeaders() + .forEach( + (name, values) -> { + if (!isHopByHopHeader(name)) { + values.forEach(value -> request.header(name, value)); + } + }); + + HttpResponse response; + try { + response = client.send(request.build(), HttpResponse.BodyHandlers.ofByteArray()); + } catch (IOException e) { + sendResponse( + exchange, + 502, + ("REST proxy failed: " + errorMessage(e)).getBytes(StandardCharsets.UTF_8)); + return; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + sendResponse(exchange, 502, "REST proxy interrupted".getBytes(StandardCharsets.UTF_8)); + return; + } + + response + .headers() + .map() + .forEach( + (name, values) -> { + if (!isHopByHopHeader(name)) { + values.forEach(value -> exchange.getResponseHeaders().add(name, value)); + } + }); + if (isSuccessfulCreateTableVersion(exchange, response.statusCode())) { + createTableVersionCount.incrementAndGet(); + } + sendResponse(exchange, response.statusCode(), response.body()); + } catch (RuntimeException e) { + sendResponse( + exchange, + 502, + ("REST proxy failed: " + errorMessage(e)).getBytes(StandardCharsets.UTF_8)); + } finally { + exchange.close(); + } + } + + private void handleCreateTableVersionCount(HttpExchange exchange) throws IOException { + if (!"GET".equals(exchange.getRequestMethod())) { + sendResponse(exchange, 405, new byte[0]); + return; + } + exchange.getResponseHeaders().set("Content-Type", "text/plain; charset=utf-8"); + sendResponse( + exchange, + 200, + Long.toString(createTableVersionCount.get()).getBytes(StandardCharsets.UTF_8)); + } + + private static boolean isSuccessfulCreateTableVersion(HttpExchange exchange, int statusCode) { + String path = exchange.getRequestURI().getPath(); + return "POST".equals(exchange.getRequestMethod()) + && path.startsWith("/v1/table/") + && path.endsWith("/version/create") + && statusCode >= 200 + && statusCode < 300; + } + + private static boolean isHopByHopHeader(String name) { + return HOP_BY_HOP_HEADERS.contains(name.toLowerCase()); + } + + private static String errorMessage(Exception error) { + return error.getMessage() == null ? error.getClass().getSimpleName() : error.getMessage(); + } + + private static void sendResponse(HttpExchange exchange, int statusCode, byte[] body) + throws IOException { + exchange.sendResponseHeaders(statusCode, body.length); + exchange.getResponseBody().write(body); + } + + @Override + public void close() { + server.stop(0); + executor.shutdownNow(); + } + } } diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index f8916fc2b..ee0a82c83 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -196,6 +196,11 @@ def minio(): "/home/lance/rest-data", ) LANCE_SPARK_REST_DIR_PORT = int(os.environ.get("LANCE_SPARK_REST_DIR_PORT", "10024")) +LANCE_SPARK_REST_DIR_MANAGED_VERSIONING = os.environ.get( + "LANCE_SPARK_REST_DIR_MANAGED_VERSIONING", + "", +).lower() in ("1", "true", "yes") +CREATE_TABLE_VERSION_COUNT_PATH = "/__lance_test/create_table_version_count" AWS_S3_BUCKET_NAME = os.environ.get("AWS_S3_BUCKET_NAME") AWS_REGION = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1" AWS_GLUE_CATALOG_ID = os.environ.get("AWS_GLUE_CATALOG_ID") @@ -372,6 +377,7 @@ def spark(request): session.sql("CREATE NAMESPACE IF NOT EXISTS default") # Store backend name for marker-based test skipping session._lance_backend = backend + session._lance_rest_dir = rest_dir if backend == "rest-dir" else None yield session session.stop() @@ -403,17 +409,41 @@ def rest_dir_namespace(): LANCE_SPARK_REST_DIR_ROOT, "127.0.0.1", str(port), + str(LANCE_SPARK_REST_DIR_MANAGED_VERSIONING).lower(), ], stdout=log, stderr=subprocess.STDOUT, ) try: _wait_for_tcp(host, port, proc, "Lance REST directory namespace") - yield {"uri": uri} + yield { + "uri": uri, + "create_table_version_count_uri": ( + uri.rstrip("/") + CREATE_TABLE_VERSION_COUNT_PATH + if LANCE_SPARK_REST_DIR_MANAGED_VERSIONING + else None + ), + } finally: _stop_process(proc) +@pytest.fixture +def create_table_version_count(spark): + """Return the managed REST server's successful create-table-version count.""" + rest_dir = getattr(spark, "_lance_rest_dir", None) + count_uri = rest_dir.get("create_table_version_count_uri") if rest_dir else None + if not count_uri: + return None + + def retrieve_count(): + with urllib.request.urlopen(count_uri, timeout=5) as response: + return int(response.read().decode("utf-8")) + + retrieve_count() + return retrieve_count + + @pytest.fixture def test_table(request, spark): """Provide a unique table name for each test to avoid isolation issues. diff --git a/integration-tests/test_lance_spark.py b/integration-tests/test_lance_spark.py index 9c799ef3b..fabe0e41b 100644 --- a/integration-tests/test_lance_spark.py +++ b/integration-tests/test_lance_spark.py @@ -67,25 +67,40 @@ def _read_options_builder(jvm): return getattr(jvm.org.lance, "ReadOptions$Builder")() -def _lance_index_metadata(spark, table_name, index_name): - if getattr(spark, "_lance_backend", None) == "lancedb": - return None - - jvm = spark._jvm - storage_options = _java_hash_map(spark, _lance_storage_options(spark)) +def _open_direct_dataset(spark, table_name): read_options = ( - _read_options_builder(jvm) - .setStorageOptions(storage_options) - .setSession(jvm.org.lance.spark.LanceRuntime.session(LANCE_CATALOG)) + _read_options_builder(spark._jvm) + .setStorageOptions(_java_hash_map(spark, _lance_storage_options(spark))) + .setSession(spark._jvm.org.lance.spark.LanceRuntime.session(LANCE_CATALOG)) .build() ) - dataset = ( - jvm.org.lance.Dataset.open() - .allocator(jvm.org.lance.spark.LanceRuntime.allocator()) + return ( + spark._jvm.org.lance.Dataset.open() + .allocator(spark._jvm.org.lance.spark.LanceRuntime.allocator()) .uri(_table_location(spark, table_name)) .readOptions(read_options) .build() ) + + +def _commit_dataset_config_directly(spark, table_name): + """Make a direct Dataset commit that intentionally bypasses the namespace client.""" + dataset = _open_direct_dataset(spark, table_name) + try: + version_before = dataset.version() + dataset.updateConfig( + _java_hash_map(spark, {"managed_versioning_negative_control": "true"}) + ) + assert dataset.version() == version_before + 1 + finally: + dataset.close() + + +def _lance_index_metadata(spark, table_name, index_name): + if getattr(spark, "_lance_backend", None) == "lancedb": + return None + + dataset = _open_direct_dataset(spark, table_name) try: indexes = dataset.getIndexes() for pos in range(indexes.size()): @@ -2445,7 +2460,9 @@ def test_add_column_computed_values(self, spark): @pytest.mark.requires_rest @pytest.mark.rest_dir_compatible - def test_add_column_from_view_on_rest(self, spark, test_table): + def test_add_column_from_view_on_rest( + self, spark, test_table, create_table_version_count + ): spark.sql(f""" CREATE TABLE {test_table} ( id INT, @@ -2465,10 +2482,17 @@ def test_add_column_from_view_on_rest(self, spark, test_table): FROM {test_table} """) + commit_count_before = ( + create_table_version_count() if create_table_version_count else None + ) + spark.sql(f""" ALTER TABLE {test_table} ADD COLUMNS name_copy FROM namespace_add_columns_view """) + if commit_count_before is not None: + assert create_table_version_count() == commit_count_before + 1 + rows = spark.sql(f""" SELECT id, name, name_copy FROM {test_table} @@ -2480,13 +2504,21 @@ def test_add_column_from_view_on_rest(self, spark, test_table): (2, "bravo", "bravo"), ] + if commit_count_before is not None: + # Negative control: a real direct commit must bypass the namespace counter. + direct_commit_count_before = create_table_version_count() + _commit_dataset_config_directly(spark, test_table) + assert create_table_version_count() == direct_commit_count_before + class TestDMLUpdateColumn: """Test DML UPDATE COLUMNS FROM operations for updating existing columns via backfill.""" @pytest.mark.requires_rest @pytest.mark.rest_dir_compatible - def test_update_column_from_view_on_rest(self, spark, test_table): + def test_update_column_from_view_on_rest( + self, spark, test_table, create_table_version_count + ): spark.sql(f""" CREATE TABLE {test_table} ( id INT, @@ -2508,11 +2540,18 @@ def test_update_column_from_view_on_rest(self, spark, test_table): WHERE id = 2 """) + commit_count_before = ( + create_table_version_count() if create_table_version_count else None + ) + spark.sql(f""" ALTER TABLE {test_table} UPDATE COLUMNS value FROM namespace_update_columns_view """) + if commit_count_before is not None: + assert create_table_version_count() == commit_count_before + 1 + rows = spark.sql(f""" SELECT id, name, value FROM {test_table}