thomasrebele commented on code in PR #4602: URL: https://github.com/apache/calcite/pull/4602#discussion_r2487086372
########## core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java: ########## @@ -0,0 +1,155 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.visualizer.RuleMatchVisualizer; +import org.apache.calcite.runtime.Hook; + +import java.io.File; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +/** + * Utility class to enable RuleMatchVisualizer for Calcite connections. + * + * <p>This class provides hooks to automatically attach a RuleMatchVisualizer + * to planners when a connection specifies the ruleVisualizerDir property. + * + * <p>Usage in JDBC URL: + * <blockquote><pre> + * jdbc:calcite:ruleVisualizerDir=/tmp/calcite-viz + * </pre></blockquote> + * + * <p>Or programmatically: + * <blockquote><pre> + * RuleMatchVisualizerHook.enable("/tmp/calcite-viz"); + * </pre></blockquote> + */ +public class RuleMatchVisualizerHook { + public static final RuleMatchVisualizerHook INSTANCE = new RuleMatchVisualizerHook(); + + private final Map<RelOptPlanner, RuleMatchVisualizer> visualizerMap = new HashMap<>(); + private final AtomicInteger queryCounter = new AtomicInteger(0); + + private Hook.Closeable hookCloseable = Hook.Closeable.EMPTY; + + /** Private constructor to prevent instantiation. */ + private RuleMatchVisualizerHook() {} + + /** + * Enables the visualizer for all subsequent queries with the specified output directory. + * + * @param outputDir Directory where visualization files will be created + */ + public synchronized void enable(String outputDir) { + hookCloseable.close(); + + // Ensure the output directory exists + File dir = new File(outputDir); + if (!dir.exists()) { + boolean madeDir = dir.mkdirs(); + assert madeDir : "Failed to create directory: " + outputDir; + } + + // Install the hook + hookCloseable = Hook.PLANNER.addThread((Consumer<RelOptPlanner>) planner -> { + attachVisualizer(planner, outputDir); + }); + } + + /** + * Enables the visualizer using the connection's configuration. + * This method checks if the connection has the ruleVisualizerDir property set. + * + * @param connection The Calcite connection + */ + public synchronized void enableFromConnection(CalciteConnection connection) { + CalciteConnectionConfig config = connection.config(); + String vizDir = config.ruleVisualizerDir(); + + if (vizDir != null && !vizDir.isEmpty()) { + enable(vizDir); + } + } + + /** + * Disables the visualizer. + */ + public synchronized void disable() { + hookCloseable.close(); + + // Write any pending visualizations + for (RuleMatchVisualizer viz : visualizerMap.values()) { + viz.writeToFile(); + } + visualizerMap.clear(); + } + + /** + * Attaches a visualizer to the given planner. + */ + private void attachVisualizer(RelOptPlanner planner, String outputDir) { + + // Check if we've already attached a visualizer to this planner + if (visualizerMap.containsKey(planner)) { + return; + } + + int queryNum = queryCounter.incrementAndGet(); + int queryStart = (int) System.currentTimeMillis() / 1000; + String suffix = String.format(Locale.ROOT, "query_%d_%d", queryNum, queryStart); + + // Create and attach the visualizer + RuleMatchVisualizer visualizer = new RuleMatchVisualizer(outputDir, suffix); + visualizer.attachTo(planner); + visualizerMap.put(planner, visualizer); + + // For HepPlanner, we need to manually write the output + if (planner instanceof HepPlanner) { + // Add a hook to write the visualization after the planner finishes + Hook.PLAN_BEFORE_IMPLEMENTATION.addThread(relRoot -> { + RuleMatchVisualizer viz = visualizerMap.get(planner); + if (viz != null) { + viz.writeToFile(); + visualizerMap.remove(planner); + } + }); + } + // VolcanoPlanner automatically calls writeToFile() when done + + System.out.println("RuleMatchVisualizer enabled: Output will be written to " + + outputDir + File.separator + suffix + "*"); Review Comment: Please use a LOGGER here as well. ########## core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java: ########## @@ -0,0 +1,155 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.visualizer.RuleMatchVisualizer; +import org.apache.calcite.runtime.Hook; + +import java.io.File; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +/** + * Utility class to enable RuleMatchVisualizer for Calcite connections. + * + * <p>This class provides hooks to automatically attach a RuleMatchVisualizer + * to planners when a connection specifies the ruleVisualizerDir property. + * + * <p>Usage in JDBC URL: + * <blockquote><pre> + * jdbc:calcite:ruleVisualizerDir=/tmp/calcite-viz + * </pre></blockquote> + * + * <p>Or programmatically: + * <blockquote><pre> + * RuleMatchVisualizerHook.enable("/tmp/calcite-viz"); + * </pre></blockquote> + */ +public class RuleMatchVisualizerHook { + public static final RuleMatchVisualizerHook INSTANCE = new RuleMatchVisualizerHook(); + + private final Map<RelOptPlanner, RuleMatchVisualizer> visualizerMap = new HashMap<>(); + private final AtomicInteger queryCounter = new AtomicInteger(0); + + private Hook.Closeable hookCloseable = Hook.Closeable.EMPTY; + + /** Private constructor to prevent instantiation. */ + private RuleMatchVisualizerHook() {} + + /** + * Enables the visualizer for all subsequent queries with the specified output directory. + * + * @param outputDir Directory where visualization files will be created + */ + public synchronized void enable(String outputDir) { + hookCloseable.close(); + + // Ensure the output directory exists + File dir = new File(outputDir); + if (!dir.exists()) { + boolean madeDir = dir.mkdirs(); + assert madeDir : "Failed to create directory: " + outputDir; + } + + // Install the hook + hookCloseable = Hook.PLANNER.addThread((Consumer<RelOptPlanner>) planner -> { + attachVisualizer(planner, outputDir); + }); + } + + /** + * Enables the visualizer using the connection's configuration. + * This method checks if the connection has the ruleVisualizerDir property set. + * + * @param connection The Calcite connection + */ + public synchronized void enableFromConnection(CalciteConnection connection) { + CalciteConnectionConfig config = connection.config(); + String vizDir = config.ruleVisualizerDir(); + + if (vizDir != null && !vizDir.isEmpty()) { + enable(vizDir); + } + } + + /** + * Disables the visualizer. + */ + public synchronized void disable() { Review Comment: The `disable()` should be called somewhere as well. The Handler defined in `org.apache.calcite.jdbc.Driver#createHandler` could override `onConnectionClose` to call it. I would strongly recommend to make the calls to "enable" and "disable" in the same class. So either moving the init code to `Driver#createHandler` as well, or creating another method in `CalciteConnectionImpl` to complement the `init()` method. ########## core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java: ########## @@ -0,0 +1,155 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.visualizer.RuleMatchVisualizer; +import org.apache.calcite.runtime.Hook; + +import java.io.File; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +/** + * Utility class to enable RuleMatchVisualizer for Calcite connections. + * + * <p>This class provides hooks to automatically attach a RuleMatchVisualizer + * to planners when a connection specifies the ruleVisualizerDir property. + * + * <p>Usage in JDBC URL: + * <blockquote><pre> + * jdbc:calcite:ruleVisualizerDir=/tmp/calcite-viz + * </pre></blockquote> + * + * <p>Or programmatically: + * <blockquote><pre> + * RuleMatchVisualizerHook.enable("/tmp/calcite-viz"); + * </pre></blockquote> + */ +public class RuleMatchVisualizerHook { + public static final RuleMatchVisualizerHook INSTANCE = new RuleMatchVisualizerHook(); + + private final Map<RelOptPlanner, RuleMatchVisualizer> visualizerMap = new HashMap<>(); + private final AtomicInteger queryCounter = new AtomicInteger(0); + + private Hook.Closeable hookCloseable = Hook.Closeable.EMPTY; + + /** Private constructor to prevent instantiation. */ + private RuleMatchVisualizerHook() {} + + /** + * Enables the visualizer for all subsequent queries with the specified output directory. + * + * @param outputDir Directory where visualization files will be created + */ + public synchronized void enable(String outputDir) { + hookCloseable.close(); + + // Ensure the output directory exists + File dir = new File(outputDir); + if (!dir.exists()) { + boolean madeDir = dir.mkdirs(); + assert madeDir : "Failed to create directory: " + outputDir; + } + + // Install the hook + hookCloseable = Hook.PLANNER.addThread((Consumer<RelOptPlanner>) planner -> { + attachVisualizer(planner, outputDir); + }); + } + + /** + * Enables the visualizer using the connection's configuration. + * This method checks if the connection has the ruleVisualizerDir property set. + * + * @param connection The Calcite connection + */ + public synchronized void enableFromConnection(CalciteConnection connection) { + CalciteConnectionConfig config = connection.config(); + String vizDir = config.ruleVisualizerDir(); + + if (vizDir != null && !vizDir.isEmpty()) { + enable(vizDir); + } + } + + /** + * Disables the visualizer. + */ + public synchronized void disable() { + hookCloseable.close(); + + // Write any pending visualizations + for (RuleMatchVisualizer viz : visualizerMap.values()) { + viz.writeToFile(); + } + visualizerMap.clear(); + } + + /** + * Attaches a visualizer to the given planner. + */ + private void attachVisualizer(RelOptPlanner planner, String outputDir) { + + // Check if we've already attached a visualizer to this planner + if (visualizerMap.containsKey(planner)) { + return; + } + + int queryNum = queryCounter.incrementAndGet(); + int queryStart = (int) System.currentTimeMillis() / 1000; + String suffix = String.format(Locale.ROOT, "query_%d_%d", queryNum, queryStart); + + // Create and attach the visualizer + RuleMatchVisualizer visualizer = new RuleMatchVisualizer(outputDir, suffix); + visualizer.attachTo(planner); + visualizerMap.put(planner, visualizer); + + // For HepPlanner, we need to manually write the output Review Comment: A VolcanoPlanner will automatically write the output when its done with the planning. However, the `writeToFile()` needs to be called explicitly for a HepPlanner. So maybe `explicitly` would be a better word. ########## core/src/test/java/org/apache/calcite/util/RuleMatchVisualizerHookTest.java: ########## @@ -0,0 +1,200 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.jdbc.CalciteConnection; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.apache.calcite.plan.visualizer.RuleMatchVisualizer.DATA_FILE_PREFIX; +import static org.apache.calcite.plan.visualizer.RuleMatchVisualizer.HTML_FILE_PREFIX; + +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasSize; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import static java.util.Objects.requireNonNull; + +/** + * Tests for {@link RuleMatchVisualizerHook}. + */ +class RuleMatchVisualizerHookTest { + + @TempDir + Path tempDir; + + @AfterEach + void cleanup() { + RuleMatchVisualizerHook.INSTANCE.disable(); + } + + @Test void testEnableDisable() { + String vizDir = tempDir.toString(); + + // Enable visualizer + RuleMatchVisualizerHook.INSTANCE.enable(vizDir); + + // Directory should be created + File dir = new File(vizDir); + assertTrue(dir.exists()); + assertTrue(dir.isDirectory()); + + // Disable visualizer + RuleMatchVisualizerHook.INSTANCE.disable(); + } + + @Test void testEnableFromConnection() throws Exception { + String vizDir = tempDir.resolve("viz").toString(); + + // Create connection with visualizer property + String url = "jdbc:calcite:ruleVisualizerDir=" + vizDir; + try (Connection conn = DriverManager.getConnection(url)) { + assertThat(conn, instanceOf(CalciteConnection.class)); + CalciteConnection calciteConn = (CalciteConnection) conn; + + // Check that the property is set + String configuredDir = calciteConn.config().ruleVisualizerDir(); + assertThat(configuredDir, notNullValue()); + assertThat(configuredDir, is(vizDir)); + + // The hook should be enabled automatically by the connection + // Let's run a simple query to trigger the planner + try (Statement stmt = conn.createStatement()) { + String sql = "SELECT 1 FROM (VALUES (1))"; + try (ResultSet rs = stmt.executeQuery(sql)) { + assertTrue(rs.next()); + assertThat(rs.getInt(1), is(1)); + } + } + } + + File dir = new File(vizDir); + assertTrue(dir.exists()); + assertTrue(dir.isDirectory()); + Map<Boolean, List<File>> matched = + Arrays.stream(requireNonNull(dir.listFiles())) + .collect(Collectors.partitioningBy(f -> f.getName().contains(DATA_FILE_PREFIX))); + + List<File> dataFiles = matched.get(true); + List<File> htmlFiles = matched.get(false); + + assertThat(dataFiles, hasSize(1)); + assertThat(htmlFiles, hasSize(1)); + assertThat(dataFiles.get(0).getName(), containsString(DATA_FILE_PREFIX)); + assertThat(htmlFiles.get(0).getName(), containsString(HTML_FILE_PREFIX)); Review Comment: JUnit takes care of this, see [TempDir](https://docs.junit.org/5.9.1/api/org.junit.jupiter.api/org/junit/jupiter/api/io/TempDir.html). ########## core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java: ########## @@ -0,0 +1,155 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.visualizer.RuleMatchVisualizer; +import org.apache.calcite.runtime.Hook; + +import java.io.File; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +/** + * Utility class to enable RuleMatchVisualizer for Calcite connections. + * + * <p>This class provides hooks to automatically attach a RuleMatchVisualizer + * to planners when a connection specifies the ruleVisualizerDir property. + * + * <p>Usage in JDBC URL: + * <blockquote><pre> + * jdbc:calcite:ruleVisualizerDir=/tmp/calcite-viz + * </pre></blockquote> + * + * <p>Or programmatically: + * <blockquote><pre> + * RuleMatchVisualizerHook.enable("/tmp/calcite-viz"); + * </pre></blockquote> + */ +public class RuleMatchVisualizerHook { + public static final RuleMatchVisualizerHook INSTANCE = new RuleMatchVisualizerHook(); + + private final Map<RelOptPlanner, RuleMatchVisualizer> visualizerMap = new HashMap<>(); + private final AtomicInteger queryCounter = new AtomicInteger(0); + + private Hook.Closeable hookCloseable = Hook.Closeable.EMPTY; + + /** Private constructor to prevent instantiation. */ + private RuleMatchVisualizerHook() {} + + /** + * Enables the visualizer for all subsequent queries with the specified output directory. + * + * @param outputDir Directory where visualization files will be created + */ + public synchronized void enable(String outputDir) { + hookCloseable.close(); + + // Ensure the output directory exists + File dir = new File(outputDir); + if (!dir.exists()) { + boolean madeDir = dir.mkdirs(); + assert madeDir : "Failed to create directory: " + outputDir; + } + + // Install the hook + hookCloseable = Hook.PLANNER.addThread((Consumer<RelOptPlanner>) planner -> { + attachVisualizer(planner, outputDir); + }); + } + + /** + * Enables the visualizer using the connection's configuration. + * This method checks if the connection has the ruleVisualizerDir property set. + * + * @param connection The Calcite connection + */ + public synchronized void enableFromConnection(CalciteConnection connection) { Review Comment: The code in [CalciteConnectionImpl](https://github.com/apache/calcite/pull/4602/files#diff-312ec137d81d2e343908c9c879bbf31a51f363df2e1e633f764aa6702715987eR198) already checks for an existing dir, so it could just call the `enable` method directly. ########## core/src/main/java/org/apache/calcite/jdbc/CalciteConnectionImpl.java: ########## @@ -188,6 +189,18 @@ void init() { true, true); } } + + // Enable RuleMatchVisualizer if configured + CalciteConnectionConfig cfg = config(); + String vizDir = cfg.ruleVisualizerDir(); + if (vizDir != null && !vizDir.isEmpty()) { + try { + RuleMatchVisualizerHook.INSTANCE.enableFromConnection(this); + } catch (Exception e) { + // Log but don't fail connection if visualizer setup fails + System.err.println("Warning: Failed to enable RuleMatchVisualizer: " + e.getMessage()); Review Comment: I would prefer using a LOGGER. Not sure whether a new method in `org.apache.calcite.util.trace.CalciteTrace` is needed, or whether a LOGGER can be defined as in `org.apache.calcite.runtime.ResultSetEnumerable#LOGGER`. ########## core/src/main/java/org/apache/calcite/util/RuleMatchVisualizerHook.java: ########## @@ -0,0 +1,155 @@ +/* + * 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.calcite.util; + +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.jdbc.CalciteConnection; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.visualizer.RuleMatchVisualizer; +import org.apache.calcite.runtime.Hook; + +import java.io.File; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +/** + * Utility class to enable RuleMatchVisualizer for Calcite connections. + * + * <p>This class provides hooks to automatically attach a RuleMatchVisualizer + * to planners when a connection specifies the ruleVisualizerDir property. + * + * <p>Usage in JDBC URL: + * <blockquote><pre> + * jdbc:calcite:ruleVisualizerDir=/tmp/calcite-viz + * </pre></blockquote> + * + * <p>Or programmatically: + * <blockquote><pre> + * RuleMatchVisualizerHook.enable("/tmp/calcite-viz"); + * </pre></blockquote> + */ +public class RuleMatchVisualizerHook { + public static final RuleMatchVisualizerHook INSTANCE = new RuleMatchVisualizerHook(); + + private final Map<RelOptPlanner, RuleMatchVisualizer> visualizerMap = new HashMap<>(); + private final AtomicInteger queryCounter = new AtomicInteger(0); + + private Hook.Closeable hookCloseable = Hook.Closeable.EMPTY; + + /** Private constructor to prevent instantiation. */ + private RuleMatchVisualizerHook() {} + + /** + * Enables the visualizer for all subsequent queries with the specified output directory. + * + * @param outputDir Directory where visualization files will be created + */ + public synchronized void enable(String outputDir) { + hookCloseable.close(); + + // Ensure the output directory exists + File dir = new File(outputDir); + if (!dir.exists()) { + boolean madeDir = dir.mkdirs(); + assert madeDir : "Failed to create directory: " + outputDir; + } + + // Install the hook + hookCloseable = Hook.PLANNER.addThread((Consumer<RelOptPlanner>) planner -> { + attachVisualizer(planner, outputDir); + }); + } + + /** + * Enables the visualizer using the connection's configuration. + * This method checks if the connection has the ruleVisualizerDir property set. + * + * @param connection The Calcite connection + */ + public synchronized void enableFromConnection(CalciteConnection connection) { + CalciteConnectionConfig config = connection.config(); + String vizDir = config.ruleVisualizerDir(); + + if (vizDir != null && !vizDir.isEmpty()) { + enable(vizDir); + } + } + + /** + * Disables the visualizer. + */ + public synchronized void disable() { + hookCloseable.close(); + + // Write any pending visualizations + for (RuleMatchVisualizer viz : visualizerMap.values()) { + viz.writeToFile(); + } + visualizerMap.clear(); + } + + /** + * Attaches a visualizer to the given planner. + */ + private void attachVisualizer(RelOptPlanner planner, String outputDir) { + + // Check if we've already attached a visualizer to this planner + if (visualizerMap.containsKey(planner)) { + return; + } + + int queryNum = queryCounter.incrementAndGet(); + int queryStart = (int) System.currentTimeMillis() / 1000; + String suffix = String.format(Locale.ROOT, "query_%d_%d", queryNum, queryStart); Review Comment: I would suggest using `plan_%d_%d`, as a query might be compiled in several stages and have multiple plans. As far as I know, Calcite does not provide a query id. Some projects that use Calcite may do, though, and I wonder whether there's an easy way to include this in the file name. The query id could be passed as a setter/getter on the `RelOptPlanner`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
