Add MATERIALIZED support for postgres CTEs

internal/jet/with_statement.go

@@ -58,11 +58,23 @@ func (w withImpl) projections() ProjectionList {
 	return ProjectionList{}
 }
 
+// CTEMaterialization is the materialization of a common table expression.
+type CTEMaterialization int
+
+const (
+	// CTEMaterializationDefault leaves materialization to the database.
+	CTEMaterializationDefault CTEMaterialization = iota
+	// CTEMaterializationForced forces materialization (MATERIALIZED).
+	CTEMaterializationForced
+	// CTEMaterializationDisabled disables materialization (NOT MATERIALIZED).
+	CTEMaterializationDisabled
+)
+
 // CommonTableExpression contains information about a CTE.
 type CommonTableExpression struct {
 	selectTableImpl
 
-	NotMaterialized bool
+	Materialization CTEMaterialization
 	Columns         []ColumnExpression
 }
 
@@ -90,8 +102,11 @@ func (c CommonTableExpression) serialize(statement StatementType, out *SQLBuilde
 		}
 		out.WriteString("AS")
 
-		if c.NotMaterialized {
-			out.WriteString("NOT MATERIALIZED")
+		switch c.Materialization {
+		case CTEMaterializationForced:
+			out.WriteString(" MATERIALIZED")
+		case CTEMaterializationDisabled:
+			out.WriteString(" NOT MATERIALIZED")
 		}
 
 		if c.Statement == nil {

postgres/with_statement.go

@@ -7,7 +7,8 @@ type CommonTableExpression interface {
 	SelectTable
 
 	AS(statement jet.SerializerHasProjections) CommonTableExpression
-	AS_NOT_MATERIALIZED(statement jet.SerializerStatement) CommonTableExpression
+	AS_MATERIALIZED(statement jet.SerializerHasProjections) CommonTableExpression
+	AS_NOT_MATERIALIZED(statement jet.SerializerHasProjections) CommonTableExpression
 	// ALIAS is used to create another alias of the CTE, if a CTE needs to appear multiple times in the main query.
 	ALIAS(alias string) SelectTable
 
@@ -47,9 +48,16 @@ func (c *commonTableExpression) AS(statement jet.SerializerHasProjections) Commo
 	return c
 }
 
+// AS_MATERIALIZED is used to define a materialized CTE query
+func (c *commonTableExpression) AS_MATERIALIZED(statement jet.SerializerHasProjections) CommonTableExpression {
+	c.CommonTableExpression.Materialization = jet.CTEMaterializationForced
+	c.CommonTableExpression.Statement = statement
+	return c
+}
+
 // AS_NOT_MATERIALIZED is used to define not materialized CTE query
-func (c *commonTableExpression) AS_NOT_MATERIALIZED(statement jet.SerializerStatement) CommonTableExpression {
-	c.CommonTableExpression.NotMaterialized = true
+func (c *commonTableExpression) AS_NOT_MATERIALIZED(statement jet.SerializerHasProjections) CommonTableExpression {
+	c.CommonTableExpression.Materialization = jet.CTEMaterializationDisabled
 	c.CommonTableExpression.Statement = statement
 	return c
 }

postgres/with_statement_test.go

@@ -0,0 +1,61 @@
+package postgres
+
+import (
+	"testing"
+
+	"github.com/go-jet/jet/v2/internal/testutils"
+)
+
+func TestCTE_Materialized(t *testing.T) {
+	col := IntegerColumn("colo")
+	tbl := NewTable("db", "t", "", col)
+
+	c1 := CTE("c1")
+	c2 := c1.ALIAS("c2")
+
+	stmt := WITH(
+		c1.AS_MATERIALIZED(
+			SELECT(col).FROM(tbl),
+		),
+	)(
+		SELECT(
+			c1.AllColumns().As("c1.*"),
+			c2.AllColumns().As("c2.*"),
+		).FROM(
+			c1.INNER_JOIN(c2, Bool(true)),
+		),
+	)
+
+	testutils.AssertStatementSql(t, stmt, `
+WITH c1 AS MATERIALIZED (
+     SELECT t.colo AS "t.colo"
+     FROM db.t
+)
+SELECT c1."t.colo" AS "c1.colo",
+     c2."t.colo" AS "c2.colo"
+FROM c1
+     INNER JOIN c1 AS c2 ON $1::boolean;
+`)
+}
+
+func TestCTE_NotMaterialized(t *testing.T) {
+	col := IntegerColumn("colo")
+	tbl := NewTable("db", "t", "", col)
+
+	n := CTE("n1")
+
+	stmt := WITH(
+		n.AS_NOT_MATERIALIZED(SELECT(col).FROM(tbl)),
+	)(
+		SELECT(n.AllColumns().As("n1.*")).FROM(n),
+	)
+
+	testutils.AssertStatementSql(t, stmt, `
+WITH n1 AS NOT MATERIALIZED (
+     SELECT t.colo AS "t.colo"
+     FROM db.t
+)
+SELECT n1."t.colo" AS "n1.colo"
+FROM n1;
+`)
+}

sqlite/with_statement.go

@@ -49,7 +49,7 @@ func (c *commonTableExpression) AS(statement jet.SerializerHasProjections) Commo
 
 // AS_NOT_MATERIALIZED is used to define not materialized CTE query
 func (c *commonTableExpression) AS_NOT_MATERIALIZED(statement jet.SerializerHasProjections) CommonTableExpression {
-	c.CommonTableExpression.NotMaterialized = true
+	c.CommonTableExpression.Materialization = jet.CTEMaterializationDisabled
 	c.CommonTableExpression.Statement = statement
 	return c
 }

tests/postgres/with_test.go

@@ -875,3 +875,61 @@ WHERE orders1."orders.order_id" < $1;
 	require.NoError(t, err)
 	require.Len(t, dest, 72)
 }
+
+func TestWitStatement_CTE_Materialized(t *testing.T) {
+	orders1 := CTE("orders1")
+	orders1ID := Orders.OrderID.From(orders1)
+	orders2 := orders1.ALIAS("orders2")
+	orders2ID := Orders.OrderID.From(orders2)
+
+	stmt := WITH(
+		orders1.AS_MATERIALIZED(
+			SELECT(
+				Orders.OrderID,
+				Orders.EmployeeID,
+				Orders.ShipCity,
+			).FROM(
+				Orders,
+			),
+		),
+	)(
+		SELECT(
+			orders1.AllColumns().As("orders1.*"),
+			orders2.AllColumns().As("orders2.*"),
+		).FROM(
+			orders1.
+				INNER_JOIN(orders2, orders1ID.EQ(orders2ID)),
+		).WHERE(
+			orders1ID.LT(Int(10320)),
+		),
+	)
+
+	// fmt.Println(stmt.Sql())
+
+	testutils.AssertStatementSql(t, stmt, `
+WITH orders1 AS MATERIALIZED (
+     SELECT orders.order_id AS "orders.order_id",
+          orders.employee_id AS "orders.employee_id",
+          orders.ship_city AS "orders.ship_city"
+     FROM northwind.orders
+)
+SELECT orders1."orders.order_id" AS "orders1.order_id",
+     orders1."orders.employee_id" AS "orders1.employee_id",
+     orders1."orders.ship_city" AS "orders1.ship_city",
+     orders2."orders.order_id" AS "orders2.order_id",
+     orders2."orders.employee_id" AS "orders2.employee_id",
+     orders2."orders.ship_city" AS "orders2.ship_city"
+FROM orders1
+     INNER JOIN orders1 AS orders2 ON (orders1."orders.order_id" = orders2."orders.order_id")
+WHERE orders1."orders.order_id" < $1;
+`)
+
+	var dest []struct {
+		Orders1 model.Orders `alias:"orders1.*"`
+		Orders2 model.Orders `alias:"orders2.*"`
+	}
+
+	err := stmt.Query(db, &dest)
+	require.NoError(t, err)
+	require.Len(t, dest, 72)
+}