This is an automated email from the ASF dual-hosted git repository. roryqi pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit 991788f35f0e0fd751880ccbb05500193607f591 Author: remilu <[email protected]> AuthorDate: Mon Oct 26 10:27:44 2026 +0800 [#14221] feat(tms): Built-in TMS Iceberg maintenance Spark job Fixes #14221. Squashed from fork #77 (B10). Co-authored-by: Cursor <[email protected]> --- .../optimizer-configuration.md | 2 +- maintenance/jobs/build.gradle.kts | 2 + .../jobs/BuiltInJobTemplateProvider.java | 4 +- .../jobs/MaintenanceActionHttpClient.java | 368 ++++++++++ .../jobs/iceberg/IcebergMaintenanceJob.java | 740 +++++++++++++++++++++ .../jobs/iceberg/MaintenanceTriggerGate.java | 171 +++++ .../jobs/TestBuiltInJobTemplateProvider.java | 6 + .../jobs/TestMaintenanceActionHttpClient.java | 141 ++++ .../jobs/iceberg/TestIcebergMaintenanceJob.java | 83 +++ .../jobs/iceberg/TestMaintenanceTriggerGate.java | 67 ++ .../common/util/GravitinoAuthSettings.java | 42 ++ 11 files changed, 1624 insertions(+), 2 deletions(-) diff --git a/docs/table-maintenance-service/optimizer-configuration.md b/docs/table-maintenance-service/optimizer-configuration.md index 226a21578d..7c6e8c54b4 100644 --- a/docs/table-maintenance-service/optimizer-configuration.md +++ b/docs/table-maintenance-service/optimizer-configuration.md @@ -66,7 +66,7 @@ type is still in-flight it deletes the row and stops; otherwise it claims a soft gravitino.server.rest.extensionPackages = org.apache.gravitino.maintenance.web.rest.feature gravitino.auxService.names = iceberg-rest gravitino.eventListener.names = post-commit-maintenance-listener -gravitino.eventListener.post-commit-maintenance-listener.class = org.apache.gravitino.maintenance.listener.PostCommitMaintenanceListener +gravitino.eventListener.post-commit-maintenance-listener.class = org.apache.gravitino.maintenance.listener.IcebergUpdateTableEventListener ``` Without the listener, the crontab (`timeTrigger` → expand → cron-job) path still works; commit-triggered diff --git a/maintenance/jobs/build.gradle.kts b/maintenance/jobs/build.gradle.kts index 2d40e6a663..baef5c062d 100644 --- a/maintenance/jobs/build.gradle.kts +++ b/maintenance/jobs/build.gradle.kts @@ -36,7 +36,9 @@ val sparkMajorVersion = "3.5" dependencies { implementation(project(":api")) + implementation(project(":common")) implementation(project(":clients:client-java")) + implementation(project(":maintenance:maintenance-common")) implementation(project(":maintenance:optimizer-api")) implementation(project(":maintenance:updaters")) diff --git a/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/BuiltInJobTemplateProvider.java b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/BuiltInJobTemplateProvider.java index 6b26dce55e..1cb2f8e566 100644 --- a/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/BuiltInJobTemplateProvider.java +++ b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/BuiltInJobTemplateProvider.java @@ -26,6 +26,7 @@ import java.util.stream.Collectors; import org.apache.gravitino.job.JobTemplate; import org.apache.gravitino.job.JobTemplateProvider; import org.apache.gravitino.maintenance.jobs.iceberg.IcebergExpireSnapshotsJob; +import org.apache.gravitino.maintenance.jobs.iceberg.IcebergMaintenanceJob; import org.apache.gravitino.maintenance.jobs.iceberg.IcebergRemoveOrphanFilesJob; import org.apache.gravitino.maintenance.jobs.iceberg.IcebergRewriteDataFilesJob; import org.apache.gravitino.maintenance.jobs.iceberg.IcebergRewriteManifestsJob; @@ -53,7 +54,8 @@ public class BuiltInJobTemplateProvider implements JobTemplateProvider { new IcebergUpdateStatsAndMetricsJob(), new IcebergUpdateManifestStatsJob(), new IcebergExpireSnapshotsJob(), - new IcebergRemoveOrphanFilesJob()); + new IcebergRemoveOrphanFilesJob(), + new IcebergMaintenanceJob()); @Override public List<? extends JobTemplate> jobTemplates() { diff --git a/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/MaintenanceActionHttpClient.java b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/MaintenanceActionHttpClient.java new file mode 100644 index 0000000000..b98cea2a2d --- /dev/null +++ b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/MaintenanceActionHttpClient.java @@ -0,0 +1,368 @@ +/* + * 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. + */ +package org.apache.gravitino.maintenance.jobs; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.UnsupportedEncodingException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.util.Base64; +import java.util.HashMap; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.commons.lang3.StringUtils; +import org.apache.gravitino.auth.AuthConstants; +import org.apache.gravitino.dto.responses.BaseResponse; +import org.apache.gravitino.dto.responses.OAuth2TokenResponse; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionFinishRequest; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionMetricsRequest; +import org.apache.gravitino.maintenance.optimizer.common.conf.OptimizerConfig; +import org.apache.gravitino.maintenance.optimizer.common.util.GravitinoAuthSettings; +import org.apache.gravitino.rest.RESTRequest; +import org.apache.gravitino.rest.RESTUtils; + +/** + * HTTP client for TMS Spark callbacks ({@code metrics/before}, {@code metrics/after}, {@code + * finish}). + * + * <p>Uses {@link HttpURLConnection} only (Java 8, no {@code client-java}). Authentication follows + * the same {@code --updater-options} keys as {@code builtin-iceberg-update-stats} ({@link + * GravitinoAuthSettings}): {@code none}, {@code simple}, {@code basic}, and {@code oauth}. + */ +public final class MaintenanceActionHttpClient implements AutoCloseable { + + private static final String ACCEPT = "application/vnd.gravitino.v1+json"; + private static final String JSON = "application/json"; + private static final String FORM_URLENCODED = "application/x-www-form-urlencoded"; + private static final String SIMPLE_AUTH_DUMMY_PASSWORD = "dummy"; + private static final int CONNECT_TIMEOUT_MS = 30_000; + private static final int READ_TIMEOUT_MS = 60_000; + private static final int OAUTH_EXPIRY_SKEW_SECONDS = 30; + private static final int DEFAULT_OAUTH_EXPIRES_IN_SECONDS = 300; + private static final int ERROR_BODY_TRUNCATE = 512; + + private final ObjectMapper mapper; + private final String baseUri; + private final String metalakeName; + private final GravitinoAuthSettings authSettings; + + @Nullable private String cachedBearer; + private long bearerExpiresAtMillis; + + @VisibleForTesting + MaintenanceActionHttpClient( + ObjectMapper mapper, + String baseUri, + String metalakeName, + GravitinoAuthSettings authSettings) { + this.mapper = Preconditions.checkNotNull(mapper, "mapper"); + this.baseUri = trimTrailingSlash(Preconditions.checkNotNull(baseUri, "baseUri")); + this.metalakeName = Preconditions.checkNotNull(metalakeName, "metalakeName"); + this.authSettings = Preconditions.checkNotNull(authSettings, "authSettings"); + } + + /** + * Creates a client from updater-options (same shape as iceberg-update-stats). + * + * @param updaterOptions flat map from {@code --updater-options} JSON + * @return open client + */ + public static MaintenanceActionHttpClient fromUpdaterOptions(Map<String, String> updaterOptions) { + Map<String, String> properties = normalizeUpdaterOptions(updaterOptions); + String uri = properties.get(OptimizerConfig.GRAVITINO_URI); + String metalake = properties.get(OptimizerConfig.GRAVITINO_METALAKE); + Preconditions.checkArgument( + StringUtils.isNotBlank(uri), "updater_options must contain gravitino_uri"); + Preconditions.checkArgument( + StringUtils.isNotBlank(metalake), "updater_options must contain metalake"); + GravitinoAuthSettings auth = GravitinoAuthSettings.from(new OptimizerConfig(properties)); + return new MaintenanceActionHttpClient( + JsonUtils.objectMapper(), uri.trim(), metalake.trim(), auth); + } + + /** + * Reports S0 (before-maintenance) snapshot metrics. + * + * @param jobId Gravitino job run id + * @param request metrics payload + */ + public void reportBeforeMetrics(long jobId, MaintenanceActionMetricsRequest request) { + postJson(jobId, "metrics/before", request); + } + + /** + * Reports S1 (after-maintenance) snapshot metrics. + * + * @param jobId Gravitino job run id + * @param request metrics payload + */ + public void reportAfterMetrics(long jobId, MaintenanceActionMetricsRequest request) { + postJson(jobId, "metrics/after", request); + } + + /** + * Closes the soft slot for one table without writing metrics. + * + * @param jobId Gravitino job run id + * @param request finish payload + */ + public void finish(long jobId, MaintenanceActionFinishRequest request) { + postJson(jobId, "finish", request); + } + + @Override + public void close() { + cachedBearer = null; + bearerExpiresAtMillis = 0L; + } + + private void postJson(long jobId, String action, RESTRequest request) { + Preconditions.checkArgument(request != null, "request must not be null"); + request.validate(); + String path = + String.format( + "/api/metalakes/%s/maintenance/actions/%d/%s", + RESTUtils.encodeString(metalakeName), jobId, action); + try { + HttpResult result = + post(baseUri + path, JSON, mapper.writeValueAsString(request), authorizationHeader()); + if (!isSuccess(result.status)) { + throw new IllegalStateException( + String.format( + "TMS %s failed: HTTP %d body=%s", action, result.status, truncate(result.body))); + } + if (StringUtils.isNotBlank(result.body)) { + BaseResponse parsed = mapper.readValue(result.body, BaseResponse.class); + parsed.validate(); + } + } catch (IOException e) { + throw new IllegalStateException("Failed to post TMS " + action + ": " + e.getMessage(), e); + } + } + + @Nullable + private String authorizationHeader() { + switch (authSettings.authType()) { + case GravitinoAuthSettings.TYPE_SIMPLE: + return basicAuthorization(simpleUsername(), SIMPLE_AUTH_DUMMY_PASSWORD); + case GravitinoAuthSettings.TYPE_BASIC: + requireAuthValue(authSettings.username(), "username is required for basic authentication"); + requireAuthValue(authSettings.password(), "password is required for basic authentication"); + return basicAuthorization(authSettings.username(), authSettings.password()); + case GravitinoAuthSettings.TYPE_OAUTH: + return AuthConstants.AUTHORIZATION_BEARER_HEADER + oauthAccessToken(); + case GravitinoAuthSettings.TYPE_NONE: + return null; + default: + throw new IllegalArgumentException("Unsupported auth_type: " + authSettings.authType()); + } + } + + private String simpleUsername() { + if (StringUtils.isNotBlank(authSettings.username())) { + return authSettings.username().trim(); + } + String envUser = System.getenv("GRAVITINO_USER"); + if (StringUtils.isNotBlank(envUser)) { + return envUser.trim(); + } + return System.getProperty("user.name"); + } + + private String oauthAccessToken() { + long now = System.currentTimeMillis(); + if (cachedBearer != null && now < bearerExpiresAtMillis) { + return cachedBearer; + } + requireAuthValue( + authSettings.oauthServerUri(), "oauth server URI is required for oauth authentication"); + requireAuthValue(authSettings.oauthPath(), "oauth path is required for oauth authentication"); + requireAuthValue( + authSettings.oauthCredential(), "oauth credential is required for oauth authentication"); + requireAuthValue(authSettings.oauthScope(), "oauth scope is required for oauth authentication"); + + String tokenUrl = joinUrl(authSettings.oauthServerUri(), authSettings.oauthPath()); + try { + HttpResult result = post(tokenUrl, FORM_URLENCODED, buildClientCredentialsForm(), null); + if (!isSuccess(result.status)) { + throw new IllegalStateException("OAuth token request failed: HTTP " + result.status); + } + OAuth2TokenResponse token = mapper.readValue(result.body, OAuth2TokenResponse.class); + if (StringUtils.isBlank(token.getAccessToken())) { + throw new IllegalStateException("OAuth token response did not include access_token"); + } + cachedBearer = token.getAccessToken(); + int expiresIn = + token.getExpiresIn() == null ? DEFAULT_OAUTH_EXPIRES_IN_SECONDS : token.getExpiresIn(); + bearerExpiresAtMillis = now + Math.max(1, expiresIn - OAUTH_EXPIRY_SKEW_SECONDS) * 1000L; + return cachedBearer; + } catch (IOException e) { + throw new IllegalStateException("Failed to request OAuth token: " + e.getMessage(), e); + } + } + + private String buildClientCredentialsForm() { + String credential = authSettings.oauthCredential().trim(); + String clientId = null; + String clientSecret = credential; + int split = credential.indexOf(':'); + if (split >= 0) { + clientId = credential.substring(0, split); + clientSecret = credential.substring(split + 1); + } + StringBuilder form = new StringBuilder("grant_type=client_credentials"); + if (StringUtils.isNotBlank(clientId)) { + form.append("&client_id=").append(urlEncode(clientId)); + } + form.append("&client_secret=").append(urlEncode(clientSecret)); + form.append("&scope=").append(urlEncode(authSettings.oauthScope().trim())); + return form.toString(); + } + + private static HttpResult post( + String url, String contentType, String body, @Nullable String authorization) + throws IOException { + HttpURLConnection connection = (HttpURLConnection) URI.create(url).toURL().openConnection(); + connection.setRequestMethod("POST"); + connection.setConnectTimeout(CONNECT_TIMEOUT_MS); + connection.setReadTimeout(READ_TIMEOUT_MS); + connection.setDoOutput(true); + connection.setRequestProperty("Accept", ACCEPT); + connection.setRequestProperty("Content-Type", contentType); + if (authorization != null) { + connection.setRequestProperty(AuthConstants.HTTP_HEADER_AUTHORIZATION, authorization); + } + byte[] bytes = body.getBytes(StandardCharsets.UTF_8); + connection.setFixedLengthStreamingMode(bytes.length); + try { + try (OutputStream out = connection.getOutputStream()) { + out.write(bytes); + } + int status = connection.getResponseCode(); + InputStream stream = + isSuccess(status) ? connection.getInputStream() : connection.getErrorStream(); + return new HttpResult(status, readUtf8(stream)); + } finally { + connection.disconnect(); + } + } + + @VisibleForTesting + static Map<String, String> normalizeUpdaterOptions(@Nullable Map<String, String> updaterOptions) { + Map<String, String> properties = + updaterOptions == null ? new HashMap<>() : new HashMap<>(updaterOptions); + putCanonical(properties, OptimizerConfig.GRAVITINO_URI, "gravitino_uri", "gravitino-uri"); + putCanonical(properties, OptimizerConfig.GRAVITINO_METALAKE, "metalake"); + GravitinoAuthSettings.copyAliases(properties); + return properties; + } + + private static void putCanonical( + Map<String, String> properties, String canonical, String... aliases) { + if (StringUtils.isNotBlank(properties.get(canonical))) { + return; + } + for (String alias : aliases) { + String value = properties.get(alias); + if (StringUtils.isNotBlank(value)) { + properties.put(canonical, value.trim()); + return; + } + } + } + + private static void requireAuthValue(@Nullable String value, String message) { + Preconditions.checkArgument(StringUtils.isNotBlank(value), message); + } + + private static boolean isSuccess(int status) { + return status >= 200 && status < 300; + } + + private static String basicAuthorization(String user, String password) { + String userInformation = user + ":" + password; + return AuthConstants.AUTHORIZATION_BASIC_HEADER + + Base64.getEncoder().encodeToString(userInformation.getBytes(StandardCharsets.UTF_8)); + } + + private static String joinUrl(String serverUri, String path) { + String base = trimTrailingSlash(serverUri); + String suffix = path.startsWith("/") ? path : "/" + path; + return base + suffix; + } + + private static String urlEncode(String value) { + try { + return URLEncoder.encode(value, StandardCharsets.UTF_8.name()); + } catch (UnsupportedEncodingException e) { + throw new IllegalStateException(e); + } + } + + private static String trimTrailingSlash(String uri) { + String trimmed = uri.trim(); + while (trimmed.endsWith("/")) { + trimmed = trimmed.substring(0, trimmed.length() - 1); + } + return trimmed; + } + + private static String readUtf8(@Nullable InputStream stream) throws IOException { + if (stream == null) { + return ""; + } + try (InputStream in = stream; + ByteArrayOutputStream out = new ByteArrayOutputStream()) { + byte[] buffer = new byte[4096]; + int read; + while ((read = in.read(buffer)) >= 0) { + out.write(buffer, 0, read); + } + return out.toString(StandardCharsets.UTF_8.name()); + } + } + + private static String truncate(@Nullable String body) { + if (body == null) { + return ""; + } + return body.length() <= ERROR_BODY_TRUNCATE + ? body + : body.substring(0, ERROR_BODY_TRUNCATE) + "..."; + } + + private static final class HttpResult { + private final int status; + private final String body; + + private HttpResult(int status, String body) { + this.status = status; + this.body = body; + } + } +} diff --git a/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/iceberg/IcebergMaintenanceJob.java b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/iceberg/IcebergMaintenanceJob.java new file mode 100644 index 0000000000..1a5419e70b --- /dev/null +++ b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/iceberg/IcebergMaintenanceJob.java @@ -0,0 +1,740 @@ +/* + * 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. + */ +package org.apache.gravitino.maintenance.jobs.iceberg; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.google.common.collect.ImmutableList; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.gravitino.job.JobHandle; +import org.apache.gravitino.job.JobTemplateProvider; +import org.apache.gravitino.job.SparkJobTemplate; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionFinishRequest; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionMetricsRequest; +import org.apache.gravitino.maintenance.jobs.BuiltInJob; +import org.apache.gravitino.maintenance.jobs.MaintenanceActionHttpClient; +import org.apache.gravitino.maintenance.optimizer.common.util.IcebergSparkConfigUtils; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.spark.Spark3Util; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * TMS Spark entry for cron (§5.2.2) and commit-chain (§5.2.3) {@code runJob} submits. + * + * <p>Not a user-facing CLI. Operators / optimizer keep using the four standalone {@code + * Iceberg*Job} templates. This class is launched only via the built-in job template; program + * arguments are template-expanded {@code jobConf} values. + * + * <p>Dispatches by {@code maintenance_types} (one wire name or a comma-separated commit chain). + * Reuses procedure SQL builders from the standalone {@code Iceberg*Job} classes. Continue / stop + * (snapshot drift and trigger-expr recommend) is decided in-job after S0 sample. Soft slots are + * closed via the finish callback in {@code finally}. + * + * <p>TMS-specific {@code jobConf} keys: {@code maintenance_types}, {@code policy_ids} (parallel + * comma-separated policy entity ids), {@code tables} (JSON targets with system-filled {@code + * table_id}), {@code updater_options}, {@code rules}. Shared procedure / Spark keys match other + * Iceberg built-ins. The Gravitino job id is read from the local executor staging directory ({@code + * job-} plus the numeric id), or from {@code GRAVITINO_JOB_ID}. + */ +public class IcebergMaintenanceJob implements BuiltInJob { + + private static final Logger LOG = LoggerFactory.getLogger(IcebergMaintenanceJob.class); + + /** + * Optional job-id override. Not injected by {@code JobManager}; the local executor staging + * directory is the default source. + */ + private static final String GRAVITINO_JOB_ID_ENV = "GRAVITINO_JOB_ID"; + + /** Built-in template name submitted by TMS cron / commit. */ + public static final String NAME = + JobTemplateProvider.BUILTIN_NAME_PREFIX + "iceberg-maintenance-task"; + + /** + * JobConf key for procedure dispatch ({@code maintenance_types}). One wire name, or a + * comma-separated commit chain. Required; no template default. + */ + public static final String MAINTENANCE_TYPES_KEY = "maintenance_types"; + + /** + * JobConf key for table targets ({@code tables}): JSON array of {@code table_id}, {@code + * catalog_name}, {@code table_identifier}. Single- and multi-table TMS submits use the same + * shape; {@code table_id} is system-filled by enqueue/submit. + */ + public static final String TABLES_KEY = "tables"; + + /** + * JobConf key for policy entity ids parallel to {@link #MAINTENANCE_TYPES_KEY} (comma-separated). + * Required so metrics/finish callbacks can attribute soft slots when one job binds multiple + * policies. + */ + public static final String POLICY_IDS_KEY = "policy_ids"; + + private static final String VERSION = "v1"; + + /** Supported {@code maintenance_types} wire values. */ + public enum MaintenanceType { + REWRITE_DATA_FILES("rewrite_data_files"), + REWRITE_MANIFESTS("rewrite_manifests"), + EXPIRE_SNAPSHOTS("expire_snapshots"), + REMOVE_ORPHAN_FILES("remove_orphan_files"); + + private final String wireName; + + MaintenanceType(String wireName) { + this.wireName = wireName; + } + + /** + * @return wire value used in jobConf + */ + public String wireName() { + return wireName; + } + + /** + * @param raw jobConf {@code maintenance_types} segment + * @return parsed type + */ + public static MaintenanceType fromWire(@Nullable String raw) { + if (raw == null || raw.trim().isEmpty()) { + throw new IllegalArgumentException("maintenance_types is required"); + } + String normalized = raw.trim().toLowerCase(Locale.ROOT); + for (MaintenanceType type : values()) { + if (type.wireName.equals(normalized)) { + return type; + } + } + throw new IllegalArgumentException( + "Unknown maintenance_types value '" + + raw + + "'. Expected one of: rewrite_data_files, rewrite_manifests, expire_snapshots," + + " remove_orphan_files"); + } + } + + @Override + public SparkJobTemplate jobTemplate() { + return SparkJobTemplate.builder() + .withName(NAME) + .withComment( + "TMS Iceberg maintenance: S0 sample/gate, procedure dispatch (single or commit-chain)," + + " S1 finish") + .withExecutable(resolveExecutable(IcebergMaintenanceJob.class)) + .withClassName(IcebergMaintenanceJob.class.getName()) + .withArguments(buildArguments()) + .withConfigs(IcebergSparkConfigUtils.buildTemplateSparkConfigs()) + .withCustomFields( + Collections.singletonMap(JobTemplateProvider.PROPERTY_VERSION_KEY, VERSION)) + .build(); + } + + /** + * Spark application entry used by the job executor after template expansion. Not a supported + * operator CLI. + * + * @param args template-expanded jobConf as named arguments + */ + public static void main(String[] args) { + int code = run(args); + if (code != 0) { + System.exit(code); + } + } + + /** + * @param args template-expanded jobConf as named arguments + * @return process exit code ({@code 0} success, {@code 1} failure) + */ + static int run(String[] args) { + Map<String, String> argMap = IcebergJobUtils.parseArguments(args); + + List<MaintenanceType> types; + List<Long> policyIds; + long jobId; + List<TableTarget> tables; + Map<String, String> updaterOptions; + try { + types = parseMaintenanceTypes(argMap.get("maintenance-types")); + policyIds = parsePolicyIds(argMap.get("policy-ids"), types.size()); + jobId = resolveGravitinoJobId(); + tables = parseTables(IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("tables"))); + updaterOptions = + parseUpdaterOptions( + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("updater-options"))); + } catch (IllegalArgumentException | IllegalStateException e) { + LOG.error("Invalid TMS maintenance arguments: {}", e.getMessage()); + return 1; + } + + String appLabel = + types.size() == 1 + ? types.get(0).wireName + : "chain:" + types.stream().map(t -> t.wireName).collect(Collectors.joining(",")); + SparkSession.Builder sparkBuilder = + SparkSession.builder().appName("Gravitino Iceberg Maintenance Task (" + appLabel + ")"); + String sparkConfJson = IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("spark-conf")); + if (sparkConfJson != null && !sparkConfJson.isEmpty()) { + IcebergJobUtils.parseCustomSparkConfigs(sparkConfJson).forEach(sparkBuilder::config); + } + + try (SparkSession spark = sparkBuilder.getOrCreate()) { + IcebergJobUtils.requireIcebergSparkRuntimeOrExit(spark); + try (MaintenanceActionHttpClient client = + MaintenanceActionHttpClient.fromUpdaterOptions(updaterOptions)) { + // Soft slots are per (tableId, policyId) and already bound to jobId; close every slot even + // when a procedure fails or we never reach a later type/table. + Set<String> finishedSlots = new HashSet<>(); + try { + for (TableTarget table : tables) { + for (int i = 0; i < types.size(); i++) { + MaintenanceType type = types.get(i); + long policyId = policyIds.get(i); + Long afterSnapshotId = null; + try { + SnapshotSample s0 = + sampleSnapshot(spark, table.catalogName, table.tableIdentifier, type); + reportMetricsQuietly(client, jobId, table.tableId, policyId, s0, /*after=*/ false); + + if (s0.snapshotId != s0.headSnapshotId) { + LOG.info( + "TMS type {} skipped for job {} table {} policy {} due to snapshot drift" + + " (s0={} head={})", + type.wireName, + jobId, + table.tableId, + policyId, + s0.snapshotId, + s0.headSnapshotId); + continue; + } + + if (!MaintenanceTriggerGate.shouldMaintain(type, argMap, s0.metricsJson)) { + LOG.info( + "TMS type {} skipped for job {} table {} policy {} due to in-Spark recommend" + + " skip", + type.wireName, + jobId, + table.tableId, + policyId); + continue; + } + + runProcedure(spark, type, table.catalogName, table.tableIdentifier, argMap); + SnapshotSample s1 = + sampleSnapshot(spark, table.catalogName, table.tableIdentifier, type); + afterSnapshotId = s1.snapshotId; + reportMetricsQuietly(client, jobId, table.tableId, policyId, s1, /*after=*/ true); + } finally { + finishQuietly(client, jobId, table.tableId, policyId, afterSnapshotId); + finishedSlots.add(slotKey(table.tableId, policyId)); + } + } + } + return 0; + } catch (Exception e) { + LOG.error("TMS Iceberg maintenance failed for job {}", jobId, e); + return 1; + } finally { + for (TableTarget table : tables) { + for (long policyId : policyIds) { + if (!finishedSlots.contains(slotKey(table.tableId, policyId))) { + finishQuietly(client, jobId, table.tableId, policyId, /*afterSnapshotId=*/ null); + } + } + } + } + } + } + } + + private static List<TableTarget> parseTables(@Nullable String tablesJson) { + if (tablesJson == null || tablesJson.trim().isEmpty()) { + throw new IllegalArgumentException( + "tables JSON is required (array of table_id, catalog_name, table_identifier)"); + } + try { + List<Map<String, Object>> raw = + JsonUtils.anyFieldMapper() + .readValue(tablesJson, new TypeReference<List<Map<String, Object>>>() {}); + if (raw == null || raw.isEmpty()) { + throw new IllegalArgumentException("tables must not be empty"); + } + List<TableTarget> tables = new ArrayList<>(raw.size()); + for (Map<String, Object> entry : raw) { + if (entry == null) { + throw new IllegalArgumentException("tables entry must not be null"); + } + Object tableIdObj = entry.get("table_id"); + Object catalogObj = entry.get("catalog_name"); + Object tableObj = entry.get("table_identifier"); + if (tableIdObj == null || catalogObj == null || tableObj == null) { + throw new IllegalArgumentException( + "tables entry requires table_id, catalog_name, table_identifier"); + } + long tableId = Long.parseLong(String.valueOf(tableIdObj)); + if (tableId <= 0) { + throw new IllegalArgumentException("tables table_id must be positive: " + tableId); + } + String catalogName = String.valueOf(catalogObj).trim(); + String tableIdentifier = String.valueOf(tableObj).trim(); + if (catalogName.isEmpty() || tableIdentifier.isEmpty()) { + throw new IllegalArgumentException( + "tables catalog_name and table_identifier must not be blank"); + } + tables.add(new TableTarget(tableId, catalogName, tableIdentifier)); + } + return tables; + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("tables must be valid JSON: " + e.getMessage(), e); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("tables table_id must be a long", e); + } + } + + private static ImmutableList<MaintenanceType> parseMaintenanceTypes( + @Nullable String maintenanceTypes) { + String chain = IcebergJobUtils.nullIfUnresolvedPlaceholder(maintenanceTypes); + if (chain == null || chain.trim().isEmpty()) { + throw new IllegalArgumentException("maintenance_types is required"); + } + ImmutableList.Builder<MaintenanceType> parsed = ImmutableList.builder(); + for (String part : chain.split(",")) { + String trimmed = part.trim(); + if (!trimmed.isEmpty()) { + parsed.add(MaintenanceType.fromWire(trimmed)); + } + } + ImmutableList<MaintenanceType> types = parsed.build(); + if (types.isEmpty()) { + throw new IllegalArgumentException("maintenance_types must not be empty"); + } + return types; + } + + private static List<Long> parsePolicyIds(@Nullable String policyIds, int expectedCount) { + String raw = IcebergJobUtils.nullIfUnresolvedPlaceholder(policyIds); + if (raw == null || raw.trim().isEmpty()) { + throw new IllegalArgumentException("policy_ids is required"); + } + List<Long> parsed = new ArrayList<>(); + for (String part : raw.split(",")) { + String trimmed = part.trim(); + if (trimmed.isEmpty()) { + continue; + } + try { + long policyId = Long.parseLong(trimmed); + if (policyId <= 0) { + throw new IllegalArgumentException("policy_ids must be positive longs: " + trimmed); + } + parsed.add(policyId); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("policy_ids must be longs: " + trimmed, e); + } + } + if (parsed.size() != expectedCount) { + throw new IllegalArgumentException( + String.format( + "policy_ids count %s must match maintenance_types count %s", + parsed.size(), expectedCount)); + } + return parsed; + } + + private static String slotKey(long tableId, long policyId) { + return tableId + ":" + policyId; + } + + private static void runProcedure( + SparkSession spark, + MaintenanceType type, + String catalogName, + String tableIdentifier, + Map<String, String> argMap) + throws Exception { + switch (type) { + case REWRITE_DATA_FILES: + { + String strategy = firstNonBlank(argMap.get("strategy"), "binpack"); + String sql = + IcebergRewriteDataFilesJob.buildProcedureCall( + catalogName, + tableIdentifier, + strategy, + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("sort-order")), + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("where")), + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("options"))); + LOG.info("Executing TMS rewrite_data_files: {}", sql); + spark.sql(sql).collect(); + break; + } + case REWRITE_MANIFESTS: + { + IcebergRewriteManifestsJob.validateUseCaching( + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("use-caching"))); + IcebergRewriteManifestsJob.validateSpecId( + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("spec-id"))); + String sql = + IcebergRewriteManifestsJob.buildProcedureCall( + catalogName, + tableIdentifier, + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("use-caching")), + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("spec-id"))); + LOG.info("Executing TMS rewrite_manifests: {}", sql); + spark.sql(sql).collect(); + break; + } + case EXPIRE_SNAPSHOTS: + { + String retainLast = + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("retain-last")); + IcebergExpireSnapshotsJob.validateRetainLast(retainLast); + boolean streamResults = + IcebergExpireSnapshotsJob.parseStreamResults( + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("stream-results"))); + String sql = + IcebergExpireSnapshotsJob.buildProcedureCall( + catalogName, + tableIdentifier, + IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get("older-than")), + retainLast, + streamResults); + LOG.info("Executing TMS expire_snapshots: {}", sql); + Row[] results = (Row[]) spark.sql(sql).collect(); + if (results.length > 0) { + LOG.info( + "Expire snapshots deleted data/manifest/manifest-list files: {}/{}/{}", + results[0].getLong(0), + results[0].getLong(1), + results[0].getLong(2)); + } + break; + } + case REMOVE_ORPHAN_FILES: + { + Map<String, String> orphanOpts = new HashMap<>(); + orphanOpts.put("catalog", catalogName); + orphanOpts.put("table", tableIdentifier); + putIfPresent(orphanOpts, "older-than", argMap.get("older-than")); + putIfPresent(orphanOpts, "location", argMap.get("location")); + putIfPresent(orphanOpts, "dry-run", argMap.get("dry-run")); + IcebergRemoveOrphanFilesJob.execute(spark, orphanOpts); + break; + } + default: + throw new IllegalStateException("Unhandled maintenance type: " + type); + } + } + + private static SnapshotSample sampleSnapshot( + SparkSession spark, String catalogName, String tableIdentifier, MaintenanceType type) { + // Capture S0 first, then collect metrics (may take time), then re-read HEAD for the drift gate. + long snapshotId = currentSnapshotId(spark, catalogName, tableIdentifier); + Map<String, Object> payload = new HashMap<>(); + payload.put("maintenanceType", type.wireName()); + payload.put("snapshotId", snapshotId); + if (type == MaintenanceType.REWRITE_DATA_FILES) { + payload.putAll(collectCompactionMetrics(spark, catalogName, tableIdentifier)); + } + long headSnapshotId = currentSnapshotId(spark, catalogName, tableIdentifier); + try { + return new SnapshotSample( + snapshotId, headSnapshotId, JsonUtils.anyFieldMapper().writeValueAsString(payload)); + } catch (JsonProcessingException e) { + throw new IllegalStateException("Failed to serialize metrics JSON", e); + } + } + + /** + * Compaction metrics from the Iceberg files metadata table (same formulas as {@link + * IcebergUpdateStatsAndMetricsJob}). For partitioned tables, uses max partition MSE and sum of + * delete-file counts so any hot partition can trip {@code trigger-expr}. + */ + private static Map<String, Object> collectCompactionMetrics( + SparkSession spark, String catalogName, String tableIdentifier) { + boolean partitioned = + IcebergUpdateStatsAndMetricsJob.isPartitionedTable(spark, catalogName, tableIdentifier); + if (partitioned) { + String sql = + IcebergUpdateStatsAndMetricsJob.buildPartitionStatsSql(catalogName, tableIdentifier); + Row[] rows = (Row[]) spark.sql(sql).collect(); + Map<String, Object> metrics = new HashMap<>(); + double maxMse = 0D; + long deleteFiles = 0L; + for (Row row : rows) { + maxMse = Math.max(maxMse, numberAsDouble(row, "data-file-mse")); + deleteFiles += numberAsLong(row, "delete-file-number"); + } + metrics.put(IcebergDataCompactionContent.DATA_FILE_MSE_METRIC, maxMse); + metrics.put(IcebergDataCompactionContent.DELETE_FILE_NUMBER_METRIC, deleteFiles); + return metrics; + } + + String sql = IcebergUpdateStatsAndMetricsJob.buildTableStatsSql(catalogName, tableIdentifier); + Row[] rows = (Row[]) spark.sql(sql).collect(); + Map<String, Object> metrics = new HashMap<>(); + if (rows.length == 0) { + metrics.put(IcebergDataCompactionContent.DATA_FILE_MSE_METRIC, 0D); + metrics.put(IcebergDataCompactionContent.DELETE_FILE_NUMBER_METRIC, 0L); + return metrics; + } + Row row = rows[0]; + metrics.put( + IcebergDataCompactionContent.DATA_FILE_MSE_METRIC, numberAsDouble(row, "data-file-mse")); + metrics.put( + IcebergDataCompactionContent.DELETE_FILE_NUMBER_METRIC, + numberAsLong(row, "delete-file-number")); + return metrics; + } + + private static double numberAsDouble(Row row, String field) { + Number number = row.getAs(field); + return number == null ? 0D : number.doubleValue(); + } + + private static long numberAsLong(Row row, String field) { + Number number = row.getAs(field); + return number == null ? 0L : number.longValue(); + } + + private static long currentSnapshotId( + SparkSession spark, String catalogName, String tableIdentifier) { + // Catch Exception (not Spark NoSuchTableException) so BuiltInJobTemplateProvider can load this + // class on the Gravitino server classpath, which does not ship Spark catalyst types. + try { + Table table = + Spark3Util.loadIcebergTable(spark, String.format("%s.%s", catalogName, tableIdentifier)); + Snapshot snapshot = table.currentSnapshot(); + if (snapshot == null) { + throw new IllegalStateException( + "Iceberg table has no current snapshot: " + catalogName + "." + tableIdentifier); + } + return snapshot.snapshotId(); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new IllegalStateException( + "Failed to load Iceberg snapshot for " + catalogName + "." + tableIdentifier, e); + } + } + + private static void finishQuietly( + MaintenanceActionHttpClient client, + long jobId, + long tableId, + long policyId, + @Nullable Long afterSnapshotId) { + try { + client.finish(jobId, new MaintenanceActionFinishRequest(tableId, policyId, afterSnapshotId)); + } catch (Exception e) { + LOG.warn( + "Best-effort TMS finish failed for job {} table {} policy {} afterSnapshotId={}", + jobId, + tableId, + policyId, + afterSnapshotId, + e); + } + } + + private static void reportMetricsQuietly( + MaintenanceActionHttpClient client, + long jobId, + long tableId, + long policyId, + SnapshotSample sample, + boolean after) { + if (sample == null || sample.metricsJson == null || sample.metricsJson.trim().isEmpty()) { + return; + } + try { + MaintenanceActionMetricsRequest request = + new MaintenanceActionMetricsRequest( + tableId, policyId, sample.snapshotId, sample.metricsJson); + if (after) { + client.reportAfterMetrics(jobId, request); + } else { + client.reportBeforeMetrics(jobId, request); + } + } catch (Exception e) { + LOG.warn( + "Best-effort TMS {} metrics report failed for job {} table {} policy {} snapshot {}", + after ? "after" : "before", + jobId, + tableId, + policyId, + sample.snapshotId, + e); + } + } + + private static Map<String, String> parseUpdaterOptions(@Nullable String json) { + if (json == null || json.trim().isEmpty()) { + return Collections.emptyMap(); + } + try { + @SuppressWarnings("unchecked") + Map<String, String> parsed = JsonUtils.anyFieldMapper().readValue(json, Map.class); + return parsed == null ? Collections.emptyMap() : parsed; + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("Invalid updater-options JSON: " + e.getMessage(), e); + } + } + + /** + * Prefers {@code GRAVITINO_JOB_ID}. Falls back to parsing a {@code job-} plus numeric id segment + * from {@code user.dir} (local executor staging directory). + */ + private static long resolveGravitinoJobId() { + String envValue = System.getenv(GRAVITINO_JOB_ID_ENV); + if (envValue != null && !envValue.trim().isEmpty()) { + try { + long parsed = Long.parseLong(envValue.trim()); + if (parsed <= 0) { + throw new IllegalStateException( + GRAVITINO_JOB_ID_ENV + " must be a positive long: " + envValue); + } + return parsed; + } catch (NumberFormatException e) { + throw new IllegalStateException(GRAVITINO_JOB_ID_ENV + " must be a long: " + envValue, e); + } + } + String userDir = System.getProperty("user.dir"); + if (userDir != null && !userDir.trim().isEmpty()) { + Path dir = Paths.get(userDir).toAbsolutePath().normalize(); + for (Path current = dir; current != null; current = current.getParent()) { + Path fileName = current.getFileName(); + if (fileName == null) { + continue; + } + String name = fileName.toString(); + if (!name.startsWith(JobHandle.JOB_ID_PREFIX)) { + continue; + } + try { + long parsed = Long.parseLong(name.substring(JobHandle.JOB_ID_PREFIX.length())); + if (parsed > 0) { + return parsed; + } + } catch (NumberFormatException ignored) { + // Keep walking parents; staging dirs use job-<numericId>. + } + } + } + throw new IllegalStateException( + "Cannot resolve Gravitino job id: set " + + GRAVITINO_JOB_ID_ENV + + " or run under a job-<id> staging directory"); + } + + private static void putIfPresent(Map<String, String> target, String key, @Nullable String raw) { + String value = IcebergJobUtils.nullIfUnresolvedPlaceholder(raw); + if (value != null && !value.isEmpty()) { + target.put(key, value); + } + } + + private static String firstNonBlank(@Nullable String value, String defaultValue) { + String normalized = IcebergJobUtils.nullIfUnresolvedPlaceholder(value); + return normalized == null || normalized.isEmpty() ? defaultValue : normalized; + } + + private static List<String> buildArguments() { + // Named args are the Spark job-executor wiring for expanded jobConf — not a user CLI. + return Arrays.asList( + "--maintenance-types", + "{{maintenance_types}}", + "--policy-ids", + "{{policy_ids}}", + "--tables", + "{{tables}}", + "--updater-options", + "{{updater_options:-}}", + "--strategy", + "{{strategy:-binpack}}", + "--sort-order", + "{{sort_order:-}}", + "--where", + "{{where_clause:-}}", + "--options", + "{{options:-}}", + "--use-caching", + "{{use_caching:-}}", + "--spec-id", + "{{spec_id:-}}", + "--older-than", + "{{older_than:-}}", + "--retain-last", + "{{retain_last:-}}", + "--stream-results", + "{{stream_results:-false}}", + "--location", + "{{location:-}}", + "--dry-run", + "{{dry_run:-false}}", + "--spark-conf", + "{{spark_conf:-}}", + "--rules", + "{{rules:-}}"); + } + + private static final class TableTarget { + private final long tableId; + private final String catalogName; + private final String tableIdentifier; + + private TableTarget(long tableId, String catalogName, String tableIdentifier) { + this.tableId = tableId; + this.catalogName = catalogName; + this.tableIdentifier = tableIdentifier; + } + } + + private static final class SnapshotSample { + private final long snapshotId; + private final long headSnapshotId; + private final String metricsJson; + + private SnapshotSample(long snapshotId, long headSnapshotId, String metricsJson) { + this.snapshotId = snapshotId; + this.headSnapshotId = headSnapshotId; + this.metricsJson = metricsJson; + } + } +} diff --git a/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/iceberg/MaintenanceTriggerGate.java b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/iceberg/MaintenanceTriggerGate.java new file mode 100644 index 0000000000..262785428a --- /dev/null +++ b/maintenance/jobs/src/main/java/org/apache/gravitino/maintenance/jobs/iceberg/MaintenanceTriggerGate.java @@ -0,0 +1,171 @@ +/* + * 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. + */ +package org.apache.gravitino.maintenance.jobs.iceberg; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.Map; +import java.util.Map.Entry; +import javax.annotation.Nullable; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.jobs.iceberg.IcebergMaintenanceJob.MaintenanceType; +import org.apache.gravitino.maintenance.optimizer.recommender.util.ExpressionEvaluator; +import org.apache.gravitino.maintenance.optimizer.recommender.util.QLExpressionEvaluator; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.IcebergOrphanFileRemovalContent; +import org.apache.gravitino.policy.PolicyContents; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * In-Spark maintain-or-skip gate after S0 metrics (§5.2.2). + * + * <p>Evaluates enqueue-time policy rules locally with the same QL stack as the optimizer. Does not + * call Gravitino. Expire / rewrite-manifests placeholders always maintain (optimizer-only {@code + * trigger-expr=false} on those types). + */ +final class MaintenanceTriggerGate { + + private static final Logger LOG = LoggerFactory.getLogger(MaintenanceTriggerGate.class); + private static final ExpressionEvaluator EVALUATOR = new QLExpressionEvaluator(); + + /** Spark job argument key: JSON map of wire maintenance type → policy {@code rules()} object. */ + static final String RULES_ARG = "rules"; + + private MaintenanceTriggerGate() {} + + /** + * @param type maintenance type for this sample + * @param argMap Spark job arguments (includes optional {@link #RULES_ARG}) + * @param metricsJson S0 metrics JSON + * @return whether the procedure should run + */ + static boolean shouldMaintain( + MaintenanceType type, Map<String, String> argMap, @Nullable String metricsJson) { + if (type == MaintenanceType.EXPIRE_SNAPSHOTS || type == MaintenanceType.REWRITE_MANIFESTS) { + return true; + } + + Map<String, Object> rules = rulesForType(type, argMap); + Map<String, Object> metrics = parseNumericMetrics(metricsJson); + if (type == MaintenanceType.REWRITE_DATA_FILES + && (metrics.isEmpty() + || !metrics.containsKey(IcebergDataCompactionContent.DATA_FILE_MSE_METRIC) + || !metrics.containsKey(IcebergDataCompactionContent.DELETE_FILE_NUMBER_METRIC))) { + LOG.warn("Compaction recommend skip: metrics JSON missing compaction fields"); + return false; + } + + boolean maintain = evaluateTrigger(rules, metrics); + LOG.info( + "In-Spark recommend type={} expr={} -> {}", + type.wireName(), + rules.get(IcebergDataCompactionContent.TRIGGER_EXPR_KEY), + maintain ? "MAINTAIN" : "SKIP"); + return maintain; + } + + /** + * Same binding rules as {@code CompactionTriggerEvaluator}: numeric policy rules become QL + * bindings, then metric values overlay them, then {@code trigger-expr} is evaluated. + */ + private static boolean evaluateTrigger( + Map<String, Object> rules, Map<String, Object> metricValues) { + Map<String, Object> context = new HashMap<>(); + rules.forEach( + (key, value) -> { + if (value == null) { + return; + } + try { + context.put(key, Long.parseLong(value.toString())); + } catch (NumberFormatException ignored) { + // Non-numeric rule values (expressions, strings) are not QL bindings. + } + }); + context.putAll(metricValues); + + Object trigger = rules.get(IcebergDataCompactionContent.TRIGGER_EXPR_KEY); + String expression = + trigger == null || trigger.toString().trim().isEmpty() + ? "false" + : trigger.toString().trim(); + return EVALUATOR.evaluateBool(expression, context); + } + + private static Map<String, Object> rulesForType( + MaintenanceType type, Map<String, String> argMap) { + Map<String, Map<String, Object>> byType = Collections.emptyMap(); + if (argMap != null) { + String raw = IcebergJobUtils.nullIfUnresolvedPlaceholder(argMap.get(RULES_ARG)); + if (raw != null && !raw.trim().isEmpty()) { + try { + Map<String, Map<String, Object>> parsed = + JsonUtils.anyFieldMapper() + .readValue(raw, new TypeReference<Map<String, Map<String, Object>>>() {}); + if (parsed != null) { + byType = parsed; + } + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to parse {}: {}", RULES_ARG, e.toString()); + } + } + } + Map<String, Object> rules = byType.get(type.wireName()); + if (rules != null && !rules.isEmpty()) { + return rules; + } + if (type == MaintenanceType.REWRITE_DATA_FILES) { + return new HashMap<>(PolicyContents.icebergDataCompaction().rules()); + } + if (type == MaintenanceType.REMOVE_ORPHAN_FILES) { + return new HashMap<>(PolicyContents.icebergOrphanFileRemoval().rules()); + } + Map<String, Object> empty = new HashMap<>(); + empty.put(IcebergOrphanFileRemovalContent.TRIGGER_EXPR_KEY, "false"); + return empty; + } + + private static Map<String, Object> parseNumericMetrics(@Nullable String metricsJson) { + Map<String, Object> metrics = new HashMap<>(); + if (metricsJson == null || metricsJson.trim().isEmpty()) { + return metrics; + } + try { + JsonNode root = JsonUtils.anyFieldMapper().readTree(metricsJson); + Iterator<Entry<String, JsonNode>> fields = root.fields(); + while (fields.hasNext()) { + Entry<String, JsonNode> field = fields.next(); + JsonNode value = field.getValue(); + if (value == null || value.isNull() || !value.isNumber()) { + continue; + } + metrics.put( + field.getKey(), value.isIntegralNumber() ? value.longValue() : value.doubleValue()); + } + } catch (IOException | RuntimeException e) { + LOG.warn("Failed to parse metrics JSON: {}", e.toString()); + } + return metrics; + } +} diff --git a/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestBuiltInJobTemplateProvider.java b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestBuiltInJobTemplateProvider.java index 4ed6c75093..c89e3ef4a9 100644 --- a/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestBuiltInJobTemplateProvider.java +++ b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestBuiltInJobTemplateProvider.java @@ -111,6 +111,12 @@ public class TestBuiltInJobTemplateProvider { expected.put( "builtin-iceberg-rewrite-data-files", ImmutableSet.<String>builder().addAll(icebergCatalog).add("where_clause").build()); + expected.put( + "builtin-iceberg-maintenance-task", + ImmutableSet.<String>builder() + .add("catalog_name", "catalog_type", "catalog_uri", "warehouse_location") + .add("maintenance_types", "policy_ids", "tables") + .build()); for (JobTemplate template : new BuiltInJobTemplateProvider().jobTemplates()) { Set<String> required = new JobTemplateResolver(toEntity(template)).requiredParameters(); diff --git a/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestMaintenanceActionHttpClient.java b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestMaintenanceActionHttpClient.java new file mode 100644 index 0000000000..c7fdf4d8ed --- /dev/null +++ b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/TestMaintenanceActionHttpClient.java @@ -0,0 +1,141 @@ +/* + * 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. + */ +package org.apache.gravitino.maintenance.jobs; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.sun.net.httpserver.HttpServer; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.util.Base64; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionFinishRequest; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionMetricsRequest; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link MaintenanceActionHttpClient}. */ +public class TestMaintenanceActionHttpClient { + + private HttpServer server; + private final AtomicReference<String> auth = new AtomicReference<>(); + private final AtomicReference<String> path = new AtomicReference<>(); + private final AtomicReference<String> body = new AtomicReference<>(); + + @BeforeEach + public void startServer() throws Exception { + server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + server.createContext( + "/", + exchange -> { + auth.set(exchange.getRequestHeaders().getFirst("Authorization")); + path.set(exchange.getRequestURI().getPath()); + body.set(new String(exchange.getRequestBody().readAllBytes(), StandardCharsets.UTF_8)); + byte[] payload = "{\"code\":0}".getBytes(StandardCharsets.UTF_8); + int status = 200; + if (exchange.getRequestURI().getPath().endsWith("/token")) { + payload = + "{\"access_token\":\"tok-1\",\"token_type\":\"bearer\",\"expires_in\":3600}" + .getBytes(StandardCharsets.UTF_8); + } + if (exchange.getRequestURI().getPath().contains("/fail")) { + status = 500; + payload = "{\"code\":-1}".getBytes(StandardCharsets.UTF_8); + } + exchange.sendResponseHeaders(status, payload.length); + exchange.getResponseBody().write(payload); + exchange.close(); + }); + server.start(); + } + + @AfterEach + public void stopServer() { + if (server != null) { + server.stop(0); + } + } + + @Test + public void testSimpleAuthPostsMetricsAndFinish() { + Map<String, String> options = baseOptions(); + options.put("auth_type", "simple"); + options.put("username", "alice"); + try (MaintenanceActionHttpClient client = + MaintenanceActionHttpClient.fromUpdaterOptions(options)) { + client.reportBeforeMetrics(7L, new MaintenanceActionMetricsRequest(1L, 20L, 2L, "{}")); + assertEquals("/api/metalakes/demo/maintenance/actions/7/metrics/before", path.get()); + assertTrue(body.get().contains("\"tableId\":1")); + assertTrue(body.get().contains("\"policyId\":20")); + client.finish(7L, new MaintenanceActionFinishRequest(1L, 20L, 11L)); + assertEquals("/api/metalakes/demo/maintenance/actions/7/finish", path.get()); + } + assertEquals(basic("alice", "dummy"), auth.get()); + } + + @Test + public void testOauthFetchesBearerToken() { + Map<String, String> options = baseOptions(); + options.put("auth_type", "oauth"); + options.put("oauth_server_uri", gravitinoUri()); + options.put("oauth_path", "oauth/token"); + options.put("oauth_credential", "id:secret"); + options.put("oauth_scope", "openid"); + try (MaintenanceActionHttpClient client = + MaintenanceActionHttpClient.fromUpdaterOptions(options)) { + client.reportBeforeMetrics(1L, new MaintenanceActionMetricsRequest(1L, 20L, 2L, "{}")); + } + assertEquals("Bearer tok-1", auth.get()); + assertTrue(path.get().endsWith("/metrics/before")); + } + + @Test + public void testHttpErrorFails() { + Map<String, String> options = baseOptions(); + options.put("gravitino_uri", gravitinoUri() + "/fail"); + MaintenanceActionHttpClient client = MaintenanceActionHttpClient.fromUpdaterOptions(options); + assertThrows( + IllegalStateException.class, + () -> + client.reportBeforeMetrics(1L, new MaintenanceActionMetricsRequest(1L, 20L, 2L, "{}"))); + client.close(); + } + + private Map<String, String> baseOptions() { + Map<String, String> options = new HashMap<>(); + options.put("gravitino_uri", gravitinoUri()); + options.put("metalake", "demo"); + return options; + } + + private String gravitinoUri() { + return "http://127.0.0.1:"; + server.getAddress().getPort(); + } + + private static String basic(String user, String password) { + return "Basic " + + Base64.getEncoder() + .encodeToString((user + ":" + password).getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/iceberg/TestIcebergMaintenanceJob.java b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/iceberg/TestIcebergMaintenanceJob.java new file mode 100644 index 0000000000..f942ca3968 --- /dev/null +++ b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/iceberg/TestIcebergMaintenanceJob.java @@ -0,0 +1,83 @@ +/* + * 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. + */ +package org.apache.gravitino.maintenance.jobs.iceberg; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import org.apache.gravitino.job.SparkJobTemplate; +import org.apache.gravitino.maintenance.jobs.iceberg.IcebergMaintenanceJob.MaintenanceType; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link IcebergMaintenanceJob}. */ +public class TestIcebergMaintenanceJob { + + @Test + public void testTemplatePlaceholders() { + SparkJobTemplate template = new IcebergMaintenanceJob().jobTemplate(); + assertEquals(IcebergMaintenanceJob.NAME, template.name()); + List<String> args = template.arguments(); + assertTrue(args.contains("--maintenance-types")); + assertTrue(args.contains("{{maintenance_types}}")); + assertFalse(args.contains("{{maintenance_types:-}}")); + assertTrue(args.contains("--policy-ids")); + assertTrue(args.contains("{{policy_ids}}")); + assertTrue(args.contains("--tables")); + assertTrue(args.contains("{{tables}}")); + assertTrue(args.contains("--rules")); + assertTrue(args.contains("{{rules:-}}")); + assertFalse(args.contains("--table-id")); + assertFalse(args.contains("--gravitino-job-id")); + } + + @Test + public void testMaintenanceTypeFromWire() { + assertEquals(MaintenanceType.EXPIRE_SNAPSHOTS, MaintenanceType.fromWire("expire_snapshots")); + assertEquals( + MaintenanceType.REWRITE_DATA_FILES, MaintenanceType.fromWire("REWRITE_DATA_FILES")); + assertThrows(IllegalArgumentException.class, () -> MaintenanceType.fromWire("unknown")); + assertThrows(IllegalArgumentException.class, () -> MaintenanceType.fromWire(null)); + } + + @Test + public void testRunRejectsInvalidArgs() { + assertEquals( + 1, IcebergMaintenanceJob.run(new String[] {"--maintenance-types", "expire_snapshots"})); + assertEquals( + 1, + IcebergMaintenanceJob.run( + new String[] { + "--maintenance-types", + "expire_snapshots,rewrite_manifests", + "--policy-ids", + "10", + "--tables", + "[{\"table_id\":1,\"catalog_name\":\"c\",\"table_identifier\":\"db.t\"}]" + })); + assertEquals( + 1, + IcebergMaintenanceJob.run( + new String[] { + "--maintenance-types", "expire_snapshots", "--policy-ids", "10", "--tables", "[]" + })); + } +} diff --git a/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/iceberg/TestMaintenanceTriggerGate.java b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/iceberg/TestMaintenanceTriggerGate.java new file mode 100644 index 0000000000..10abbe887f --- /dev/null +++ b/maintenance/jobs/src/test/java/org/apache/gravitino/maintenance/jobs/iceberg/TestMaintenanceTriggerGate.java @@ -0,0 +1,67 @@ +/* + * 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. + */ +package org.apache.gravitino.maintenance.jobs.iceberg; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.HashMap; +import java.util.Map; +import org.apache.gravitino.maintenance.jobs.iceberg.IcebergMaintenanceJob.MaintenanceType; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link MaintenanceTriggerGate}. */ +public class TestMaintenanceTriggerGate { + + @Test + public void testNonCompactionAlwaysMaintain() { + assertTrue( + MaintenanceTriggerGate.shouldMaintain( + MaintenanceType.EXPIRE_SNAPSHOTS, new HashMap<>(), "{\"snapshotId\":1}")); + } + + @Test + public void testCompactionUsesEnqueueRules() { + Map<String, String> args = new HashMap<>(); + args.put( + MaintenanceTriggerGate.RULES_ARG, + "{\"rewrite_data_files\":{" + + "\"trigger-expr\":\"custom-data-file-mse >= minDataFileMse\"," + + "\"minDataFileMse\":10}}"); + + String below = + "{\"" + + IcebergDataCompactionContent.DATA_FILE_MSE_METRIC + + "\":5,\"" + + IcebergDataCompactionContent.DELETE_FILE_NUMBER_METRIC + + "\":0}"; + assertFalse( + MaintenanceTriggerGate.shouldMaintain(MaintenanceType.REWRITE_DATA_FILES, args, below)); + + String above = + "{\"" + + IcebergDataCompactionContent.DATA_FILE_MSE_METRIC + + "\":11,\"" + + IcebergDataCompactionContent.DELETE_FILE_NUMBER_METRIC + + "\":0}"; + assertTrue( + MaintenanceTriggerGate.shouldMaintain(MaintenanceType.REWRITE_DATA_FILES, args, above)); + } +} diff --git a/maintenance/optimizer-api/src/main/java/org/apache/gravitino/maintenance/optimizer/common/util/GravitinoAuthSettings.java b/maintenance/optimizer-api/src/main/java/org/apache/gravitino/maintenance/optimizer/common/util/GravitinoAuthSettings.java index fb9640f9cc..2180308c69 100644 --- a/maintenance/optimizer-api/src/main/java/org/apache/gravitino/maintenance/optimizer/common/util/GravitinoAuthSettings.java +++ b/maintenance/optimizer-api/src/main/java/org/apache/gravitino/maintenance/optimizer/common/util/GravitinoAuthSettings.java @@ -166,6 +166,48 @@ public final class GravitinoAuthSettings { return normalizedType(); } + /** + * @return simple/basic username, or null when unset + */ + public String username() { + return username; + } + + /** + * @return basic password, or null when unset + */ + public String password() { + return password; + } + + /** + * @return OAuth token server URI, or null when unset + */ + public String oauthServerUri() { + return oauthServerUri; + } + + /** + * @return OAuth token path, or null when unset + */ + public String oauthPath() { + return oauthPath; + } + + /** + * @return OAuth client credential ({@code client_id:secret} or secret), or null when unset + */ + public String oauthCredential() { + return oauthCredential; + } + + /** + * @return OAuth scope, or null when unset + */ + public String oauthScope() { + return oauthScope; + } + private String normalizedType() { if (StringUtils.isBlank(authType)) { return TYPE_NONE;
