Retry listing operations on deadlock in SQL Server

Author: ebyhr

plugin/trino-sqlserver/src/main/java/io/trino/plugin/sqlserver/SqlServerClient.java

@@ -77,6 +77,8 @@
 import io.trino.spi.connector.JoinCondition;
 import io.trino.spi.connector.JoinStatistics;
 import io.trino.spi.connector.JoinType;
+import io.trino.spi.connector.RelationColumnsMetadata;
+import io.trino.spi.connector.RelationCommentMetadata;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.connector.TableNotFoundException;
 import io.trino.spi.expression.ConnectorExpression;
@@ -121,6 +123,7 @@
 import java.time.format.DateTimeFormatterBuilder;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -428,6 +431,24 @@ protected boolean isTableLockNeeded(ConnectorSession session)
         return isBulkCopyForWrite(session) && isBulkCopyForWriteLockDestinationTable(session);
     }
 
+    @Override
+    public List<RelationCommentMetadata> getAllTableComments(ConnectorSession session, Optional<String> schema)
+    {
+        return retryOnDeadlock(() -> super.getAllTableComments(session, schema), "error when getting all table comments for '%s'".formatted(schema));
+    }
+
+    @Override
+    public Iterator<RelationColumnsMetadata> getAllTableColumns(ConnectorSession session, Optional<String> schema)
+    {
+        return retryOnDeadlock(() -> super.getAllTableColumns(session, schema), "error when getting all table columns for '%s'".formatted(schema));
+    }
+
+    @Override
+    public List<JdbcColumnHandle> getColumns(ConnectorSession session, SchemaTableName schemaTableName, RemoteTableName remoteTableName)
+    {
+        return retryOnDeadlock(() -> super.getColumns(session, schemaTableName, remoteTableName), "error when getting columns for table '%s'".formatted(remoteTableName));
+    }
+
     @Override
     protected void verifyTableName(DatabaseMetaData databaseMetadata, String tableName)
             throws SQLException

plugin/trino-sqlserver/src/test/java/io/trino/plugin/sqlserver/TestSqlServerConnectorTest.java

@@ -276,28 +276,4 @@ private List<String> testTableNameTestData()
                 .add("close]bracket")
                 .build();
     }
-
-    @Test
-    @Override
-    public void testSelectInformationSchemaTables()
-    {
-        // Isolate this test to avoid problem described in https://github.com/trinodb/trino/issues/10846
-        executeExclusively(super::testSelectInformationSchemaTables);
-    }
-
-    @Test
-    @Override
-    public void testSelectInformationSchemaColumns()
-    {
-        // Isolate this test to avoid problem described in https://github.com/trinodb/trino/issues/10846
-        executeExclusively(super::testSelectInformationSchemaColumns);
-    }
-
-    @Test
-    @Override
-    public void testBulkColumnListingOptions()
-    {
-        // Isolate this test to avoid problem described in https://github.com/trinodb/trino/issues/10846
-        executeExclusively(super::testBulkColumnListingOptions);
-    }
 }