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 02c746141ec6c4a19119472503dee2ce2d3a4787 Author: remilu <[email protected]> AuthorDate: Mon Oct 26 07:55:09 2026 +0800 [#14216] feat(tms): Build enqueue-time jobConf specs for cron Fixes #14216. Squashed from fork #75 (B5). --- .../maintenance/policy/MaintenanceAuthConfigs.java | 127 +++++++- .../policy/TestMaintenanceAuthConfigs.java | 49 +++ .../maintenance/policy/PolicyExpandRegistrar.java | 17 +- .../scheduler/AutocommitDataSources.java | 129 ++++++++ .../maintenance/scheduler/MaintenanceJobSpec.java | 349 ++++++++++++++++++++ .../scheduler/MaintenanceJobSpecFactory.java | 285 +++++++++++++++++ .../job/cron/MaintenanceJobConfBuilder.java | 356 +++++++++++++++++++++ .../job/cron/MaintenanceJobTemplates.java | 96 ++++++ .../scheduler/job/cron/MaintenanceTableRefs.java | 169 ++++++++++ .../scheduler/policy/expand/PolicyExpandCrons.java | 32 +- .../scheduler/TestMaintenanceJobSpecFactory.java | 180 +++++++++++ .../job/cron/TestMaintenanceJobConfBuilder.java | 260 +++++++++++++++ .../job/cron/TestMaintenanceJobTemplates.java | 58 ++++ 13 files changed, 2067 insertions(+), 40 deletions(-) diff --git a/core/src/main/java/org/apache/gravitino/maintenance/policy/MaintenanceAuthConfigs.java b/core/src/main/java/org/apache/gravitino/maintenance/policy/MaintenanceAuthConfigs.java index a90451bbd3..d963c74506 100644 --- a/core/src/main/java/org/apache/gravitino/maintenance/policy/MaintenanceAuthConfigs.java +++ b/core/src/main/java/org/apache/gravitino/maintenance/policy/MaintenanceAuthConfigs.java @@ -36,15 +36,19 @@ import org.apache.gravitino.config.ConfigEntry; * Cluster-static TMS client auth in {@code gravitino.conf}. * * <p>{@link #GRAVITINO_AUTH_PREFIX} is a <strong>fixed</strong> camelCase bag in {@code - * gravitino.conf} ({@code authType}, {@code oauthServerUri}, …). {@link #loadGravitinoAuth} maps - * those suffixes onto built-in updater short names ({@code auth_type}, {@code oauth_server_uri}, …) - * before overlaying onto policy {@code updaterOptions}. Unknown or blank suffixes are rejected, and - * {@code authType} plus mode-required fields are validated at load time. + * gravitino.conf} ({@code uri}, {@code authType}, {@code oauthServerUri}, …), symmetric with {@link + * #IRC_AUTH_PREFIX}. {@link #loadGravitinoAuth} maps those suffixes onto built-in updater short + * names ({@code gravitino_uri}, {@code auth_type}, {@code oauth_server_uri}, …) before overlaying + * onto policy {@code updaterOptions}. Unknown or blank suffixes are rejected, and {@code authType} + * plus mode-required fields are validated at load time ({@code uri} alone is allowed without {@code + * authType}). * * <p>{@link #IRC_AUTH_PREFIX} is an <strong>open</strong> passthrough for Spark / Iceberg REST * properties (for example {@code rest.auth.type}); any non-empty suffix under the prefix is loaded - * and overlaid onto policy {@code jobOptions}. Auth wins on duplicate keys. Live {@code runJob} - * wiring is a follow-up. + * and overlaid onto policy {@code jobOptions} with suffixes unchanged (for example {@code + * rest.auth.type}). Cron {@code runJob} later expands those short names into {@code spark_conf} + * under {@code spark.sql.catalog.<catalog>.*}; Spark resource keys such as {@code + * spark.executor.memory} stay as Spark settings, not catalog properties. */ public final class MaintenanceAuthConfigs { @@ -58,6 +62,12 @@ public final class MaintenanceAuthConfigs { */ public static final String IRC_AUTH_PREFIX = "gravitino.maintenance.ircAuth."; + /** + * Base URI for Spark → Gravitino TMS callbacks ({@code updater_options.gravitino_uri}). Under + * {@link #GRAVITINO_AUTH_PREFIX} for symmetry with {@link #IRC_AUTH_PREFIX}. + */ + public static final String URI_SUFFIX = "uri"; + /** Auth mode: {@code none}, {@code simple}, {@code basic}, or {@code oauth}. */ public static final String AUTH_TYPE_SUFFIX = "authType"; @@ -93,6 +103,7 @@ public final class MaintenanceAuthConfigs { */ private static final Map<String, String> GRAVITINO_AUTH_TO_UPDATER_OPTIONS = ImmutableMap.<String, String>builder() + .put(URI_SUFFIX, "gravitino_uri") .put(AUTH_TYPE_SUFFIX, "auth_type") .put(USERNAME_SUFFIX, "username") .put(PASSWORD_SUFFIX, "password") @@ -102,6 +113,24 @@ public final class MaintenanceAuthConfigs { .put(OAUTH_SCOPE_SUFFIX, "oauth_scope") .build(); + /** + * Updater-option short names that belong to the Gravitino auth bag (after camelCase → short-name + * remap). + */ + public static final Set<String> GRAVITINO_AUTH_UPDATER_OPTION_KEYS = + ImmutableSet.copyOf(GRAVITINO_AUTH_TO_UPDATER_OPTIONS.values()); + + /** + * Cluster-static base URI for TMS sample/finish callbacks. Policy {@code updaterOptions} may + * override via {@code gravitino_uri}. + */ + public static final ConfigEntry<Optional<String>> GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI = + optionalAuth( + GRAVITINO_AUTH_PREFIX, + URI_SUFFIX, + "Base URI for Spark → Gravitino TMS callbacks (sample/finish). Injected into cron " + + "updater_options when the policy does not set gravitino_uri."); + public static final ConfigEntry<Optional<String>> GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_TYPE = optionalAuth( GRAVITINO_AUTH_PREFIX, @@ -165,6 +194,30 @@ public final class MaintenanceAuthConfigs { return ImmutableMap.copyOf(updaterOptions); } + /** + * Like {@link #loadGravitinoAuth} but values are the full {@code gravitino.conf} keys (not secret + * values). For enqueue {@code task_data}; resolve with {@link Config#getRawString(String)} at + * submit. + * + * @param config gravitino server config + * @return updater short name → conf key; empty when unset + */ + public static Map<String, String> loadGravitinoAuthAsConfigKeys(Config config) { + Map<String, String> loaded = loadAuthBag(config, GRAVITINO_AUTH_PREFIX, /*fixedSchema=*/ true); + if (loaded.isEmpty()) { + return Collections.emptyMap(); + } + validateGravitinoAuth(loaded); + Map<String, String> keys = new LinkedHashMap<>(); + GRAVITINO_AUTH_TO_UPDATER_OPTIONS.forEach( + (suffix, updateKey) -> { + if (loaded.containsKey(suffix)) { + keys.put(updateKey, GRAVITINO_AUTH_PREFIX + suffix); + } + }); + return ImmutableMap.copyOf(keys); + } + /** * Loads arbitrary IRC / Spark property suffixes under {@link #IRC_AUTH_PREFIX}. * @@ -175,6 +228,42 @@ public final class MaintenanceAuthConfigs { return loadAuthBag(config, IRC_AUTH_PREFIX, /*fixedSchema=*/ false); } + /** + * Like {@link #loadIrcAuth} but values are full conf keys under {@link #IRC_AUTH_PREFIX}. + * + * @param config gravitino server config + * @return IRC short suffix → conf key; empty when unset + */ + public static Map<String, String> loadIrcAuthAsConfigKeys(Config config) { + Map<String, String> loaded = loadAuthBag(config, IRC_AUTH_PREFIX, /*fixedSchema=*/ false); + if (loaded.isEmpty()) { + return Collections.emptyMap(); + } + Map<String, String> keys = new LinkedHashMap<>(); + loaded.forEach((suffix, ignored) -> keys.put(suffix, IRC_AUTH_PREFIX + suffix)); + return ImmutableMap.copyOf(keys); + } + + /** + * Resolves a value that may be a {@code gravitino.maintenance.*} conf key into the live config + * value. Non-key literals are returned unchanged. + * + * @param config server config + * @param value literal or conf key + * @return resolved value + */ + public static String resolveConfigReference(Config config, @Nullable String value) { + if (value == null || !value.startsWith("gravitino.maintenance.")) { + return value; + } + String raw = config.getRawString(value); + if (raw == null || raw.trim().isEmpty()) { + throw new IllegalArgumentException( + "Missing gravitino.conf value for referenced key: " + value); + } + return raw.trim(); + } + /** * Merges policy {@code jobOptions} with IRC auth passthrough; auth wins on duplicate keys. * @@ -188,7 +277,8 @@ public final class MaintenanceAuthConfigs { } /** - * Merges policy {@code updaterOptions} with Gravitino auth; auth wins on duplicate keys. + * Merges policy {@code updaterOptions} with Gravitino auth. Auth wins on credential keys; policy + * {@code gravitino_uri} wins over cluster {@code gravitinoAuth.uri} when both are set. * * @param policyUpdaterOptions options from the effective policy (may be null) * @param gravitinoAuth updater-short-name Gravitino auth bag (may be null) @@ -197,7 +287,24 @@ public final class MaintenanceAuthConfigs { public static Map<String, String> overlayUpdaterOptions( @Nullable Map<String, String> policyUpdaterOptions, @Nullable Map<String, String> gravitinoAuth) { - return overlay(policyUpdaterOptions, gravitinoAuth); + if ((policyUpdaterOptions == null || policyUpdaterOptions.isEmpty()) + && (gravitinoAuth == null || gravitinoAuth.isEmpty())) { + return Collections.emptyMap(); + } + Map<String, String> merged = new LinkedHashMap<>(); + if (policyUpdaterOptions != null) { + merged.putAll(policyUpdaterOptions); + } + if (gravitinoAuth != null) { + for (Map.Entry<String, String> entry : gravitinoAuth.entrySet()) { + String key = entry.getKey(); + if ("gravitino_uri".equals(key) && merged.containsKey("gravitino_uri")) { + continue; + } + merged.put(key, entry.getValue()); + } + } + return ImmutableMap.copyOf(merged); } /** @@ -256,6 +363,10 @@ public final class MaintenanceAuthConfigs { private static void validateGravitinoAuth(Map<String, String> loaded) { String authType = loaded.get(AUTH_TYPE_SUFFIX); if (authType == null) { + // uri alone is connection target, not a credential mode — allow without authType. + if (loaded.keySet().stream().allMatch(URI_SUFFIX::equals)) { + return; + } throw new IllegalArgumentException( String.format( "Configuration %s%s is required when any Gravitino auth setting is present", diff --git a/core/src/test/java/org/apache/gravitino/maintenance/policy/TestMaintenanceAuthConfigs.java b/core/src/test/java/org/apache/gravitino/maintenance/policy/TestMaintenanceAuthConfigs.java index fed3289e83..7e0e957a4d 100644 --- a/core/src/test/java/org/apache/gravitino/maintenance/policy/TestMaintenanceAuthConfigs.java +++ b/core/src/test/java/org/apache/gravitino/maintenance/policy/TestMaintenanceAuthConfigs.java @@ -41,6 +41,9 @@ class TestMaintenanceAuthConfigs { Assertions.assertEquals( "gravitino.maintenance.gravitinoAuth.oauthServerUri", MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_OAUTH_SERVER_URI.getKey()); + Assertions.assertEquals( + "gravitino.maintenance.gravitinoAuth.uri", + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey()); } @Test @@ -69,6 +72,8 @@ class TestMaintenanceAuthConfigs { Config config = new Config(false) {}; config.loadFromMap( Map.of( + MaintenanceAuthConfigs.GRAVITINO_AUTH_PREFIX + "uri", + "http://gravitino.example:8090";, MaintenanceAuthConfigs.GRAVITINO_AUTH_PREFIX + "authType", "oauth", MaintenanceAuthConfigs.GRAVITINO_AUTH_PREFIX + "username", @@ -87,6 +92,7 @@ class TestMaintenanceAuthConfigs { Assertions.assertEquals( ImmutableMap.<String, String>builder() + .put("gravitino_uri", "http://gravitino.example:8090";) .put("auth_type", "oauth") .put("username", "tms") .put("password", "secret") @@ -98,6 +104,49 @@ class TestMaintenanceAuthConfigs { MaintenanceAuthConfigs.loadGravitinoAuth(config)); } + @Test + void loadGravitinoAuthAllowsUriWithoutAuthType() { + Config config = new Config(false) {}; + config.loadFromMap( + Map.of( + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey(), + "http://gravitino.example:8090";), + key -> key.startsWith("gravitino.")); + + Assertions.assertEquals( + ImmutableMap.of("gravitino_uri", "http://gravitino.example:8090";), + MaintenanceAuthConfigs.loadGravitinoAuth(config)); + } + + @Test + void loadAuthAsConfigKeysStoresKeysNotValues() { + Config config = new Config(false) {}; + config.loadFromMap( + Map.of( + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey(), + "http://gravitino.example:8090";, + MaintenanceAuthConfigs.GRAVITINO_AUTH_PREFIX + "authType", + "simple", + MaintenanceAuthConfigs.IRC_AUTH_PREFIX + "rest.auth.type", + "oauth2"), + key -> key.startsWith("gravitino.")); + + Assertions.assertEquals( + ImmutableMap.of( + "gravitino_uri", + "gravitino.maintenance.gravitinoAuth.uri", + "auth_type", + "gravitino.maintenance.gravitinoAuth.authType"), + MaintenanceAuthConfigs.loadGravitinoAuthAsConfigKeys(config)); + Assertions.assertEquals( + ImmutableMap.of("rest.auth.type", "gravitino.maintenance.ircAuth.rest.auth.type"), + MaintenanceAuthConfigs.loadIrcAuthAsConfigKeys(config)); + Assertions.assertEquals( + "http://gravitino.example:8090";, + MaintenanceAuthConfigs.resolveConfigReference( + config, "gravitino.maintenance.gravitinoAuth.uri")); + } + @Test void loadIrcAuthPassesThroughArbitrarySparkKeys() { Config config = new Config(false) {}; diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java index 80b766fd37..2ee07368f3 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java @@ -41,6 +41,7 @@ import java.util.Optional; import java.util.Set; import javax.annotation.Nullable; import javax.sql.DataSource; +import org.apache.gravitino.maintenance.scheduler.AutocommitDataSources; import org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools; import org.apache.gravitino.storage.relational.session.SqlSessions; import org.slf4j.Logger; @@ -90,13 +91,18 @@ public final class PolicyExpandRegistrar { } /** - * Opens a client that reads committed rows from the entity-store pool. + * Opens a client against the entity-store pool with JDBC autocommit forced on. + * + * <p>Used for post-commit reads and test-only writes such as {@link #scheduleAt}. The + * entity-store pool often disables autocommit; without this wrap, SchedulerClient mutations would + * not commit and the expand pool would never observe them. * * @param dataSource relational JDBC pool shared with the entity store - * @return registrar that can read expand rows after the policy transaction commits + * @return registrar that can read (and test-write) expand rows after the policy transaction + * commits */ public static PolicyExpandRegistrar forReads(DataSource dataSource) { - return new PolicyExpandRegistrar(dataSource); + return new PolicyExpandRegistrar(AutocommitDataSources.wrap(dataSource)); } private SchedulerClient client() { @@ -130,7 +136,7 @@ public final class PolicyExpandRegistrar { * @param due execution time to store */ @VisibleForTesting - void scheduleAt(long policyId, Instant due) { + public void scheduleAt(long policyId, Instant due) { runDespiteExecution(() -> scheduleOnce(policyId, due)); } @@ -147,7 +153,8 @@ public final class PolicyExpandRegistrar { * @param policyId persisted policy id * @return next execution time of the expand row, or empty when the row is absent */ - Optional<Instant> scheduledExecutionTime(long policyId) { + @VisibleForTesting + public Optional<Instant> scheduledExecutionTime(long policyId) { return client() .getScheduledExecution(instanceId(policyId)) .map(execution -> execution.getExecutionTime()); diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/AutocommitDataSources.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/AutocommitDataSources.java new file mode 100644 index 0000000000..c8d12e22bb --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/AutocommitDataSources.java @@ -0,0 +1,129 @@ +/* + * 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.scheduler; + +import java.io.PrintWriter; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import javax.sql.DataSource; + +/** + * Wraps the entity-store {@link DataSource} so db-scheduler {@code SchedulerClient} calls commit + * even when the pool disables JDBC autocommit. + */ +public final class AutocommitDataSources { + + private AutocommitDataSources() {} + + /** + * @param delegate entity-store pool + * @return proxy that forces {@code autoCommit=true} on borrowed connections + */ + public static DataSource wrap(DataSource delegate) { + return new AutocommitDataSource(delegate); + } + + private static final class AutocommitDataSource implements DataSource { + + private final DataSource delegate; + + AutocommitDataSource(DataSource delegate) { + this.delegate = delegate; + } + + @Override + public Connection getConnection() throws SQLException { + return withAutocommit(delegate.getConnection()); + } + + @Override + public Connection getConnection(String username, String password) throws SQLException { + return withAutocommit(delegate.getConnection(username, password)); + } + + @Override + public PrintWriter getLogWriter() throws SQLException { + return delegate.getLogWriter(); + } + + @Override + public void setLogWriter(PrintWriter out) throws SQLException { + delegate.setLogWriter(out); + } + + @Override + public void setLoginTimeout(int seconds) throws SQLException { + delegate.setLoginTimeout(seconds); + } + + @Override + public int getLoginTimeout() throws SQLException { + return delegate.getLoginTimeout(); + } + + @Override + public java.util.logging.Logger getParentLogger() throws SQLFeatureNotSupportedException { + return delegate.getParentLogger(); + } + + @Override + public <T> T unwrap(Class<T> iface) throws SQLException { + if (iface.isInstance(this)) { + return iface.cast(this); + } + return delegate.unwrap(iface); + } + + @Override + public boolean isWrapperFor(Class<?> iface) throws SQLException { + return iface.isInstance(this) || delegate.isWrapperFor(iface); + } + + private static Connection withAutocommit(Connection connection) throws SQLException { + Connection proxy = + (Connection) + Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class<?>[] {Connection.class}, + suppressingSetAutocommit(connection)); + connection.setAutoCommit(true); + return proxy; + } + + private static InvocationHandler suppressingSetAutocommit(Connection connection) { + return (proxy, method, args) -> { + if ("setAutoCommit".equals(method.getName())) { + return null; + } + if ("getAutoCommit".equals(method.getName())) { + return true; + } + try { + return method.invoke(connection, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + }; + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/MaintenanceJobSpec.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/MaintenanceJobSpec.java new file mode 100644 index 0000000000..516d70b3ea --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/MaintenanceJobSpec.java @@ -0,0 +1,349 @@ +/* + * 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.scheduler; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.google.common.base.Preconditions; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; +import javax.annotation.Nullable; + +/** + * Enqueue-time spec stored as JSON in db-scheduler {@code task_data} for {@code cron-job} and + * {@code commit-job} rows. + * + * <p>Thin envelope only: {@link #jobConf} is the same {@code Map<String,String>} passed to {@code + * runJob} (including {@code updater_options} as a JSON string). Tables carry {@code tableId} plus + * NameIdentifier levels (metalake first); operations carry {@code policyId} plus the wire + * maintenance type. Metalake for {@code runJob} is derived from {@link #tables}. + */ +@JsonIgnoreProperties(ignoreUnknown = true) +public final class MaintenanceJobSpec { + + private List<TableTarget> tables = new ArrayList<>(); + private List<Operation> operations = new ArrayList<>(); + private Map<String, String> jobConf = new LinkedHashMap<>(); + + /** Jackson / empty constructor. */ + public MaintenanceJobSpec() {} + + /** + * @param tables table targets (order preserved) + * @param operations maintenance operations (order preserved; cron: one; commit: chain) + * @param jobConf runJob configuration map + */ + public MaintenanceJobSpec( + List<TableTarget> tables, List<Operation> operations, Map<String, String> jobConf) { + this.tables = new ArrayList<>(Preconditions.checkNotNull(tables, "tables")); + this.operations = new ArrayList<>(Preconditions.checkNotNull(operations, "operations")); + this.jobConf = new LinkedHashMap<>(Preconditions.checkNotNull(jobConf, "jobConf")); + } + + /** + * @return table targets + */ + public List<TableTarget> getTables() { + return tables; + } + + /** + * @param tables table targets + */ + public void setTables(List<TableTarget> tables) { + this.tables = tables == null ? new ArrayList<>() : new ArrayList<>(tables); + } + + /** + * @return maintenance operations + */ + public List<Operation> getOperations() { + return operations; + } + + /** + * @param operations maintenance operations + */ + public void setOperations(List<Operation> operations) { + this.operations = operations == null ? new ArrayList<>() : new ArrayList<>(operations); + } + + /** + * @return runJob configuration + */ + public Map<String, String> getJobConf() { + return jobConf; + } + + /** + * @param jobConf runJob configuration + */ + public void setJobConf(Map<String, String> jobConf) { + this.jobConf = jobConf == null ? new LinkedHashMap<>() : new LinkedHashMap<>(jobConf); + } + + /** + * @return unmodifiable table entity ids in spec order + */ + @JsonIgnore + public List<Long> tableIds() { + if (tables == null || tables.isEmpty()) { + return List.of(); + } + return Collections.unmodifiableList( + tables.stream().map(TableTarget::getTableId).collect(Collectors.toList())); + } + + /** + * Metalake for {@code runJob} / cancel, taken from the first table's NameIdentifier. + * + * @return metalake name, or {@code null} when tables are missing / incomplete + */ + @JsonIgnore + @Nullable + public String metalake() { + if (tables == null || tables.isEmpty()) { + return null; + } + List<String> identifier = tables.get(0).getIdentifier(); + if (identifier == null || identifier.isEmpty()) { + return null; + } + String name = identifier.get(0); + return name == null || name.isBlank() ? null : name; + } + + /** + * @return true when tables, operations, and jobConf look usable for submit + */ + @JsonIgnore + public boolean isValid() { + return metalake() != null + && tables != null + && !tables.isEmpty() + && operations != null + && !operations.isEmpty() + && jobConf != null + && !jobConf.isEmpty() + && operations.get(0).getPolicyId() > 0 + && tables.stream().allMatch(TableTarget::isValid) + && operations.stream().allMatch(Operation::isValid); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof MaintenanceJobSpec)) { + return false; + } + MaintenanceJobSpec that = (MaintenanceJobSpec) o; + return Objects.equals(tables, that.tables) + && Objects.equals(operations, that.operations) + && Objects.equals(jobConf, that.jobConf); + } + + @Override + public int hashCode() { + return Objects.hash(tables, operations, jobConf); + } + + @Override + public String toString() { + return "MaintenanceJobSpec{tables=" + + tableIds() + + ", operations=" + + (operations == null ? 0 : operations.size()) + + ", jobConfKeys=" + + (jobConf == null ? 0 : jobConf.size()) + + '}'; + } + + /** One table covered by this spec. */ + @JsonIgnoreProperties(ignoreUnknown = true) + public static final class TableTarget { + + private long tableId; + private List<String> identifier = new ArrayList<>(); + + /** Jackson constructor. */ + public TableTarget() {} + + /** + * @param tableId table entity id + * @param identifier NameIdentifier levels (metalake, catalog, …, table) + */ + public TableTarget(long tableId, List<String> identifier) { + Preconditions.checkArgument(tableId > 0, "tableId must be positive"); + this.tableId = tableId; + this.identifier = new ArrayList<>(Preconditions.checkNotNull(identifier, "identifier")); + } + + /** + * @return table entity id + */ + public long getTableId() { + return tableId; + } + + /** + * @param tableId table entity id + */ + public void setTableId(long tableId) { + this.tableId = tableId; + } + + /** + * @return NameIdentifier levels + */ + public List<String> getIdentifier() { + return identifier; + } + + /** + * @param identifier NameIdentifier levels + */ + public void setIdentifier(List<String> identifier) { + this.identifier = identifier == null ? new ArrayList<>() : new ArrayList<>(identifier); + } + + /** + * @return Spark catalog name ({@code identifier[1]}) + */ + @JsonIgnore + @Nullable + public String catalogName() { + return identifier != null && identifier.size() >= 2 ? identifier.get(1) : null; + } + + /** + * @return {@code schema.table} (or nested) Spark table identifier + */ + @JsonIgnore + @Nullable + public String sparkTableIdentifier() { + if (identifier == null || identifier.size() < 3) { + return null; + } + return String.join(".", identifier.subList(2, identifier.size())); + } + + @JsonIgnore + boolean isValid() { + return tableId > 0 + && identifier != null + && identifier.size() >= 3 + && identifier.stream().noneMatch(part -> part == null || part.isBlank()); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof TableTarget)) { + return false; + } + TableTarget that = (TableTarget) o; + return tableId == that.tableId && Objects.equals(identifier, that.identifier); + } + + @Override + public int hashCode() { + return Objects.hash(tableId, identifier); + } + } + + /** One maintenance operation (cron: single type; commit: chain member). */ + @JsonIgnoreProperties(ignoreUnknown = true) + public static final class Operation { + + private long policyId; + private String maintenanceType; + + /** Jackson constructor. */ + public Operation() {} + + /** + * @param policyId policy entity id + * @param maintenanceType Spark wire name + */ + public Operation(long policyId, String maintenanceType) { + Preconditions.checkArgument(policyId > 0, "policyId must be positive"); + this.policyId = policyId; + this.maintenanceType = Preconditions.checkNotNull(maintenanceType, "maintenanceType"); + } + + /** + * @return policy entity id + */ + public long getPolicyId() { + return policyId; + } + + /** + * @param policyId policy entity id + */ + public void setPolicyId(long policyId) { + this.policyId = policyId; + } + + /** + * @return Spark wire maintenance type + */ + public String getMaintenanceType() { + return maintenanceType; + } + + /** + * @param maintenanceType Spark wire maintenance type + */ + public void setMaintenanceType(String maintenanceType) { + this.maintenanceType = maintenanceType; + } + + @JsonIgnore + boolean isValid() { + return policyId > 0 && maintenanceType != null && !maintenanceType.isBlank(); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Operation)) { + return false; + } + Operation that = (Operation) o; + return policyId == that.policyId && Objects.equals(maintenanceType, that.maintenanceType); + } + + @Override + public int hashCode() { + return Objects.hash(policyId, maintenanceType); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/MaintenanceJobSpecFactory.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/MaintenanceJobSpecFactory.java new file mode 100644 index 0000000000..37c55059e0 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/MaintenanceJobSpecFactory.java @@ -0,0 +1,285 @@ +/* + * 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.scheduler; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.Operation; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.TableTarget; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobConfBuilder; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobTemplates; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceTableRefs.ResolvedTableRef; +import org.apache.gravitino.meta.PolicyEntity; + +/** + * Builds enqueue-time {@link MaintenanceJobSpec} values for the cron path and shared jobConf + * helpers used by cron pick / commit merge. + */ +public final class MaintenanceJobSpecFactory { + + private static final Pattern RESOURCE_NUMBER = Pattern.compile("^(\\d+)"); + private static final Pattern MEMORY_VALUE = + Pattern.compile("^(\\d+(?:\\.\\d+)?)\\s*([kmgt]?i?b?)$", Pattern.CASE_INSENSITIVE); + + private final MaintenanceJobConfBuilder jobConfBuilder; + + /** + * @param gravitinoConfig server config for auth overlay and gravitino URI + */ + public MaintenanceJobSpecFactory(Config gravitinoConfig) { + this.jobConfBuilder = + new MaintenanceJobConfBuilder(Preconditions.checkNotNull(gravitinoConfig)); + } + + /** + * Cron spec for one or more tables under a single policy. Shared {@code jobConf} omits per-table + * keys; pick overlays {@code catalog_name} / {@code table_identifier} for each {@link + * TableTarget}. + * + * <p>All tables must share the same Spark catalog name so catalog-scoped {@code spark_conf} keys + * stay correct for every member of a batch. + * + * @param policy effective maintenance policy + * @param tables resolved table refs (same order as enqueue) + * @return spec ready for {@code task_data} + */ + public MaintenanceJobSpec forCron(PolicyEntity policy, List<ResolvedTableRef> tables) { + Preconditions.checkNotNull(policy, "policy"); + Preconditions.checkArgument(tables != null && !tables.isEmpty(), "tables must not be empty"); + Optional<String> wire = MaintenanceJobTemplates.maintenanceType(policy.policyType()); + Preconditions.checkArgument( + wire.isPresent(), "unsupported policy type %s", policy.policyType()); + + String catalogName = tables.get(0).catalogName(); + for (ResolvedTableRef ref : tables) { + Preconditions.checkArgument( + catalogName.equals(ref.catalogName()), + "cron batch tables must share catalogName; got %s and %s", + catalogName, + ref.catalogName()); + } + + String metalake = policy.namespace().level(0); + Map<String, String> jobConf = jobConfBuilder.build(policy, catalogName); + // Per-table keys are applied at pick so one batch spec can cover many tables. + jobConf.remove("catalog_name"); + jobConf.remove("table_identifier"); + + List<TableTarget> targets = new ArrayList<>(tables.size()); + for (ResolvedTableRef ref : tables) { + targets.add(toTarget(metalake, ref)); + } + List<Operation> operations = List.of(new Operation(policy.id(), wire.get())); + return new MaintenanceJobSpec(targets, operations, jobConf); + } + + /** + * Copies {@code shared} and overlays catalog plus a one-element {@code tables} JSON for submit. + * + * <p>Entity {@code table_id} stays inside {@code tables} (system-filled). It is not a top-level + * template placeholder — users submitting the built-in template do not know entity ids. + * + * @param shared jobConf + * @param target table to submit + * @return mutable per-table jobConf + */ + public static Map<String, String> jobConfForTable( + Map<String, String> shared, TableTarget target) { + Preconditions.checkNotNull(shared, "shared"); + Preconditions.checkNotNull(target, "target"); + String catalogName = target.catalogName(); + String tableIdentifier = target.sparkTableIdentifier(); + Preconditions.checkArgument( + catalogName != null && tableIdentifier != null, + "table target %s is missing catalog/table identifier levels", + target.getTableId()); + Map<String, String> jobConf = new LinkedHashMap<>(shared); + jobConf.put("catalog_name", catalogName); + jobConf.put("tables", toTablesJson(List.of(target))); + jobConf.remove("table_identifier"); + jobConf.remove("table_id"); + return jobConf; + } + + private static String toTablesJson(List<TableTarget> tables) { + List<Map<String, Object>> payload = new ArrayList<>(tables.size()); + for (TableTarget target : tables) { + String catalogName = target.catalogName(); + String tableIdentifier = target.sparkTableIdentifier(); + Preconditions.checkArgument( + catalogName != null && tableIdentifier != null, + "table target %s is missing catalog/table identifier levels", + target.getTableId()); + Map<String, Object> entry = new LinkedHashMap<>(); + entry.put("table_id", target.getTableId()); + entry.put("catalog_name", catalogName); + entry.put("table_identifier", tableIdentifier); + payload.add(entry); + } + try { + return JsonUtils.anyFieldMapper().writeValueAsString(payload); + } catch (JsonProcessingException e) { + throw new IllegalStateException("Failed to serialize tables for runJob", e); + } + } + + /** + * For resource-shaped keys ({@code *memory*}, {@code *cores*}, {@code *instances*}), keep the + * larger value. Memory keys compare normalized byte sizes ({@code 1g} > {@code 512m}); cores / + * instances compare the leading integer. Other keys from the first policy are left unchanged. + * + * @param into baseline jobConf (mutated) + * @param from candidate jobConf + */ + public static void mergeResourceKeys(Map<String, String> into, Map<String, String> from) { + for (Map.Entry<String, String> entry : from.entrySet()) { + String key = entry.getKey(); + String value = entry.getValue(); + if (key == null || value == null || !isResourceKey(key)) { + continue; + } + String existing = into.get(key); + if (existing == null) { + into.put(key, value); + continue; + } + Optional<Long> existingNum = resourceMagnitude(key, existing); + Optional<Long> candidateNum = resourceMagnitude(key, value); + if (existingNum.isPresent() + && candidateNum.isPresent() + && candidateNum.get() > existingNum.get()) { + into.put(key, value); + } + } + } + + /** + * @param metalake metalake name + * @param ref resolved Spark table ref + * @return table target with NameIdentifier levels + */ + public static TableTarget toTarget(String metalake, ResolvedTableRef ref) { + List<String> levels = new ArrayList<>(); + levels.add(metalake); + levels.add(ref.catalogName()); + levels.addAll(Arrays.asList(ref.tableIdentifier().split("\\."))); + return new TableTarget(ref.tableId(), levels); + } + + private static boolean isResourceKey(String key) { + String lower = key.toLowerCase(Locale.ROOT); + return lower.contains("memory") || lower.contains("cores") || lower.contains("instances"); + } + + private static Optional<Long> resourceMagnitude(String key, String value) { + if (key.toLowerCase(Locale.ROOT).contains("memory")) { + return parseMemoryBytes(value); + } + return leadingLong(value); + } + + /** + * Parses Spark-style memory strings ({@code 512m}, {@code 1g}, {@code 2048}) into bytes for max + * merge. + */ + @VisibleForTesting + static Optional<Long> parseMemoryBytes(String value) { + if (value == null) { + return Optional.empty(); + } + Matcher matcher = MEMORY_VALUE.matcher(value.trim()); + if (!matcher.matches()) { + return leadingLong(value); + } + try { + double amount = Double.parseDouble(matcher.group(1)); + String unit = matcher.group(2) == null ? "" : matcher.group(2).toLowerCase(Locale.ROOT); + long multiplier; + switch (unit) { + case "": + case "b": + multiplier = 1L; + break; + case "k": + case "kb": + case "ki": + case "kib": + multiplier = 1024L; + break; + case "m": + case "mb": + case "mi": + case "mib": + multiplier = 1024L * 1024L; + break; + case "g": + case "gb": + case "gi": + case "gib": + multiplier = 1024L * 1024L * 1024L; + break; + case "t": + case "tb": + case "ti": + case "tib": + multiplier = 1024L * 1024L * 1024L * 1024L; + break; + default: + return Optional.empty(); + } + return Optional.of(Math.round(amount * multiplier)); + } catch (NumberFormatException e) { + return Optional.empty(); + } + } + + private static Optional<Long> leadingLong(String value) { + Matcher matcher = RESOURCE_NUMBER.matcher(value.trim()); + if (!matcher.find()) { + return Optional.empty(); + } + try { + return Optional.of(Long.parseLong(matcher.group(1))); + } catch (NumberFormatException e) { + return Optional.empty(); + } + } + + /** Visible for tests that assert resource merge. */ + @VisibleForTesting + public static Map<String, String> newLinkedMap(String... kv) { + Map<String, String> map = new LinkedHashMap<>(); + for (int i = 0; i + 1 < kv.length; i += 2) { + map.put(kv[i], kv[i + 1]); + } + return map; + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceJobConfBuilder.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceJobConfBuilder.java new file mode 100644 index 0000000000..1fc53dd5b9 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceJobConfBuilder.java @@ -0,0 +1,356 @@ +/* + * 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.scheduler.job.cron; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import java.time.Clock; +import java.time.DateTimeException; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.time.temporal.ChronoUnit; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.TreeMap; +import javax.annotation.Nullable; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.optimizer.common.util.IcebergSparkConfigUtils; +import org.apache.gravitino.maintenance.policy.MaintenanceAuthConfigs; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.policy.ExecutablePolicyContent; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.IcebergOrphanFileRemovalContent; +import org.apache.gravitino.policy.PolicyContent; +import org.apache.gravitino.policy.PolicyExecutionSpec; + +/** + * Builds the {@code runJob} configuration map for TMS Spark submits. + * + * <p>Shape matches Iceberg update-stats submit: flatten template placeholders into the top-level + * map, pack {@code updater_options} as one JSON string, and pack extra Spark / IRC settings into + * {@code spark_conf} as another JSON string. Policy {@code jobOptions} (including overlay from + * {@code gravitino.maintenance.ircAuth.*}) keep short IRC keys such as {@code rest.auth.type}; + * those are expanded to {@code spark.sql.catalog.<catalog>.*} only when writing {@code spark_conf}. + * Dot-form Spark resource keys ({@code spark.executor.memory}) become the matching underscore + * placeholders ({@code spark_executor_memory}) instead of catalog properties. + * + * <p>{@link #build} is for enqueue {@code task_data}: cluster {@code gravitinoAuth} / {@code + * ircAuth} are stored as conf <em>keys</em> (not secret values). Call {@link #resolveForSubmit} + * immediately before {@code runJob} to expand those keys from the live {@code gravitino.conf}. + */ +public final class MaintenanceJobConfBuilder { + + private static final DateTimeFormatter ORPHAN_CUTOFF = + DateTimeFormatter.ofPattern("uuuu-MM-dd HH:mm:ssXXX").withZone(ZoneOffset.UTC); + + private static final String SPARK_SQL_CATALOG_PREFIX = "spark.sql.catalog."; + + /** Top-level {@code runJob} placeholders consumed by Iceberg / TMS Spark templates. */ + private static final Set<String> FLAT_PLACEHOLDER_KEYS = + ImmutableSet.of( + "catalog_name", + "table_identifier", + "catalog_type", + "catalog_uri", + "warehouse_location", + "spark_master", + "spark_executor_instances", + "spark_executor_cores", + "spark_executor_memory", + "spark_driver_memory", + "strategy", + "sort_order", + "where_clause", + "options", + "use_caching", + "spec_id", + "older_than", + "retain_last", + "stream_results", + "location", + "dry_run", + MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY, + MaintenanceJobTemplates.RULES_KEY, + "gravitino_job_id"); + + /** + * Spark resource / session keys that map onto flat template placeholders (same as update-stats). + */ + private static final Map<String, String> SPARK_KEY_TO_PLACEHOLDER = + ImmutableMap.of( + "spark.master", + "spark_master", + "spark.executor.instances", + "spark_executor_instances", + "spark.executor.cores", + "spark_executor_cores", + "spark.executor.memory", + "spark_executor_memory", + "spark.driver.memory", + "spark_driver_memory"); + + private final Config gravitinoConfig; + private final Clock clock; + + /** + * @param gravitinoConfig server config for auth overlay and gravitino URI + */ + public MaintenanceJobConfBuilder(Config gravitinoConfig) { + this(gravitinoConfig, Clock.systemUTC()); + } + + MaintenanceJobConfBuilder(Config gravitinoConfig, Clock clock) { + this.gravitinoConfig = gravitinoConfig; + this.clock = clock; + } + + /** + * Enqueue-time jobConf for db-scheduler {@code task_data}. Cluster auth overlays store conf keys + * (for example {@code gravitino.maintenance.gravitinoAuth.uri}), not secret values. + * + * @param policy effective maintenance policy + * @param catalogName Spark catalog alias for IRC property scoping + * @return mutable jobConf (per-table keys applied later by {@code jobConfForTable}) + */ + public Map<String, String> build(PolicyEntity policy, String catalogName) { + PolicyContent content = policy.content(); + PolicyExecutionSpec spec = executionSpec(content); + // ircAuth suffixes stay short-named; values are conf keys until resolveForSubmit. + Map<String, String> jobOptions = + MaintenanceAuthConfigs.overlayJobOptions( + spec.jobOptions(), MaintenanceAuthConfigs.loadIrcAuthAsConfigKeys(gravitinoConfig)); + Map<String, String> updaterOptions = + ensureUpdaterCallbackOptions( + policy, + MaintenanceAuthConfigs.overlayUpdaterOptions( + spec.updaterOptions(), + MaintenanceAuthConfigs.loadGravitinoAuthAsConfigKeys(gravitinoConfig))); + + Map<String, String> jobConf = new LinkedHashMap<>(); + Map<String, String> sparkConf = new LinkedHashMap<>(); + + String existingSparkConfJson = jobOptions.get("spark_conf"); + if (existingSparkConfJson != null && !existingSparkConfJson.trim().isEmpty()) { + sparkConf.putAll( + IcebergSparkConfigUtils.parseFlatJsonMap(existingSparkConfJson, "spark_conf")); + } + + for (Map.Entry<String, String> entry : jobOptions.entrySet()) { + String key = entry.getKey(); + String value = entry.getValue(); + if (key == null + || value == null + || "spark_conf".equals(key) + || "updater_options".equals(key)) { + continue; + } + routeJobOption(jobConf, sparkConf, catalogName, key, value); + } + + fillPlaceholdersFromSparkConf(jobConf, sparkConf, catalogName); + putTypedProcedureParams(jobConf, content); + + Optional<String> maintenanceType = MaintenanceJobTemplates.maintenanceType(policy.policyType()); + maintenanceType.ifPresent( + type -> { + jobConf.put(MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY, type); + jobConf.put(MaintenanceJobTemplates.POLICY_IDS_KEY, Long.toString(policy.id())); + }); + if (content != null && maintenanceType.isPresent()) { + Map<String, Map<String, Object>> rulesByType = new LinkedHashMap<>(); + rulesByType.put(maintenanceType.get(), new LinkedHashMap<>(content.rules())); + jobConf.put( + MaintenanceJobTemplates.RULES_KEY, + writeJson(rulesByType, MaintenanceJobTemplates.RULES_KEY)); + } + + jobConf.put( + MaintenanceJobTemplates.UPDATER_OPTIONS_JOB_CONF_KEY, + writeJson(updaterOptions, "updater_options")); + if (!sparkConf.isEmpty()) { + jobConf.put("spark_conf", writeJson(new TreeMap<>(sparkConf), "spark_conf")); + } + return jobConf; + } + + /** + * Expands {@code gravitino.maintenance.*} conf-key references in an enqueue snapshot into live + * values from {@code gravitino.conf}. Call immediately before {@code runJob}. + * + * @param enqueueJobConf jobConf from {@link #build} / {@code task_data} + * @return new map with secrets and URIs resolved + */ + public Map<String, String> resolveForSubmit(Map<String, String> enqueueJobConf) { + Map<String, String> resolved = new LinkedHashMap<>(); + for (Map.Entry<String, String> entry : enqueueJobConf.entrySet()) { + String key = entry.getKey(); + String value = entry.getValue(); + if (key == null || value == null) { + continue; + } + if (MaintenanceJobTemplates.UPDATER_OPTIONS_JOB_CONF_KEY.equals(key) + || "spark_conf".equals(key)) { + resolved.put(key, resolveJsonMapValues(value, key)); + } else { + resolved.put(key, MaintenanceAuthConfigs.resolveConfigReference(gravitinoConfig, value)); + } + } + return resolved; + } + + private String resolveJsonMapValues(String json, String field) { + Map<String, String> parsed = IcebergSparkConfigUtils.parseFlatJsonMap(json, field); + Map<String, String> out = new TreeMap<>(); + for (Map.Entry<String, String> entry : parsed.entrySet()) { + out.put( + entry.getKey(), + MaintenanceAuthConfigs.resolveConfigReference(gravitinoConfig, entry.getValue())); + } + return writeJson(out, field); + } + + private static void routeJobOption( + Map<String, String> jobConf, + Map<String, String> sparkConf, + String catalogName, + String key, + String value) { + if (FLAT_PLACEHOLDER_KEYS.contains(key)) { + jobConf.put(key, value); + return; + } + String placeholder = SPARK_KEY_TO_PLACEHOLDER.get(key); + if (placeholder != null) { + jobConf.put(placeholder, value); + sparkConf.put(key, value); + return; + } + if (key.startsWith("spark.")) { + // Spark session / SQL settings stay as full keys in spark_conf (never catalog-prefixed). + sparkConf.put(key, value); + return; + } + if (key.startsWith(SPARK_SQL_CATALOG_PREFIX)) { + sparkConf.put(key, value); + return; + } + // IRC / Iceberg catalog props from jobOptions / ircAuth (short names). + sparkConf.put(SPARK_SQL_CATALOG_PREFIX + catalogName + "." + key, value); + } + + private static void fillPlaceholdersFromSparkConf( + Map<String, String> jobConf, Map<String, String> sparkConf, String catalogName) { + for (Map.Entry<String, String> entry : SPARK_KEY_TO_PLACEHOLDER.entrySet()) { + String sparkKey = entry.getKey(); + String placeholder = entry.getValue(); + if (!jobConf.containsKey(placeholder) && sparkConf.containsKey(sparkKey)) { + jobConf.put(placeholder, sparkConf.get(sparkKey)); + } + } + String catalogPrefix = SPARK_SQL_CATALOG_PREFIX + catalogName; + putIfAbsentFromSpark(jobConf, sparkConf, "catalog_type", catalogPrefix + ".type"); + putIfAbsentFromSpark(jobConf, sparkConf, "catalog_uri", catalogPrefix + ".uri"); + putIfAbsentFromSpark(jobConf, sparkConf, "warehouse_location", catalogPrefix + ".warehouse"); + } + + private static void putIfAbsentFromSpark( + Map<String, String> jobConf, + Map<String, String> sparkConf, + String placeholder, + String sparkKey) { + if (!jobConf.containsKey(placeholder) && sparkConf.containsKey(sparkKey)) { + jobConf.put(placeholder, sparkConf.get(sparkKey)); + } + } + + /** + * Ensures Spark → Gravitino sample/finish callbacks have {@code metalake} (from the policy + * namespace) and {@code gravitino_uri} (from {@code gravitino.maintenance.gravitinoAuth.uri} via + * auth overlay, overridable by policy updaterOptions). + */ + private Map<String, String> ensureUpdaterCallbackOptions( + PolicyEntity policy, Map<String, String> overlaid) { + Map<String, String> updaterOptions = new LinkedHashMap<>(); + if (overlaid != null) { + updaterOptions.putAll(overlaid); + } + String metalake = policy.namespace().level(0); + updaterOptions.put("metalake", metalake); + + String gravitinoUri = updaterOptions.get("gravitino_uri"); + if (gravitinoUri == null || gravitinoUri.trim().isEmpty()) { + throw new IllegalArgumentException( + "updater_options.gravitino_uri is required for TMS Spark submit; set policy " + + "updaterOptions.gravitino_uri or gravitino.maintenance.gravitinoAuth.uri"); + } + return updaterOptions; + } + + private void putTypedProcedureParams( + Map<String, String> jobConf, @Nullable PolicyContent content) { + if (content instanceof IcebergDataCompactionContent) { + IcebergDataCompactionContent compaction = (IcebergDataCompactionContent) content; + jobConf.put("strategy", compaction.rewriteStrategy()); + if (compaction.sortOrder() != null && !compaction.sortOrder().isEmpty()) { + jobConf.put("sort_order", compaction.sortOrder()); + } + if (!compaction.rewriteOptions().isEmpty()) { + jobConf.put("options", writeJson(new TreeMap<>(compaction.rewriteOptions()), "options")); + } + return; + } + if (content instanceof IcebergOrphanFileRemovalContent) { + IcebergOrphanFileRemovalContent orphan = (IcebergOrphanFileRemovalContent) content; + IcebergOrphanFileRemovalContent.validateOlderThanDays(orphan.olderThanDays()); + jobConf.put("older_than", orphanCutoff(orphan.olderThanDays())); + jobConf.put("dry_run", Boolean.toString(orphan.dryRun())); + if (orphan.location() != null) { + jobConf.put("location", orphan.location()); + } + } + } + + private String orphanCutoff(long days) { + try { + return ORPHAN_CUTOFF.format(clock.instant().minus(days, ChronoUnit.DAYS)); + } catch (DateTimeException | ArithmeticException e) { + throw new IllegalArgumentException( + "olderThanDays cannot be represented as a cleanup timestamp", e); + } + } + + private static PolicyExecutionSpec executionSpec(@Nullable PolicyContent content) { + if (content instanceof ExecutablePolicyContent) { + return ((ExecutablePolicyContent) content).executionSpec(); + } + return PolicyExecutionSpec.empty(); + } + + private static String writeJson(Map<String, ?> map, String field) { + try { + return JsonUtils.anyFieldMapper().writeValueAsString(map); + } catch (JsonProcessingException e) { + throw new IllegalStateException("Failed to serialize " + field + " for runJob", e); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceJobTemplates.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceJobTemplates.java new file mode 100644 index 0000000000..07d116cef6 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceJobTemplates.java @@ -0,0 +1,96 @@ +/* + * 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.scheduler.job.cron; + +import java.util.Optional; +import org.apache.gravitino.job.JobTemplateProvider; +import org.apache.gravitino.policy.Policy; + +/** + * Maps maintenance {@link Policy.BuiltInType} to the TMS cron Spark job template and procedure + * kind. + * + * <p>Cron always submits {@link #TMS_ICEBERG_MAINTENANCE_TEMPLATE}; standalone {@code + * builtin-iceberg-*} templates remain for CLI / optimizer. + */ +public final class MaintenanceJobTemplates { + + /** + * Built-in template name for the TMS Iceberg maintenance Spark job. Must match {@code + * IcebergMaintenanceJob.NAME} ({@code builtin-iceberg-maintenance-task}). + */ + public static final String TMS_ICEBERG_MAINTENANCE_TEMPLATE = + JobTemplateProvider.BUILTIN_NAME_PREFIX + "iceberg-maintenance-task"; + + /** + * JobConf key for procedure dispatch inside the TMS template. One wire name, or a comma-separated + * commit chain. Required by {@link #TMS_ICEBERG_MAINTENANCE_TEMPLATE}. + */ + public static final String MAINTENANCE_TYPES_KEY = "maintenance_types"; + + /** + * JobConf key for policy entity ids parallel to {@link #MAINTENANCE_TYPES_KEY} (comma-separated). + * Required so Spark callbacks can attribute soft slots by {@code policyId}. + */ + public static final String POLICY_IDS_KEY = "policy_ids"; + + /** JobConf key for overlaid updater options (JSON object string). */ + public static final String UPDATER_OPTIONS_JOB_CONF_KEY = "updater_options"; + + /** + * JobConf key for enqueue-time policy {@code rules()}, JSON map of wire maintenance type → rules + * object. Consumed in-Spark for trigger-expr recommend (no Gravitino round-trip). + */ + public static final String RULES_KEY = "rules"; + + private MaintenanceJobTemplates() {} + + /** + * @param builtInType maintenance policy type + * @return TMS cron job template name when known + */ + public static Optional<String> templateName(Policy.BuiltInType builtInType) { + if (maintenanceType(builtInType).isEmpty()) { + return Optional.empty(); + } + return Optional.of(TMS_ICEBERG_MAINTENANCE_TEMPLATE); + } + + /** + * @param builtInType maintenance policy type + * @return wire value for {@link #MAINTENANCE_TYPES_KEY} + */ + public static Optional<String> maintenanceType(Policy.BuiltInType builtInType) { + if (builtInType == null) { + return Optional.empty(); + } + switch (builtInType) { + case ICEBERG_COMPACTION: + return Optional.of("rewrite_data_files"); + case ICEBERG_REWRITE_MANIFESTS: + return Optional.of("rewrite_manifests"); + case ICEBERG_SNAPSHOT_EXPIRATION: + return Optional.of("expire_snapshots"); + case ICEBERG_ORPHAN_FILE_REMOVAL: + return Optional.of("remove_orphan_files"); + default: + return Optional.empty(); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceTableRefs.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceTableRefs.java new file mode 100644 index 0000000000..6b855c50b1 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/MaintenanceTableRefs.java @@ -0,0 +1,169 @@ +/* + * 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.scheduler.job.cron; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.gravitino.storage.relational.mapper.CatalogMetaMapper; +import org.apache.gravitino.storage.relational.mapper.SchemaMetaMapper; +import org.apache.gravitino.storage.relational.mapper.TableMetaMapper; +import org.apache.gravitino.storage.relational.po.CatalogPO; +import org.apache.gravitino.storage.relational.po.SchemaPO; +import org.apache.gravitino.storage.relational.po.TablePO; +import org.apache.gravitino.storage.relational.utils.SessionUtils; + +/** Resolves table entity ids to catalog + {@code schema.table} identifiers for Spark jobConf. */ +public final class MaintenanceTableRefs { + + private MaintenanceTableRefs() {} + + /** + * @param tableIds table entity ids (order preserved for known ids) + * @return resolved refs; unknown / incomplete rows are omitted + */ + public static List<ResolvedTableRef> resolve(List<Long> tableIds) { + if (tableIds == null || tableIds.isEmpty()) { + return List.of(); + } + List<TablePO> tables = + SessionUtils.getWithoutCommit( + TableMetaMapper.class, mapper -> mapper.listTablePOsByTableIds(tableIds)); + if (tables == null || tables.isEmpty()) { + return List.of(); + } + Map<Long, TablePO> byId = new LinkedHashMap<>(); + for (TablePO table : tables) { + if (table != null && table.getTableId() != null) { + byId.put(table.getTableId(), table); + } + } + + Map<Long, String> catalogNames = new LinkedHashMap<>(); + Map<Long, String> schemaNames = new LinkedHashMap<>(); + List<ResolvedTableRef> resolved = new ArrayList<>(); + for (Long tableId : tableIds) { + TablePO table = byId.get(tableId); + if (table == null || table.getTableName() == null) { + continue; + } + String catalogName = catalogName(table.getCatalogId(), catalogNames); + String schemaName = schemaName(table.getSchemaId(), schemaNames); + if (catalogName == null || schemaName == null) { + continue; + } + resolved.add( + new ResolvedTableRef(tableId, catalogName, schemaName + "." + table.getTableName())); + } + return resolved; + } + + @Nullable + private static String catalogName(Long catalogId, Map<Long, String> cache) { + if (catalogId == null) { + return null; + } + if (cache.containsKey(catalogId)) { + return cache.get(catalogId); + } + CatalogPO catalog = + SessionUtils.getWithoutCommit( + CatalogMetaMapper.class, mapper -> mapper.selectCatalogMetaById(catalogId)); + String name = catalog == null ? null : catalog.getCatalogName(); + cache.put(catalogId, name); + return name; + } + + @Nullable + private static String schemaName(Long schemaId, Map<Long, String> cache) { + if (schemaId == null) { + return null; + } + if (cache.containsKey(schemaId)) { + return cache.get(schemaId); + } + SchemaPO schema = + SessionUtils.getWithoutCommit( + SchemaMetaMapper.class, mapper -> mapper.selectSchemaMetaById(schemaId)); + String name = schema == null ? null : schema.getSchemaName(); + cache.put(schemaId, name); + return name; + } + + /** One table ready for Spark expire-snapshots jobConf. */ + public static final class ResolvedTableRef { + private final long tableId; + private final String catalogName; + private final String tableIdentifier; + + /** + * @param tableId table entity id + * @param catalogName catalog name + * @param tableIdentifier {@code schema.table} Spark identifier + */ + public ResolvedTableRef(long tableId, String catalogName, String tableIdentifier) { + this.tableId = tableId; + this.catalogName = catalogName; + this.tableIdentifier = tableIdentifier; + } + + /** + * @return table entity id + */ + public long tableId() { + return tableId; + } + + /** + * @return catalog name + */ + public String catalogName() { + return catalogName; + } + + /** + * @return {@code schema.table} (or nested-namespace) Spark identifier + */ + public String tableIdentifier() { + return tableIdentifier; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ResolvedTableRef)) { + return false; + } + ResolvedTableRef that = (ResolvedTableRef) o; + return tableId == that.tableId + && Objects.equals(catalogName, that.catalogName) + && Objects.equals(tableIdentifier, that.tableIdentifier); + } + + @Override + public int hashCode() { + return Objects.hash(tableId, catalogName, tableIdentifier); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/policy/expand/PolicyExpandCrons.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/policy/expand/PolicyExpandCrons.java index 2efbde275e..4a6fedbfce 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/policy/expand/PolicyExpandCrons.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/policy/expand/PolicyExpandCrons.java @@ -22,45 +22,23 @@ import com.github.kagkarlsson.scheduler.task.ExecutionComplete; import com.github.kagkarlsson.scheduler.task.schedule.CronSchedule; import com.github.kagkarlsson.scheduler.task.schedule.CronStyle; import com.github.kagkarlsson.scheduler.task.schedule.Schedule; -import com.google.common.annotations.VisibleForTesting; import java.time.Instant; import java.time.ZoneId; -import java.time.ZoneOffset; -/** - * Quartz cron helpers for {@code policy-expand} next-due calculation. - * - * <p>Always evaluates cron expressions in UTC so every scheduler node computes the same next-due - * instant for a given {@code timeTrigger}, regardless of the JVM default time zone. - */ +/** Quartz cron helpers for {@code policy-expand} next-due calculation. */ final class PolicyExpandCrons { - /** Zone used for all policy-expand cron evaluations. */ - static final ZoneId SCHEDULER_ZONE = ZoneOffset.UTC; - private PolicyExpandCrons() {} /** - * Next fire time for a Quartz cron, strictly after now, in {@link #SCHEDULER_ZONE}. + * Next fire time for a Quartz cron, strictly after now. * * @param timeTrigger six-field Quartz cron expression - * @return next execution instant + * @return next execution instant in the system default zone */ static Instant nextExecution(String timeTrigger) { - return nextExecution(timeTrigger, Instant.now()); - } - - /** - * Next fire time for a Quartz cron, strictly after {@code after}, in {@link #SCHEDULER_ZONE}. - * - * @param timeTrigger six-field Quartz cron expression - * @param after reference instant - * @return next execution instant - */ - @VisibleForTesting - static Instant nextExecution(String timeTrigger, Instant after) { - CronSchedule schedule = new CronSchedule(timeTrigger, SCHEDULER_ZONE, CronStyle.QUARTZ); - Instant next = schedule.getNextExecutionTime(ExecutionComplete.simulatedSuccess(after)); + CronSchedule schedule = new CronSchedule(timeTrigger, ZoneId.systemDefault(), CronStyle.QUARTZ); + Instant next = schedule.getNextExecutionTime(ExecutionComplete.simulatedSuccess(Instant.now())); if (next == null || Schedule.NEVER.equals(next)) { throw new IllegalArgumentException("timeTrigger has no next execution: " + timeTrigger); } diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestMaintenanceJobSpecFactory.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestMaintenanceJobSpecFactory.java new file mode 100644 index 0000000000..74027d4b9e --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestMaintenanceJobSpecFactory.java @@ -0,0 +1,180 @@ +/* + * 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.scheduler; + +import com.github.kagkarlsson.scheduler.serializer.JacksonSerializer; +import com.github.kagkarlsson.scheduler.serializer.Serializer; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.policy.MaintenanceAuthConfigs; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.TableTarget; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceTableRefs.ResolvedTableRef; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContents; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TestMaintenanceJobSpecFactory { + + private static final String GRAVITINO_URI = "http://localhost:8090";; + + @Test + void testForCronBuildsSharedJobConfAndIdentifierLevels() { + MaintenanceJobSpecFactory factory = new MaintenanceJobSpecFactory(configWithUri()); + PolicyEntity policy = compaction(202L); + MaintenanceJobSpec spec = + factory.forCron( + policy, + List.of( + new ResolvedTableRef(10L, "iceberg", "db.t1"), + new ResolvedTableRef(11L, "iceberg", "db.t2"))); + + Assertions.assertTrue(spec.isValid()); + Assertions.assertEquals("ml", spec.metalake()); + Assertions.assertEquals(List.of(10L, 11L), spec.tableIds()); + Assertions.assertEquals(1, spec.getOperations().size()); + Assertions.assertEquals(202L, spec.getOperations().get(0).getPolicyId()); + Assertions.assertEquals("rewrite_data_files", spec.getOperations().get(0).getMaintenanceType()); + Assertions.assertNull(spec.getJobConf().get("catalog_name")); + Assertions.assertNull(spec.getJobConf().get("table_identifier")); + Assertions.assertFalse(spec.getJobConf().containsKey("tms_table_id")); + Assertions.assertFalse(spec.getJobConf().containsKey("tms_policy_id")); + Assertions.assertTrue( + spec.getJobConf() + .get("updater_options") + .contains("gravitino.maintenance.gravitinoAuth.uri")); + Assertions.assertFalse(spec.getJobConf().get("updater_options").contains(GRAVITINO_URI)); + + TableTarget first = spec.getTables().get(0); + Assertions.assertEquals(List.of("ml", "iceberg", "db", "t1"), first.getIdentifier()); + Map<String, String> perTable = + MaintenanceJobSpecFactory.jobConfForTable(spec.getJobConf(), first); + Assertions.assertEquals("iceberg", perTable.get("catalog_name")); + Assertions.assertNull(perTable.get("table_identifier")); + Assertions.assertNull(perTable.get("table_id")); + Assertions.assertTrue(perTable.get("tables").contains("\"table_id\":10")); + Assertions.assertTrue(perTable.get("tables").contains("\"table_identifier\":\"db.t1\"")); + Assertions.assertFalse(perTable.containsKey("tms_table_id")); + } + + @Test + void testForCronRejectsMixedCatalogs() { + MaintenanceJobSpecFactory factory = new MaintenanceJobSpecFactory(configWithUri()); + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + factory.forCron( + compaction(1L), + List.of( + new ResolvedTableRef(10L, "iceberg", "db.t1"), + new ResolvedTableRef(11L, "other", "db.t2")))); + } + + @Test + void testJacksonRoundTrip() { + MaintenanceJobSpecFactory factory = new MaintenanceJobSpecFactory(configWithUri()); + MaintenanceJobSpec original = + factory.forCron( + compaction(1L), List.of(new ResolvedTableRef(9L, "iceberg", "schema.table"))); + Serializer serializer = new JacksonSerializer(JsonUtils.anyFieldMapper()); + byte[] bytes = serializer.serialize(original); + MaintenanceJobSpec restored = serializer.deserialize(MaintenanceJobSpec.class, bytes); + Assertions.assertEquals(original, restored); + String json = new String(bytes, StandardCharsets.UTF_8); + Assertions.assertFalse(json.contains("\"version\"")); + Assertions.assertFalse(json.contains("\"metalake\"")); + Assertions.assertTrue(json.contains("\"operations\"")); + Assertions.assertTrue(json.contains("updater_options")); + } + + @Test + void testMergeResourceKeysTakesMax() { + Map<String, String> into = + MaintenanceJobSpecFactory.newLinkedMap( + "spark_executor_memory", "4g", "spark_executor_cores", "2", "catalog_uri", "http://a";); + Map<String, String> from = + MaintenanceJobSpecFactory.newLinkedMap( + "spark_executor_memory", "7g", "spark_executor_cores", "1", "catalog_uri", "http://b";); + MaintenanceJobSpecFactory.mergeResourceKeys(into, from); + Assertions.assertEquals("7g", into.get("spark_executor_memory")); + Assertions.assertEquals("2", into.get("spark_executor_cores")); + Assertions.assertEquals("http://a";, into.get("catalog_uri")); + } + + @Test + void testMergeResourceKeysComparesMemoryUnits() { + Map<String, String> into = + MaintenanceJobSpecFactory.newLinkedMap("spark_executor_memory", "1g"); + Map<String, String> from = + MaintenanceJobSpecFactory.newLinkedMap("spark_executor_memory", "512m"); + MaintenanceJobSpecFactory.mergeResourceKeys(into, from); + Assertions.assertEquals("1g", into.get("spark_executor_memory")); + + Map<String, String> mixed = + MaintenanceJobSpecFactory.newLinkedMap("spark_executor_memory", "512m"); + MaintenanceJobSpecFactory.mergeResourceKeys( + mixed, MaintenanceJobSpecFactory.newLinkedMap("spark_executor_memory", "1500m")); + Assertions.assertEquals("1500m", mixed.get("spark_executor_memory")); + } + + private static Config configWithUri() { + Config config = new Config(false) {}; + config.loadFromMap( + Map.of( + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey(), + GRAVITINO_URI), + t -> true); + return config; + } + + private static PolicyEntity compaction(long id) { + return PolicyEntity.builder() + .withId(id) + .withName("p-" + id) + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + PolicyContents.icebergDataCompaction( + IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE, + IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER, + IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT, + IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT, + IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM, + "binpack", + "", + Map.of(), + null, + true, + null, + Map.of("spark_executor_memory", "4g"), + Map.of())) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestMaintenanceJobConfBuilder.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestMaintenanceJobConfBuilder.java new file mode 100644 index 0000000000..b60bc5bd04 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestMaintenanceJobConfBuilder.java @@ -0,0 +1,260 @@ +/* + * 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.scheduler.job.cron; + +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.Map; +import org.apache.gravitino.Config; +import org.apache.gravitino.maintenance.policy.MaintenanceAuthConfigs; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContents; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TestMaintenanceJobConfBuilder { + + private static final long POLICY_ID = 9L; + private static final Instant FIXED = Instant.parse("2024-06-15T12:00:00Z"); + private static final String GRAVITINO_URI = "http://localhost:8090";; + + @Test + void testCompactionMapsTypedFieldsAndIrcIntoSparkConf() { + Config config = configWithUri(); + config.loadFromMap( + Map.of( + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey(), + GRAVITINO_URI, + "gravitino.maintenance.ircAuth.rest.auth.type", + "basic"), + t -> true); + + PolicyEntity policy = + PolicyEntity.builder() + .withId(POLICY_ID) + .withName("c") + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + PolicyContents.icebergDataCompaction( + IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE, + IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER, + IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT, + IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT, + IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM, + IcebergDataCompactionContent.REWRITE_STRATEGY_SORT, + "id ASC NULLS LAST", + Map.of("min-input-files", "4"), + null, + null, + null, + Map.of("catalog_type", "rest", "catalog_uri", "http://localhost:9001/iceberg";), + Map.of())) + .withAuditInfo( + AuditInfo.builder().withCreator("t").withCreateTime(Instant.now()).build()) + .build(); + + Map<String, String> conf = + new MaintenanceJobConfBuilder(config, Clock.fixed(FIXED, ZoneOffset.UTC)) + .build(policy, "iceberg"); + + Assertions.assertEquals("sort", conf.get("strategy")); + Assertions.assertEquals("id ASC NULLS LAST", conf.get("sort_order")); + Assertions.assertTrue(conf.get("options").contains("min-input-files")); + Assertions.assertNull(conf.get("where_clause")); + Assertions.assertEquals("rest", conf.get("catalog_type")); + Assertions.assertEquals("rewrite_data_files", conf.get("maintenance_types")); + Assertions.assertEquals(Long.toString(policy.id()), conf.get("policy_ids")); + Assertions.assertTrue( + conf.get(MaintenanceJobTemplates.RULES_KEY).contains("\"rewrite_data_files\"")); + Assertions.assertTrue(conf.get(MaintenanceJobTemplates.RULES_KEY).contains("trigger-expr")); + Assertions.assertTrue( + conf.get("spark_conf") + .contains( + "\"spark.sql.catalog.iceberg.rest.auth.type\":\"gravitino.maintenance.ircAuth.rest.auth.type\"")); + Assertions.assertFalse(conf.get("spark_conf").contains("\"basic\"")); + Assertions.assertNull(conf.get("rest.auth.type")); + Assertions.assertFalse(conf.get("spark_conf").contains("\"rest.auth.type\"")); + Assertions.assertTrue(conf.get("updater_options").contains("\"metalake\":\"ml\"")); + Assertions.assertTrue( + conf.get("updater_options") + .contains("\"gravitino_uri\":\"gravitino.maintenance.gravitinoAuth.uri\"")); + Assertions.assertFalse(conf.get("updater_options").contains(GRAVITINO_URI)); + + Map<String, String> resolved = + new MaintenanceJobConfBuilder(config, Clock.fixed(FIXED, ZoneOffset.UTC)) + .resolveForSubmit(conf); + Assertions.assertTrue( + resolved.get("updater_options").contains("\"gravitino_uri\":\"" + GRAVITINO_URI + "\"")); + Assertions.assertTrue( + resolved + .get("spark_conf") + .contains("\"spark.sql.catalog.iceberg.rest.auth.type\":\"basic\"")); + } + + @Test + void testSparkResourceKeysFlattenToPlaceholdersNotCatalogProps() { + PolicyEntity policy = + PolicyEntity.builder() + .withId(POLICY_ID) + .withName("c") + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + PolicyContents.icebergDataCompaction( + IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE, + IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER, + IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT, + IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT, + IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + null, + null, + null, + Map.of( + "spark.executor.memory", + "7g", + "spark_executor_cores", + "4", + "catalog_type", + "rest"), + Map.of())) + .withAuditInfo( + AuditInfo.builder().withCreator("t").withCreateTime(Instant.now()).build()) + .build(); + + Map<String, String> conf = + new MaintenanceJobConfBuilder(configWithUri(), Clock.fixed(FIXED, ZoneOffset.UTC)) + .build(policy, "iceberg"); + + Assertions.assertEquals("7g", conf.get("spark_executor_memory")); + Assertions.assertEquals("4", conf.get("spark_executor_cores")); + Assertions.assertEquals("rest", conf.get("catalog_type")); + Assertions.assertTrue(conf.get("spark_conf").contains("\"spark.executor.memory\":\"7g\"")); + Assertions.assertFalse( + conf.get("spark_conf").contains("spark.sql.catalog.iceberg.spark.executor.memory")); + } + + @Test + void testOrphanMapsRetentionCutoff() { + PolicyEntity policy = + PolicyEntity.builder() + .withId(POLICY_ID) + .withName("o") + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_ORPHAN_FILE_REMOVAL) + .withEnabled(true) + .withContent(PolicyContents.icebergOrphanFileRemoval(7L, null, false)) + .withAuditInfo( + AuditInfo.builder().withCreator("t").withCreateTime(Instant.now()).build()) + .build(); + + Map<String, String> conf = + new MaintenanceJobConfBuilder(configWithUri(), Clock.fixed(FIXED, ZoneOffset.UTC)) + .build(policy, "iceberg"); + + Assertions.assertEquals("remove_orphan_files", conf.get("maintenance_types")); + Assertions.assertEquals(Long.toString(policy.id()), conf.get("policy_ids")); + Assertions.assertEquals("false", conf.get("dry_run")); + Assertions.assertEquals("2024-06-08 12:00:00Z", conf.get("older_than")); + } + + @Test + void testMissingGravitinoUriFailsClosed() { + PolicyEntity policy = + PolicyEntity.builder() + .withId(POLICY_ID) + .withName("c") + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent(PolicyContents.icebergDataCompaction()) + .withAuditInfo( + AuditInfo.builder().withCreator("t").withCreateTime(Instant.now()).build()) + .build(); + + IllegalArgumentException thrown = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + new MaintenanceJobConfBuilder( + new Config(false) {}, Clock.fixed(FIXED, ZoneOffset.UTC)) + .build(policy, "iceberg")); + Assertions.assertTrue(thrown.getMessage().contains("gravitino_uri")); + } + + @Test + void testPolicyUpdaterOptionsOverrideClusterUriButNotMetalake() { + Config config = configWithUri(); + PolicyEntity policy = + PolicyEntity.builder() + .withId(POLICY_ID) + .withName("c") + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + PolicyContents.icebergDataCompaction( + IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE, + IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER, + IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT, + IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT, + IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + null, + null, + null, + Map.of(), + Map.of("gravitino_uri", "http://policy-uri:8090";, "metalake", "policy-ml"))) + .withAuditInfo( + AuditInfo.builder().withCreator("t").withCreateTime(Instant.now()).build()) + .build(); + + Map<String, String> conf = + new MaintenanceJobConfBuilder(config, Clock.fixed(FIXED, ZoneOffset.UTC)) + .build(policy, "iceberg"); + + Assertions.assertTrue( + conf.get("updater_options").contains("\"gravitino_uri\":\"http://policy-uri:8090\"";)); + // Path metalake always comes from the policy namespace, never policy updaterOptions. + Assertions.assertTrue(conf.get("updater_options").contains("\"metalake\":\"ml\"")); + Assertions.assertFalse(conf.get("updater_options").contains("policy-ml")); + } + + private static Config configWithUri() { + Config config = new Config(false) {}; + config.loadFromMap( + Map.of( + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey(), + GRAVITINO_URI), + t -> true); + return config; + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestMaintenanceJobTemplates.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestMaintenanceJobTemplates.java new file mode 100644 index 0000000000..123124f004 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestMaintenanceJobTemplates.java @@ -0,0 +1,58 @@ +/* + * 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.scheduler.job.cron; + +import org.apache.gravitino.policy.Policy; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link MaintenanceJobTemplates}. */ +class TestMaintenanceJobTemplates { + + @Test + void testAllFourTypesMapToSingleTmsTemplate() { + Assertions.assertEquals( + "builtin-iceberg-maintenance-task", + MaintenanceJobTemplates.TMS_ICEBERG_MAINTENANCE_TEMPLATE); + for (Policy.BuiltInType type : + new Policy.BuiltInType[] { + Policy.BuiltInType.ICEBERG_COMPACTION, + Policy.BuiltInType.ICEBERG_REWRITE_MANIFESTS, + Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION, + Policy.BuiltInType.ICEBERG_ORPHAN_FILE_REMOVAL + }) { + Assertions.assertEquals( + MaintenanceJobTemplates.TMS_ICEBERG_MAINTENANCE_TEMPLATE, + MaintenanceJobTemplates.templateName(type).orElseThrow()); + Assertions.assertTrue(MaintenanceJobTemplates.maintenanceType(type).isPresent()); + } + } + + @Test + void testMaintenanceTypeWireValues() { + Assertions.assertEquals( + "rewrite_data_files", + MaintenanceJobTemplates.maintenanceType(Policy.BuiltInType.ICEBERG_COMPACTION) + .orElseThrow()); + Assertions.assertEquals( + "expire_snapshots", + MaintenanceJobTemplates.maintenanceType(Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION) + .orElseThrow()); + } +}
