Test IR print with QueryRunner

Author: kasiafi

testing/trino-testing/src/main/java/io/trino/testing/AbstractDistributedEngineOnlyQueries.java

@@ -403,4 +403,29 @@ public void testRowConstructorColumnLimit()
         @Language("SQL") String query = "SELECT row(" + rowFields + ") FROM (select * from tpch.tiny.orders limit 1) t(" + colNames + ")";
         assertThat(getQueryRunner().execute(query).getOnlyValue()).isNotNull();
     }
+
+    @Test
+    public void testProgramAssembly()
+    {
+        assertAssembly(
+                "SELECT name FROM nation",
+                """
+                IR version = 1
+                %0 = query() : () -> "boolean" ({
+                    ^query
+                        %1 = table_scan() : () -> "multiset(row(""f_1"" varchar(25)))" ()
+                            {table_handle = "{""catalogHandle"":""memory:normal:21a29a35ed877cb4ea566f4b08371b5a0a3c0588f07ba25bd0881def395049cd"",""connectorHandle"":{""@type"":""system:io.trino.plugin.memory.MemoryTableHandle"",""id"":2},""transaction"":[""system:io.trino.plugin.memory.MemoryTransactionHandle"",""INSTANCE""]}", column_handles = "[{""@type"":""system:io.trino.plugin.memory.MemoryColumnHandle"",""columnIndex"":1,""type"":""varchar(25)""}]", constraint = "{""columnDomains"":[]}", update_target = "false", use_connector_node_partitioning = "false"}
+                        %2 = output(%1) : ("multiset(row(""f_1"" varchar(25)))") -> "boolean" ({
+                            ^outputFieldSelector (%3 : "row(""f_1"" varchar(25))")
+                                %4 = field_selection(%3) : ("row(""f_1"" varchar(25))") -> "varchar(25)" ()
+                                    {field_name = "f_1"}
+                                %5 = row(%4) : ("varchar(25)") -> "row(varchar(25))" ()
+                                %6 = return(%5) : ("row(varchar(25))") -> "row(varchar(25))" ()
+                                    {ir.terminal = "true"}
+                            })
+                            {output_names = "[""name""]", ir.terminal = "true"}
+                    })
+                    {ir.terminal = "true"}
+                """);
+    }
 }

testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java

@@ -18,6 +18,7 @@
 import com.google.common.collect.MoreCollectors;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
+import com.google.inject.Key;
 import io.airlift.log.Level;
 import io.airlift.log.Logging;
 import io.airlift.units.Duration;
@@ -47,6 +48,8 @@
 import io.trino.spi.connector.CatalogSchemaTableName;
 import io.trino.spi.type.Type;
 import io.trino.sql.analyzer.QueryExplainer;
+import io.trino.sql.dialect.trino.ProgramBuilder;
+import io.trino.sql.newir.FormatOptions;
 import io.trino.sql.parser.SqlParser;
 import io.trino.sql.planner.OptimizerConfig.JoinDistributionType;
 import io.trino.sql.planner.Plan;
@@ -595,6 +598,13 @@ protected void assertNoDataRead(@Language("SQL") String sql)
                 results -> assertThat(results.getRowCount()).isEqualTo(0));
     }
 
+    protected void assertAssembly(@Language("SQL") String sql, String expectedAssembly)
+    {
+        Plan plan = queryRunner.executeWithPlan(getSession(), sql).queryPlan().orElseThrow();
+        String actualAssembly = ProgramBuilder.buildProgram(plan.getRoot()).print(1, getQueryRunner().getCoordinator().getInstance(Key.get(FormatOptions.class)));
+        assertThat(actualAssembly).isEqualTo(expectedAssembly);
+    }
+
     protected MaterializedResult computeExpected(@Language("SQL") String sql, List<? extends Type> resultTypes)
     {
         return h2QueryRunner.execute(getSession(), sql, resultTypes);