Enable query and task retries in MariaDbClient

Copy test overrides from 7b80852, as MariaDB also needs them.

docs/src/main/sphinx/admin/fault-tolerant-execution.md

@@ -44,6 +44,7 @@ connector. The following connectors support fault-tolerant execution:
 - {ref}`Delta Lake connector <delta-lake-fte-support>`
 - {ref}`Hive connector <hive-fte-support>`
 - {ref}`Iceberg connector <iceberg-fte-support>`
+- {ref}`MariaDB connector <mariadb-fte-support>`
 - {ref}`MongoDB connector <mongodb-fte-support>`
 - {ref}`MySQL connector <mysql-fte-support>`
 - {ref}`Oracle connector <oracle-fte-support>`

docs/src/main/sphinx/connector/mariadb.md

@@ -299,6 +299,12 @@ statements, the connector supports the following features:
 ```{include} sql-delete-limitation.fragment
 ```
 
+(mariadb-fte-support)=
+## Fault-tolerant execution support
+
+The connector supports {doc}`/admin/fault-tolerant-execution` of query
+processing. Read and write operations are both supported with any retry policy.
+
 ## Table functions
 
 The connector provides specific {doc}`table functions </functions/table>` to

plugin/trino-mariadb/pom.xml

@@ -130,6 +130,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-exchange-filesystem</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-main</artifactId>

plugin/trino-mariadb/src/main/java/io/trino/plugin/mariadb/MariaDbClient.java

@@ -187,7 +187,7 @@ public MariaDbClient(
             IdentifierMapping identifierMapping,
             RemoteQueryModifier queryModifier)
     {
-        super("`", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, false);
+        super("`", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, true);
 
         JdbcTypeHandle bigintTypeHandle = new JdbcTypeHandle(Types.BIGINT, Optional.of("bigint"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
         this.connectorExpressionRewriter = JdbcConnectorExpressionRewriterBuilder.newBuilder()

plugin/trino-mariadb/src/test/java/io/trino/plugin/mariadb/BaseMariaDbFailureRecoveryTest.java

@@ -0,0 +1,79 @@
+/*
+ * Licensed 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 io.trino.plugin.mariadb;
+
+import com.google.common.collect.ImmutableMap;
+import io.trino.operator.RetryPolicy;
+import io.trino.plugin.exchange.filesystem.FileSystemExchangePlugin;
+import io.trino.plugin.jdbc.BaseJdbcFailureRecoveryTest;
+import io.trino.testing.QueryRunner;
+import io.trino.tpch.TpchTable;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assumptions.abort;
+
+public abstract class BaseMariaDbFailureRecoveryTest
+        extends BaseJdbcFailureRecoveryTest
+{
+    public BaseMariaDbFailureRecoveryTest(RetryPolicy retryPolicy)
+    {
+        super(retryPolicy);
+    }
+
+    @Override
+    protected QueryRunner createQueryRunner(List<TpchTable<?>> requiredTpchTables, Map<String, String> configProperties, Map<String, String> coordinatorProperties)
+            throws Exception
+    {
+        TestingMariaDbServer server = closeAfterClass(new TestingMariaDbServer());
+        return MariaDbQueryRunner.builder(server)
+                .setInitialTables(requiredTpchTables)
+                .setExtraProperties(configProperties)
+                .setCoordinatorProperties(coordinatorProperties)
+                .setAdditionalSetup(runner -> {
+                    runner.installPlugin(new FileSystemExchangePlugin());
+                    runner.loadExchangeManager("filesystem", ImmutableMap.of(
+                            "exchange.base-directories", System.getProperty("java.io.tmpdir") + "/trino-local-file-system-exchange-manager"));
+                })
+                .build();
+    }
+
+    @Test
+    @Override
+    protected void testUpdateWithSubquery()
+    {
+        assertThatThrownBy(super::testUpdateWithSubquery).hasMessageContaining("Unexpected Join over for-update table scan");
+        abort("skipped");
+    }
+
+    @Test
+    @Override
+    protected void testUpdate()
+    {
+        // This simple update on JDBC ends up as a very simple, single-fragment, coordinator-only plan,
+        // which has no ability to recover from errors. This test simply verifies that's still the case.
+        Optional<String> setupQuery = Optional.of("CREATE TABLE <table> AS SELECT * FROM orders");
+        String testQuery = "UPDATE <table> SET shippriority = 101 WHERE custkey = 1";
+        Optional<String> cleanupQuery = Optional.of("DROP TABLE <table>");
+
+        assertThatQuery(testQuery)
+                .withSetupQuery(setupQuery)
+                .withCleanupQuery(cleanupQuery)
+                .isCoordinatorOnly();
+    }
+}

plugin/trino-mariadb/src/test/java/io/trino/plugin/mariadb/TestMariaDbQueryFailureRecovery.java

@@ -0,0 +1,25 @@
+/*
+ * Licensed 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 io.trino.plugin.mariadb;
+
+import io.trino.operator.RetryPolicy;
+
+public class TestMariaDbQueryFailureRecovery
+        extends BaseMariaDbFailureRecoveryTest
+{
+    public TestMariaDbQueryFailureRecovery()
+    {
+        super(RetryPolicy.QUERY);
+    }
+}

plugin/trino-mariadb/src/test/java/io/trino/plugin/mariadb/TestMariaDbTaskFailureRecovery.java

@@ -0,0 +1,25 @@
+/*
+ * Licensed 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 io.trino.plugin.mariadb;
+
+import io.trino.operator.RetryPolicy;
+
+public class TestMariaDbTaskFailureRecovery
+        extends BaseMariaDbFailureRecoveryTest
+{
+    public TestMariaDbTaskFailureRecovery()
+    {
+        super(RetryPolicy.TASK);
+    }
+}