Add procedures support for Lakehouse

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseConnector.java

@@ -26,6 +26,7 @@
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplitManager;
 import io.trino.spi.connector.ConnectorTransactionHandle;
+import io.trino.spi.procedure.Procedure;
 import io.trino.spi.session.PropertyMetadata;
 import io.trino.spi.transaction.IsolationLevel;
 
@@ -51,6 +52,7 @@ public class LakehouseConnector
     private final LakehouseSessionProperties sessionProperties;
     private final LakehouseTableProperties tableProperties;
     private final IcebergMaterializedViewProperties materializedViewProperties;
+    private final Set<Procedure> procedures;
 
     @Inject
     public LakehouseConnector(
@@ -62,7 +64,8 @@ public LakehouseConnector(
             LakehouseNodePartitioningProvider nodePartitioningProvider,
             LakehouseSessionProperties sessionProperties,
             LakehouseTableProperties tableProperties,
-            IcebergMaterializedViewProperties materializedViewProperties)
+            IcebergMaterializedViewProperties materializedViewProperties,
+            Set<Procedure> procedures)
     {
         this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null");
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
@@ -73,6 +76,7 @@ public LakehouseConnector(
         this.sessionProperties = requireNonNull(sessionProperties, "sessionProperties is null");
         this.tableProperties = requireNonNull(tableProperties, "tableProperties is null");
         this.materializedViewProperties = requireNonNull(materializedViewProperties, "materializedViewProperties is null");
+        this.procedures = requireNonNull(procedures, "procedures is null");
     }
 
     @Override
@@ -148,6 +152,12 @@ public List<PropertyMetadata<?>> getMaterializedViewProperties()
         return materializedViewProperties.getMaterializedViewProperties();
     }
 
+    @Override
+    public Set<Procedure> getProcedures()
+    {
+        return procedures;
+    }
+
     @Override
     public void shutdown()
     {

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseHiveModule.java

@@ -58,6 +58,7 @@
 import io.trino.plugin.hive.line.SimpleTextFileWriterFactory;
 import io.trino.plugin.hive.metastore.HiveMetastoreConfig;
 import io.trino.plugin.hive.metastore.HiveMetastoreModule;
+import io.trino.plugin.hive.metastore.glue.GlueCache;
 import io.trino.plugin.hive.orc.OrcFileWriterFactory;
 import io.trino.plugin.hive.orc.OrcPageSourceFactory;
 import io.trino.plugin.hive.parquet.ParquetFileWriterFactory;
@@ -67,6 +68,7 @@
 import java.util.Optional;
 
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConfigBinder.configBinder;
 import static io.airlift.json.JsonCodecBinder.jsonCodecBinder;
 import static org.weakref.jmx.guice.ExportBinder.newExporter;
@@ -138,5 +140,8 @@ protected void setup(Binder binder)
 
         binder.install(new HiveExecutorModule());
         install(new ParquetEncryptionModule());
+
+        newOptionalBinder(binder, GlueCache.class);
+        newOptionalBinder(binder, DirectoryLister.class);
     }
 }

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseModule.java

@@ -16,15 +16,33 @@
 import com.google.inject.Binder;
 import com.google.inject.Key;
 import com.google.inject.Scopes;
+import com.google.inject.multibindings.Multibinder;
 import io.airlift.configuration.AbstractConfigurationAwareModule;
 import io.trino.plugin.base.metrics.FileFormatDataSourceStats;
+import io.trino.plugin.deltalake.procedure.DropExtendedStatsProcedure;
+import io.trino.plugin.deltalake.procedure.FlushMetadataCacheProcedure;
+import io.trino.plugin.deltalake.procedure.RegisterTableProcedure;
+import io.trino.plugin.deltalake.procedure.UnregisterTableProcedure;
+import io.trino.plugin.deltalake.procedure.VacuumProcedure;
 import io.trino.plugin.hive.HideDeltaLakeTables;
 import io.trino.plugin.hive.SortingFileWriterConfig;
 import io.trino.plugin.hive.orc.OrcReaderConfig;
 import io.trino.plugin.hive.orc.OrcWriterConfig;
 import io.trino.plugin.hive.parquet.ParquetReaderConfig;
 import io.trino.plugin.hive.parquet.ParquetWriterConfig;
+import io.trino.plugin.hive.procedure.CreateEmptyPartitionProcedure;
+import io.trino.plugin.hive.procedure.DropStatsProcedure;
+import io.trino.plugin.hive.procedure.RegisterPartitionProcedure;
+import io.trino.plugin.hive.procedure.SyncPartitionMetadataProcedure;
+import io.trino.plugin.hive.procedure.UnregisterPartitionProcedure;
+import io.trino.plugin.iceberg.procedure.RollbackToSnapshotProcedure;
+import io.trino.plugin.lakehouse.procedures.LakehouseDropStatsProcedure;
+import io.trino.plugin.lakehouse.procedures.LakehouseFlushMetadataCacheProcedure;
+import io.trino.plugin.lakehouse.procedures.LakehouseRegisterTableProcedure;
+import io.trino.plugin.lakehouse.procedures.LakehouseUnregisterTableProcedure;
+import io.trino.spi.procedure.Procedure;
 
+import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static io.airlift.configuration.ConfigBinder.configBinder;
 import static org.weakref.jmx.guice.ExportBinder.newExporter;
 
@@ -53,6 +71,33 @@ protected void setup(Binder binder)
         binder.bind(FileFormatDataSourceStats.class).in(Scopes.SINGLETON);
         newExporter(binder).export(FileFormatDataSourceStats.class).withGeneratedName();
 
+        Multibinder<Procedure> procedures = newSetBinder(binder, Procedure.class);
+        // DeltaLake procedures
+        procedures.addBinding().toProvider(VacuumProcedure.class).in(Scopes.SINGLETON);
+        // Hive procedures
+        procedures.addBinding().toProvider(CreateEmptyPartitionProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(RegisterPartitionProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(UnregisterPartitionProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(SyncPartitionMetadataProcedure.class).in(Scopes.SINGLETON);
+        // Iceberg procedures
+        procedures.addBinding().toProvider(RollbackToSnapshotProcedure.class).in(Scopes.SINGLETON);
+        // Mixed procedures
+        binder.bind(DropExtendedStatsProcedure.class).in(Scopes.SINGLETON);
+        binder.bind(DropStatsProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(LakehouseDropStatsProcedure.class).in(Scopes.SINGLETON);
+
+        binder.bind(RegisterTableProcedure.class).in(Scopes.SINGLETON);
+        binder.bind(io.trino.plugin.iceberg.procedure.RegisterTableProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(LakehouseRegisterTableProcedure.class).in(Scopes.SINGLETON);
+
+        binder.bind(UnregisterTableProcedure.class).in(Scopes.SINGLETON);
+        binder.bind(io.trino.plugin.iceberg.procedure.UnregisterTableProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(LakehouseUnregisterTableProcedure.class).in(Scopes.SINGLETON);
+
+        binder.bind(FlushMetadataCacheProcedure.class).in(Scopes.SINGLETON);
+        binder.bind(io.trino.plugin.hive.procedure.FlushMetadataCacheProcedure.class).in(Scopes.SINGLETON);
+        procedures.addBinding().toProvider(LakehouseFlushMetadataCacheProcedure.class).in(Scopes.SINGLETON);
+
         binder.bind(Key.get(boolean.class, HideDeltaLakeTables.class)).toInstance(false);
     }
 }

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/procedures/LakehouseDropStatsProcedure.java

@@ -0,0 +1,104 @@
+/*
+ * 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.lakehouse.procedures;
+
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Inject;
+import com.google.inject.Provider;
+import io.trino.plugin.deltalake.procedure.DropExtendedStatsProcedure;
+import io.trino.plugin.hive.procedure.DropStatsProcedure;
+import io.trino.plugin.lakehouse.TableType;
+import io.trino.spi.connector.ConnectorAccessControl;
+import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.procedure.Procedure;
+import io.trino.spi.type.ArrayType;
+
+import java.lang.invoke.MethodHandle;
+import java.util.List;
+
+import static io.trino.spi.type.VarcharType.VARCHAR;
+import static java.lang.invoke.MethodHandles.lookup;
+import static java.util.Objects.requireNonNull;
+
+/**
+ * A procedure that drops statistics.
+ * <p>
+ * It is delegated to the appropriate underlying procedure based on the table type.
+ * Currently, it supports Delta Lake and Hive table types.
+ */
+public class LakehouseDropStatsProcedure
+        implements Provider<Procedure>
+{
+    private static final MethodHandle DROP_STATS;
+
+    private static final String SYSTEM_SCHEMA = "system";
+    private static final String PROCEDURE_NAME = "drop_stats";
+
+    private static final String TABLE_TYPE = "TABLE_TYPE";
+    private static final String SCHEMA_NAME = "SCHEMA_NAME";
+    private static final String TABLE_NAME = "TABLE_NAME";
+    private static final String PARTITION_VALUES = "PARTITION_VALUES";
+
+    static {
+        try {
+            DROP_STATS = lookup().unreflect(LakehouseDropStatsProcedure.class.getMethod(
+                    "dropStats", ConnectorSession.class, ConnectorAccessControl.class, String.class, String.class, String.class, List.class));
+        }
+        catch (ReflectiveOperationException e) {
+            throw new AssertionError(e);
+        }
+    }
+
+    private final DropExtendedStatsProcedure deltaLakeDropStatsProcedure;
+    private final DropStatsProcedure hiveDropStatsProcedure;
+
+    @Inject
+    public LakehouseDropStatsProcedure(
+            DropExtendedStatsProcedure deltaLakeDropStatsProcedure,
+            DropStatsProcedure hiveDropStatsProcedure)
+    {
+        this.deltaLakeDropStatsProcedure = requireNonNull(deltaLakeDropStatsProcedure, "deltaLakeDropStatsProcedure is null");
+        this.hiveDropStatsProcedure = requireNonNull(hiveDropStatsProcedure, "hiveDropStatsProcedure is null");
+    }
+
+    @Override
+    public Procedure get()
+    {
+        return new Procedure(
+                SYSTEM_SCHEMA,
+                PROCEDURE_NAME,
+                ImmutableList.of(
+                        new Procedure.Argument(TABLE_TYPE, VARCHAR),
+                        new Procedure.Argument(SCHEMA_NAME, VARCHAR),
+                        new Procedure.Argument(TABLE_NAME, VARCHAR),
+                        new Procedure.Argument(PARTITION_VALUES, new ArrayType(new ArrayType(VARCHAR)), false, null)),
+                DROP_STATS.bindTo(this));
+    }
+
+    public void dropStats(ConnectorSession session, ConnectorAccessControl accessControl, String tableType, String schema, String table, List<?> partitionValues)
+    {
+        if (TableType.DELTA.name().equals(tableType)) {
+            if (partitionValues != null) {
+                throw new IllegalArgumentException("Partition values are not supported for Delta Lake procedure");
+            }
+            deltaLakeDropStatsProcedure.dropStats(session, accessControl, schema, table);
+        }
+        else if (TableType.HIVE.name().equals(tableType)) {
+            hiveDropStatsProcedure.dropStats(session, accessControl, schema, table, partitionValues);
+        }
+        else {
+            throw new IllegalArgumentException("Unsupported table type: " + tableType);
+        }
+    }
+}

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/procedures/LakehouseFlushMetadataCacheProcedure.java

@@ -0,0 +1,107 @@
+/*
+ * 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.lakehouse.procedures;
+
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Inject;
+import com.google.inject.Provider;
+import io.trino.plugin.deltalake.procedure.FlushMetadataCacheProcedure;
+import io.trino.plugin.lakehouse.TableType;
+import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.procedure.Procedure;
+import io.trino.spi.type.ArrayType;
+
+import java.lang.invoke.MethodHandle;
+import java.util.List;
+
+import static io.trino.spi.type.VarcharType.VARCHAR;
+import static java.lang.invoke.MethodHandles.lookup;
+import static java.util.Objects.requireNonNull;
+
+/**
+ * A procedure that flushes the metadata cache for a table or a specific partition of a table.
+ * <p>
+ * It is delegated to the appropriate underlying procedure based on the table type.
+ * Currently, it supports Delta Lake and Hive table types.
+ */
+public class LakehouseFlushMetadataCacheProcedure
+        implements Provider<Procedure>
+{
+    private static final MethodHandle FLUSH_METADATA_CACHE;
+
+    private static final String SYSTEM_SCHEMA = "system";
+    private static final String PROCEDURE_NAME = "flush_metadata_cache";
+
+    private static final String TABLE_TYPE = "TABLE_TYPE";
+    private static final String SCHEMA_NAME = "SCHEMA_NAME";
+    private static final String TABLE_NAME = "TABLE_NAME";
+    private static final String PARAM_PARTITION_COLUMNS = "PARTITION_COLUMNS";
+    private static final String PARAM_PARTITION_VALUES = "PARTITION_VALUES";
+
+    static {
+        try {
+            FLUSH_METADATA_CACHE = lookup().unreflect(LakehouseFlushMetadataCacheProcedure.class.getMethod(
+                    "flushMetadataCache", ConnectorSession.class, String.class, String.class, String.class, List.class, List.class));
+        }
+        catch (ReflectiveOperationException e) {
+            throw new AssertionError(e);
+        }
+    }
+
+    private final FlushMetadataCacheProcedure deltaLakeFlushMetadataCacheProcedure;
+    private final io.trino.plugin.hive.procedure.FlushMetadataCacheProcedure hiveFlushMetadataCacheProcedure;
+
+    @Inject
+    public LakehouseFlushMetadataCacheProcedure(
+            FlushMetadataCacheProcedure deltaLakeFlushMetadataCacheProcedure,
+            io.trino.plugin.hive.procedure.FlushMetadataCacheProcedure hiveFlushMetadataCacheProcedure)
+    {
+        this.deltaLakeFlushMetadataCacheProcedure = requireNonNull(deltaLakeFlushMetadataCacheProcedure, "deltaLakeFlushMetadataCacheProcedure is null");
+        this.hiveFlushMetadataCacheProcedure = requireNonNull(hiveFlushMetadataCacheProcedure, "hiveFlushMetadataCacheProcedure is null");
+    }
+
+    @Override
+    public Procedure get()
+    {
+        return new Procedure(
+                SYSTEM_SCHEMA,
+                PROCEDURE_NAME,
+                ImmutableList.of(
+                        new Procedure.Argument(TABLE_TYPE, VARCHAR),
+                        new Procedure.Argument(SCHEMA_NAME, VARCHAR),
+                        new Procedure.Argument(TABLE_NAME, VARCHAR),
+                        new Procedure.Argument(PARAM_PARTITION_COLUMNS, new ArrayType(VARCHAR), false, null),
+                        new Procedure.Argument(PARAM_PARTITION_VALUES, new ArrayType(VARCHAR), false, null)),
+                FLUSH_METADATA_CACHE.bindTo(this));
+    }
+
+    public void flushMetadataCache(ConnectorSession session, String tableType, String schema, String table, List<String> partitionColumns, List<String> partitionValues)
+    {
+        if (TableType.DELTA.name().equals(tableType)) {
+            if (partitionColumns != null && !partitionColumns.isEmpty()) {
+                throw new IllegalArgumentException("Partition columns are not supported for Delta Lake tables");
+            }
+            if (partitionValues != null && !partitionValues.isEmpty()) {
+                throw new IllegalArgumentException("Partition values are not supported for Delta Lake tables");
+            }
+            deltaLakeFlushMetadataCacheProcedure.flushMetadataCache(schema, table);
+        }
+        else if (TableType.HIVE.name().equals(tableType)) {
+            hiveFlushMetadataCacheProcedure.flushMetadataCache(session, schema, table, partitionColumns, partitionValues);
+        }
+        else {
+            throw new IllegalArgumentException("Unsupported table type: " + tableType);
+        }
+    }
+}