Add support for TRUNCATE statement in memory connector

trinodb · Jun 10, 2024 · 22da322 · 22da322
1 parent 08fb411
commit 22da322
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/docs/src/main/sphinx/connector/memory.md b/docs/src/main/sphinx/connector/memory.md
@@ -59,6 +59,7 @@ stored in memory. In addition to the {ref}`globally available
 statements, the connector supports the following features:
 
 - {doc}`/sql/insert`
+- {doc}`/sql/truncate`
 - {doc}`/sql/create-table`
 - {doc}`/sql/create-table-as`
 - {doc}`/sql/drop-table`
@@ -67,9 +68,9 @@ statements, the connector supports the following features:
 - {doc}`/sql/comment`
 - [](sql-routine-management)
 
-### DROP TABLE
+### TRUNCATE and DROP TABLE
 
-Upon execution of a `DROP TABLE` operation, memory is not released
+Upon execution of a `TRUNCATE` and a `DROP TABLE` operation, memory is not released
 immediately. It is instead released after the next write operation to the
 catalog.
 

diff --git a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java
@@ -368,6 +368,15 @@ public synchronized Optional<ConnectorOutputMetadata> finishInsert(
         return Optional.empty();
     }
 
+    @Override
+    public synchronized void truncateTable(ConnectorSession session, ConnectorTableHandle tableHandle)
+    {
+        MemoryTableHandle handle = (MemoryTableHandle) tableHandle;
+        long tableId = handle.id();
+        TableInfo info = tables.get(handle.id());
+        tables.put(tableId, new TableInfo(tableId, info.schemaName(), info.tableName(), info.columns(), ImmutableMap.of(), info.comment()));
+    }
+
     @Override
     public synchronized void createView(ConnectorSession session, SchemaTableName viewName, ConnectorViewDefinition definition, boolean replace)
     {

diff --git a/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java b/plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryConnectorTest.java
@@ -86,6 +86,7 @@ protected QueryRunner createQueryRunner()
     protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
     {
         return switch (connectorBehavior) {
+            case SUPPORTS_TRUNCATE -> true;
             case SUPPORTS_ADD_COLUMN,
                  SUPPORTS_AGGREGATION_PUSHDOWN,
                  SUPPORTS_CREATE_MATERIALIZED_VIEW,