Ensure cleanupQuery gets called for procedures

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/util/UncheckedCloseable.java

@@ -0,0 +1,21 @@
+/*
+ * 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.base.util;
+
+public interface UncheckedCloseable
+        extends AutoCloseable
+{
+    @Override
+    void close();
+}

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/functions/tablechanges/TableChangesFunction.java

@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import io.airlift.slice.Slice;
+import io.trino.plugin.base.util.UncheckedCloseable;
 import io.trino.plugin.deltalake.CorruptedDeltaLakeTableHandle;
 import io.trino.plugin.deltalake.DeltaLakeColumnHandle;
 import io.trino.plugin.deltalake.DeltaLakeMetadata;
@@ -104,42 +105,45 @@ public TableFunctionAnalysis analyze(
         long firstReadVersion = sinceVersion + 1; // +1 to ensure that the since_version is exclusive; may overflow
 
         DeltaLakeMetadata deltaLakeMetadata = deltaLakeMetadataFactory.create(session.getIdentity());
-        SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName);
-        ConnectorTableHandle connectorTableHandle = deltaLakeMetadata.getTableHandle(session, schemaTableName);
-        if (connectorTableHandle == null) {
-            throw new TableNotFoundException(schemaTableName);
-        }
-        if (connectorTableHandle instanceof CorruptedDeltaLakeTableHandle corruptedTableHandle) {
-            throw corruptedTableHandle.createException();
-        }
-        DeltaLakeTableHandle tableHandle = (DeltaLakeTableHandle) connectorTableHandle;
+        deltaLakeMetadata.beginQuery(session);
+        try (UncheckedCloseable ignore = () -> deltaLakeMetadata.cleanupQuery(session)) {
+            SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName);
+            ConnectorTableHandle connectorTableHandle = deltaLakeMetadata.getTableHandle(session, schemaTableName);
+            if (connectorTableHandle == null) {
+                throw new TableNotFoundException(schemaTableName);
+            }
+            if (connectorTableHandle instanceof CorruptedDeltaLakeTableHandle corruptedTableHandle) {
+                throw corruptedTableHandle.createException();
+            }
+            DeltaLakeTableHandle tableHandle = (DeltaLakeTableHandle) connectorTableHandle;
 
-        if (sinceVersion > tableHandle.getReadVersion()) {
-            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, format("since_version: %d is higher then current table version: %d", sinceVersion, tableHandle.getReadVersion()));
-        }
-        List<DeltaLakeColumnHandle> columnHandles = deltaLakeMetadata.getColumnHandles(session, tableHandle)
-                .values().stream()
-                .map(DeltaLakeColumnHandle.class::cast)
-                .filter(column -> column.getColumnType() != SYNTHESIZED)
-                .collect(toImmutableList());
-        accessControl.checkCanSelectFromColumns(null, schemaTableName, columnHandles.stream()
-                // Lowercase column names because users don't know the original names
-                .map(column -> column.getColumnName().toLowerCase(ENGLISH))
-                .collect(toImmutableSet()));
+            if (sinceVersion > tableHandle.getReadVersion()) {
+                throw new TrinoException(INVALID_FUNCTION_ARGUMENT, format("since_version: %d is higher then current table version: %d", sinceVersion, tableHandle.getReadVersion()));
+            }
+            List<DeltaLakeColumnHandle> columnHandles = deltaLakeMetadata.getColumnHandles(session, tableHandle)
+                    .values().stream()
+                    .map(DeltaLakeColumnHandle.class::cast)
+                    .filter(column -> column.getColumnType() != SYNTHESIZED)
+                    .collect(toImmutableList());
+            accessControl.checkCanSelectFromColumns(null, schemaTableName, columnHandles.stream()
+                    // Lowercase column names because users don't know the original names
+                    .map(column -> column.getColumnName().toLowerCase(ENGLISH))
+                    .collect(toImmutableSet()));
 
-        ImmutableList.Builder<Descriptor.Field> outputFields = ImmutableList.builder();
-        columnHandles.stream()
-                .map(columnHandle -> new Descriptor.Field(columnHandle.getColumnName(), Optional.of(columnHandle.getType())))
-                .forEach(outputFields::add);
+            ImmutableList.Builder<Descriptor.Field> outputFields = ImmutableList.builder();
+            columnHandles.stream()
+                    .map(columnHandle -> new Descriptor.Field(columnHandle.getColumnName(), Optional.of(columnHandle.getType())))
+                    .forEach(outputFields::add);
 
-        // add at the end to follow Delta Lake convention
-        outputFields.add(new Descriptor.Field(CHANGE_TYPE_COLUMN_NAME, Optional.of(VARCHAR)));
-        outputFields.add(new Descriptor.Field(COMMIT_VERSION_COLUMN_NAME, Optional.of(BIGINT)));
-        outputFields.add(new Descriptor.Field(COMMIT_TIMESTAMP_COLUMN_NAME, Optional.of(TIMESTAMP_TZ_MILLIS)));
+            // add at the end to follow Delta Lake convention
+            outputFields.add(new Descriptor.Field(CHANGE_TYPE_COLUMN_NAME, Optional.of(VARCHAR)));
+            outputFields.add(new Descriptor.Field(COMMIT_VERSION_COLUMN_NAME, Optional.of(BIGINT)));
+            outputFields.add(new Descriptor.Field(COMMIT_TIMESTAMP_COLUMN_NAME, Optional.of(TIMESTAMP_TZ_MILLIS)));
 
-        return TableFunctionAnalysis.builder()
-                .handle(new TableChangesTableFunctionHandle(schemaTableName, firstReadVersion, tableHandle.getReadVersion(), tableHandle.getLocation(), columnHandles))
-                .returnedType(new Descriptor(outputFields.build()))
-                .build();
+            return TableFunctionAnalysis.builder()
+                    .handle(new TableChangesTableFunctionHandle(schemaTableName, firstReadVersion, tableHandle.getReadVersion(), tableHandle.getLocation(), columnHandles))
+                    .returnedType(new Descriptor(outputFields.build()))
+                    .build();
+        }
     }
 }

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/procedure/DropExtendedStatsProcedure.java

@@ -15,6 +15,7 @@
 
 import com.google.inject.Inject;
 import com.google.inject.Provider;
+import io.trino.plugin.base.util.UncheckedCloseable;
 import io.trino.plugin.deltalake.DeltaLakeMetadata;
 import io.trino.plugin.deltalake.DeltaLakeMetadataFactory;
 import io.trino.plugin.deltalake.LocatedTableHandle;
@@ -79,11 +80,14 @@ public void dropStats(ConnectorSession session, ConnectorAccessControl accessCon
 
         SchemaTableName name = new SchemaTableName(schema, table);
         DeltaLakeMetadata metadata = metadataFactory.create(session.getIdentity());
-        LocatedTableHandle tableHandle = metadata.getTableHandle(session, name);
-        if (tableHandle == null) {
-            throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, format("Table '%s' does not exist", name));
+        metadata.beginQuery(session);
+        try (UncheckedCloseable ignore = () -> metadata.cleanupQuery(session)) {
+            LocatedTableHandle tableHandle = metadata.getTableHandle(session, name);
+            if (tableHandle == null) {
+                throw new TrinoException(INVALID_PROCEDURE_ARGUMENT, format("Table '%s' does not exist", name));
+            }
+            accessControl.checkCanInsertIntoTable(null, name);
+            statsAccess.deleteExtendedStatistics(session, name, tableHandle.location());
         }
-        accessControl.checkCanInsertIntoTable(null, name);
-        statsAccess.deleteExtendedStatistics(session, name, tableHandle.location());
     }
 }

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/procedure/RegisterTableProcedure.java

@@ -19,7 +19,9 @@
 import io.trino.filesystem.Location;
 import io.trino.filesystem.TrinoFileSystem;
 import io.trino.filesystem.TrinoFileSystemFactory;
+import io.trino.plugin.base.util.UncheckedCloseable;
 import io.trino.plugin.deltalake.DeltaLakeConfig;
+import io.trino.plugin.deltalake.DeltaLakeMetadata;
 import io.trino.plugin.deltalake.DeltaLakeMetadataFactory;
 import io.trino.plugin.deltalake.metastore.DeltaLakeMetastore;
 import io.trino.plugin.deltalake.statistics.CachingExtendedStatisticsAccess;
@@ -140,60 +142,64 @@ private void doRegisterTable(
         checkProcedureArgument(!isNullOrEmpty(tableLocation), "table_location cannot be null or empty");
 
         SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName);
-        DeltaLakeMetastore metastore = metadataFactory.create(session.getIdentity()).getMetastore();
+        DeltaLakeMetadata metadata = metadataFactory.create(session.getIdentity());
+        metadata.beginQuery(session);
+        try (UncheckedCloseable ignore = () -> metadata.cleanupQuery(session)) {
+            DeltaLakeMetastore metastore = metadata.getMetastore();
 
-        if (metastore.getDatabase(schemaName).isEmpty()) {
-            throw new SchemaNotFoundException(schemaTableName.getSchemaName());
-        }
-
-        TrinoFileSystem fileSystem = fileSystemFactory.create(session);
-        try {
-            Location transactionLogDir = Location.of(getTransactionLogDir(tableLocation));
-            if (!fileSystem.listFiles(transactionLogDir).hasNext()) {
-                throw new TrinoException(GENERIC_USER_ERROR, format("No transaction log found in location %s", transactionLogDir));
+            if (metastore.getDatabase(schemaName).isEmpty()) {
+                throw new SchemaNotFoundException(schemaTableName.getSchemaName());
             }
-        }
-        catch (IOException e) {
-            throw new TrinoException(DELTA_LAKE_FILESYSTEM_ERROR, format("Failed checking table location %s", tableLocation), e);
-        }
 
-        Table table = buildTable(session, schemaTableName, tableLocation, true);
+            TrinoFileSystem fileSystem = fileSystemFactory.create(session);
+            try {
+                Location transactionLogDir = Location.of(getTransactionLogDir(tableLocation));
+                if (!fileSystem.listFiles(transactionLogDir).hasNext()) {
+                    throw new TrinoException(GENERIC_USER_ERROR, format("No transaction log found in location %s", transactionLogDir));
+                }
+            }
+            catch (IOException e) {
+                throw new TrinoException(DELTA_LAKE_FILESYSTEM_ERROR, format("Failed checking table location %s", tableLocation), e);
+            }
 
-        PrincipalPrivileges principalPrivileges = buildInitialPrivilegeSet(table.getOwner().orElseThrow());
-        statisticsAccess.invalidateCache(schemaTableName, Optional.of(tableLocation));
-        transactionLogAccess.invalidateCache(schemaTableName, Optional.of(tableLocation));
-        // Verify we're registering a location with a valid table
-        try {
-            TableSnapshot tableSnapshot = transactionLogAccess.loadSnapshot(session, table.getSchemaTableName(), tableLocation);
-            transactionLogAccess.getMetadataEntry(tableSnapshot, session); // verify metadata exists
-        }
-        catch (TrinoException e) {
-            throw e;
-        }
-        catch (IOException | RuntimeException e) {
-            throw new TrinoException(DELTA_LAKE_INVALID_TABLE, "Failed to access table location: " + tableLocation, e);
-        }
+            Table table = buildTable(session, schemaTableName, tableLocation, true);
 
-        // Ensure the table has queryId set. This is relied on for exception handling
-        String queryId = session.getQueryId();
-        verify(
-                getQueryId(table).orElseThrow(() -> new IllegalArgumentException("Query id is not present")).equals(queryId),
-                "Table '%s' does not have correct query id set",
-                table);
-        try {
-            metastore.createTable(
-                    session,
-                    table,
-                    principalPrivileges);
-        }
-        catch (TableAlreadyExistsException e) {
-            // Ignore TableAlreadyExistsException when table looks like created by us.
-            // This may happen when an actually successful metastore create call is retried
-            // e.g. because of a timeout on our side.
-            Optional<Table> existingTable = metastore.getRawMetastoreTable(schemaName, tableName);
-            if (existingTable.isEmpty() || !isCreatedBy(existingTable.get(), queryId)) {
+            PrincipalPrivileges principalPrivileges = buildInitialPrivilegeSet(table.getOwner().orElseThrow());
+            statisticsAccess.invalidateCache(schemaTableName, Optional.of(tableLocation));
+            transactionLogAccess.invalidateCache(schemaTableName, Optional.of(tableLocation));
+            // Verify we're registering a location with a valid table
+            try {
+                TableSnapshot tableSnapshot = transactionLogAccess.loadSnapshot(session, table.getSchemaTableName(), tableLocation);
+                transactionLogAccess.getMetadataEntry(tableSnapshot, session); // verify metadata exists
+            }
+            catch (TrinoException e) {
                 throw e;
             }
+            catch (IOException | RuntimeException e) {
+                throw new TrinoException(DELTA_LAKE_INVALID_TABLE, "Failed to access table location: " + tableLocation, e);
+            }
+
+            // Ensure the table has queryId set. This is relied on for exception handling
+            String queryId = session.getQueryId();
+            verify(
+                    getQueryId(table).orElseThrow(() -> new IllegalArgumentException("Query id is not present")).equals(queryId),
+                    "Table '%s' does not have correct query id set",
+                    table);
+            try {
+                metastore.createTable(
+                        session,
+                        table,
+                        principalPrivileges);
+            }
+            catch (TableAlreadyExistsException e) {
+                // Ignore TableAlreadyExistsException when table looks like created by us.
+                // This may happen when an actually successful metastore create call is retried
+                // e.g. because of a timeout on our side.
+                Optional<Table> existingTable = metastore.getRawMetastoreTable(schemaName, tableName);
+                if (existingTable.isEmpty() || !isCreatedBy(existingTable.get(), queryId)) {
+                    throw e;
+                }
+            }
         }
     }
 }

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/procedure/UnregisterTableProcedure.java

@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 import com.google.inject.Provider;
+import io.trino.plugin.base.util.UncheckedCloseable;
 import io.trino.plugin.deltalake.DeltaLakeMetadata;
 import io.trino.plugin.deltalake.DeltaLakeMetadataFactory;
 import io.trino.plugin.deltalake.LocatedTableHandle;
@@ -96,14 +97,16 @@ private void doUnregisterTable(ConnectorAccessControl accessControl, ConnectorSe
 
         accessControl.checkCanDropTable(null, schemaTableName);
         DeltaLakeMetadata metadata = metadataFactory.create(session.getIdentity());
-
-        LocatedTableHandle tableHandle = metadata.getTableHandle(session, schemaTableName);
-        if (tableHandle == null) {
-            throw new TableNotFoundException(schemaTableName);
+        metadata.beginQuery(session);
+        try (UncheckedCloseable ignore = () -> metadata.cleanupQuery(session)) {
+            LocatedTableHandle tableHandle = metadata.getTableHandle(session, schemaTableName);
+            if (tableHandle == null) {
+                throw new TableNotFoundException(schemaTableName);
+            }
+            metadata.getMetastore().dropTable(session, schemaTableName, tableHandle.location(), false);
+            // As a precaution, clear the caches
+            statisticsAccess.invalidateCache(schemaTableName, Optional.of(tableHandle.location()));
+            transactionLogAccess.invalidateCache(schemaTableName, Optional.of(tableHandle.location()));
         }
-        metadata.getMetastore().dropTable(session, schemaTableName, tableHandle.location(), false);
-        // As a precaution, clear the caches
-        statisticsAccess.invalidateCache(schemaTableName, Optional.of(tableHandle.location()));
-        transactionLogAccess.invalidateCache(schemaTableName, Optional.of(tableHandle.location()));
     }
 }