dialect/sql: added support for some yql select clauses

Author: LostImagin4tion

dialect/sql/builder.go

@@ -1702,15 +1702,15 @@ func (s *SelectTable) As(alias string) *SelectTable {
 }
 
 // View sets the secondary index name for the VIEW clause (YDB-specific).
-// This allows explicit use of secondary indexes in JOIN operations.
+// This allows explicit use of secondary indexes in SELECT and JOIN operations.
 //
 //	t := Table("users").View("idx_email").As("u")
 //	Select().From(Table("orders")).Join(t).On(...)
 func (s *SelectTable) View(index string) *SelectTable {
 	if s.ydb() {
 		s.index = index
 	} else {
-		s.AddError(fmt.Errorf("JOIN VIEW: unsupported dialect: %q", s.dialect))
+		s.AddError(fmt.Errorf("VIEW: unsupported dialect: %q", s.dialect))
 	}
 	return s
 }
@@ -1792,24 +1792,25 @@ type Selector struct {
 	Builder
 	// ctx stores contextual data typically from
 	// generated code such as alternate table schemas.
-	ctx       context.Context
-	as        string
-	selection []selection
-	from      []TableView
-	joins     []join
-	collected [][]*Predicate
-	where     *Predicate
-	or        bool
-	not       bool
-	order     []any
-	group     []string
-	having    *Predicate
-	limit     *int
-	offset    *int
-	distinct  bool
-	setOps    []setOp
-	prefix    Queries
-	lock      *LockOptions
+	ctx         context.Context
+	as          string
+	selection   []selection
+	from        []TableView
+	joins       []join
+	collected   [][]*Predicate
+	where       *Predicate
+	or          bool
+	not         bool
+	order       []any
+	assumeOrder []string // YDB-specific: ASSUME ORDER BY columns
+	group       []string
+	having      *Predicate
+	limit       *int
+	offset      *int
+	distinct    bool
+	setOps      []setOp
+	prefix      Queries
+	lock        *LockOptions
 }
 
 // New returns a new Selector with the same dialect and context.
@@ -2543,21 +2544,22 @@ func (s *Selector) Clone() *Selector {
 		joins[i] = s.joins[i].clone()
 	}
 	return &Selector{
-		Builder:   s.Builder.clone(),
-		ctx:       s.ctx,
-		as:        s.as,
-		or:        s.or,
-		not:       s.not,
-		from:      s.from,
-		limit:     s.limit,
-		offset:    s.offset,
-		distinct:  s.distinct,
-		where:     s.where.clone(),
-		having:    s.having.clone(),
-		joins:     append([]join{}, joins...),
-		group:     append([]string{}, s.group...),
-		order:     append([]any{}, s.order...),
-		selection: append([]selection{}, s.selection...),
+		Builder:     s.Builder.clone(),
+		ctx:         s.ctx,
+		as:          s.as,
+		or:          s.or,
+		not:         s.not,
+		from:        s.from,
+		limit:       s.limit,
+		offset:      s.offset,
+		distinct:    s.distinct,
+		where:       s.where.clone(),
+		having:      s.having.clone(),
+		joins:       append([]join{}, joins...),
+		group:       append([]string{}, s.group...),
+		order:       append([]any{}, s.order...),
+		assumeOrder: append([]string{}, s.assumeOrder...),
+		selection:   append([]selection{}, s.selection...),
 	}
 }
 
@@ -2585,6 +2587,11 @@ func DescExpr(x Querier) Querier {
 
 // OrderBy appends the `ORDER BY` clause to the `SELECT` statement.
 func (s *Selector) OrderBy(columns ...string) *Selector {
+	if s.ydb() && len(s.assumeOrder) != 0 {
+		s.AddError(fmt.Errorf("ORDER BY: can't be used with ASSUME ORDER BY simultaneously"))
+		return s
+	}
+
 	for i := range columns {
 		s.order = append(s.order, columns[i])
 	}
@@ -2626,6 +2633,27 @@ func (s *Selector) ClearOrder() *Selector {
 	return s
 }
 
+// AssumeOrderBy appends the `ASSUME ORDER BY` clause to the `SELECT` statement (YDB-specific).
+// This tells YDB to assume the data is already sorted without actually sorting it.
+// This is an optimization hint and only works with column names (not expressions).
+//
+//	Select("*").
+//		From(Table("users")).
+//		AssumeOrderBy("first-key", Desc("second-key"))
+func (s *Selector) AssumeOrderBy(columns ...string) *Selector {
+	if s.ydb() {
+		if len(s.order) != 0 {
+			s.AddError(fmt.Errorf("ASSUME ORDER BY: can't be used with ORDER BY simultaneously"))
+			return s
+		}
+
+		s.assumeOrder = append(s.assumeOrder, columns...)
+	} else {
+		s.AddError(fmt.Errorf("ASSUME ORDER BY: unsupported dialect: %q", s.dialect))
+	}
+	return s
+}
+
 // GroupBy appends the `GROUP BY` clause to the `SELECT` statement.
 func (s *Selector) GroupBy(columns ...string) *Selector {
 	s.group = append(s.group, columns...)
@@ -2716,6 +2744,7 @@ func (s *Selector) Query() (string, []any) {
 		s.joinSetOps(&b)
 	}
 	joinOrder(s.order, &b)
+	s.joinAssumeOrder(&b)
 	if s.limit != nil {
 		b.WriteString(" LIMIT ")
 		b.WriteString(strconv.Itoa(*s.limit))
@@ -2794,6 +2823,19 @@ func joinOrder(order []any, b *Builder) {
 	}
 }
 
+func (s *Selector) joinAssumeOrder(b *Builder) {
+	if !b.ydb() || len(s.assumeOrder) == 0 {
+		return
+	}
+	b.WriteString(" ASSUME ORDER BY ")
+	for i := range s.assumeOrder {
+		if i > 0 {
+			b.Comma()
+		}
+		b.Ident(s.assumeOrder[i])
+	}
+}
+
 func joinReturning(columns []string, b *Builder) {
 	supportedByDialect := b.postgres() || b.sqlite() || b.ydb()
 	if len(columns) == 0 || !supportedByDialect {

dialect/sql/builder_test.go

@@ -1846,6 +1846,53 @@ AND "users"."id1" < "users"."id2") AND "users"."id1" <= "users"."id2"`, "\n", ""
 			}(),
 			wantQuery: "SELECT `a`.`value`, `b`.`value`, `c`.`column2` FROM `a_table` AS `a` CROSS JOIN `b_table` AS `b` LEFT JOIN `c_table` AS `c` ON `c`.`ref` = `a`.`key` AND `c`.`column1` = `b`.`value`",
 		},
+		{
+			input: Dialect(dialect.YDB).
+				Select("value").
+				From(Table("my_table")).
+				Distinct(),
+			wantQuery: "SELECT DISTINCT `value` FROM `my_table`",
+		},
+		{
+			input: Dialect(dialect.YDB).
+				Select("*").
+				From(Table("users")).
+				Distinct(),
+			wantQuery: "SELECT DISTINCT * FROM `users`",
+		},
+		{
+			input: Dialect(dialect.YDB).
+				Select("status", "type").
+				From(Table("orders")).
+				Distinct().
+				Where(GT("created_at", "2024-01-01")),
+			wantQuery: "SELECT DISTINCT `status`, `type` FROM `orders` WHERE `created_at` > $p0",
+			wantArgs:  []any{driver.NamedValue{Name: "p0", Value: "2024-01-01"}},
+		},
+		{
+			input: Dialect(dialect.YDB).
+				Select("category").
+				From(Table("products")).
+				Distinct().
+				OrderBy("category").
+				Limit(10),
+			wantQuery: "SELECT DISTINCT `category` FROM `products` ORDER BY `category` LIMIT 10",
+		},
+		{
+			input: Dialect(dialect.YDB).
+				Select("name", "age").
+				From(Table("users")).
+				Distinct(),
+			wantQuery: "SELECT DISTINCT `name`, `age` FROM `users`",
+		},
+		{
+			input: Dialect(dialect.YDB).
+				Select("email").
+				From(Table("subscribers")).
+				Where(EQ("active", true)).
+				Distinct(),
+			wantQuery: "SELECT DISTINCT `email` FROM `subscribers` WHERE `active`",
+		},
 	}
 	for i, tt := range tests {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -1969,6 +2016,62 @@ func TestSelector_Union(t *testing.T) {
 		Query()
 	require.Equal(t, `SELECT * FROM "users" WHERE "active" UNION SELECT * FROM "old_users1" WHERE "is_active" AND "age" > $1 UNION ALL SELECT * FROM "old_users2" WHERE "is_active" = $2 AND "age" < $3`, query)
 	require.Equal(t, []any{20, "true", 18}, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("key").
+		From(Table("T1")).
+		Union(
+			Dialect(dialect.YDB).
+				Select("key").
+				From(Table("T2")),
+		).
+		Query()
+	require.Equal(t, "SELECT `key` FROM `T1` UNION SELECT `key` FROM `T2`", query)
+	require.Empty(t, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("id", "name").
+		From(Table("users")).
+		Where(EQ("status", "active")).
+		Union(
+			Dialect(dialect.YDB).
+				Select("id", "name").
+				From(Table("users")).
+				Where(EQ("status", "inactive")),
+		).
+		Query()
+	require.Equal(t, "SELECT `id`, `name` FROM `users` WHERE `status` = $p0 UNION SELECT `id`, `name` FROM `users` WHERE `status` = $p1", query)
+	require.Equal(t, []any{
+		driver.NamedValue{Name: "p0", Value: "active"},
+		driver.NamedValue{Name: "p1", Value: "inactive"},
+	}, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("x").
+		UnionAll(Dialect(dialect.YDB).Select("y")).
+		UnionAll(Dialect(dialect.YDB).Select("z")).
+		Query()
+	require.Equal(t, "SELECT `x` UNION ALL SELECT `y` UNION ALL SELECT `z`", query)
+	require.Empty(t, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("id").
+		From(Table("orders")).
+		Where(EQ("status", "pending")).
+		Union(
+			Dialect(dialect.YDB).
+				Select("id").
+				From(Table("orders")).
+				Where(EQ("status", "processing")),
+		).
+		OrderBy("id").
+		Limit(100).
+		Query()
+	require.Equal(t, "SELECT `id` FROM `orders` WHERE `status` = $p0 UNION SELECT `id` FROM `orders` WHERE `status` = $p1 ORDER BY `id` LIMIT 100", query)
+	require.Equal(t, []any{
+		driver.NamedValue{Name: "p0", Value: "pending"},
+		driver.NamedValue{Name: "p1", Value: "processing"},
+	}, args)
 }
 
 func TestSelector_Except(t *testing.T) {
@@ -2031,6 +2134,111 @@ func TestSelector_Intersect(t *testing.T) {
 	require.Equal(t, []any{20, "true", 18}, args)
 }
 
+func TestSelector_AssumeOrderBy_YDB(t *testing.T) {
+	query, args := Dialect(dialect.YDB).
+		Select("*").
+		From(Table("my_table")).
+		AssumeOrderBy("key").
+		Query()
+	require.Equal(t, "SELECT * FROM `my_table` ASSUME ORDER BY `key`", query)
+	require.Empty(t, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("key", "subkey").
+		From(Table("my_table")).
+		AssumeOrderBy("key", Desc("subkey")).
+		Query()
+	require.Equal(t, "SELECT `key`, `subkey` FROM `my_table` ASSUME ORDER BY `key`, `subkey` DESC", query)
+	require.Empty(t, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("*").
+		From(Table("orders")).
+		Where(EQ("status", "active")).
+		AssumeOrderBy("created_at").
+		Query()
+	require.Equal(t, "SELECT * FROM `orders` WHERE `status` = $p0 ASSUME ORDER BY `created_at`", query)
+	require.Equal(t, []any{driver.NamedValue{Name: "p0", Value: "active"}}, args)
+
+	query, args = Dialect(dialect.YDB).
+		Select("id", "name").
+		From(Table("users")).
+		AssumeOrderBy("id").
+		Limit(10).
+		Query()
+	require.Equal(t, "SELECT `id`, `name` FROM `users` ASSUME ORDER BY `id` LIMIT 10", query)
+	require.Empty(t, args)
+
+	t.Run("OrderBy then AssumeOrderBy should error", func(t *testing.T) {
+		selector := Dialect(dialect.YDB).
+			Select("*").
+			From(Table("users")).
+			OrderBy("id").
+			AssumeOrderBy("id")
+		
+		err := selector.Err()
+		require.Error(t, err)
+	})
+
+	t.Run("AssumeOrderBy then OrderBy should error", func(t *testing.T) {
+		selector := Dialect(dialect.YDB).
+			Select("*").
+			From(Table("users")).
+			AssumeOrderBy("id").
+			OrderBy("id")
+
+		err := selector.Err()
+		require.Error(t, err)
+	})
+
+	t.Run("Non-YDB dialect with AssumeOrderBy should error", func(t *testing.T) {
+		selector := Dialect(dialect.Postgres).
+			Select("*").
+			From(Table("users")).
+			AssumeOrderBy("name")
+		
+		err := selector.Err()
+		require.Error(t, err)
+	})
+}
+
+func TestSelector_VIEW_SecondaryIndex_YDB(t *testing.T) {
+	t.Run("Simple SELECT with VIEW in FROM clause", func(t *testing.T) {
+		d := Dialect(dialect.YDB)
+		query, args := d.Select("series_id", "title", "info", "release_date", "views", "uploaded_user_id").
+			From(d.Table("series").View("views_index")).
+			Where(GTE("views", 1000)).
+			Query()
+		
+		require.Equal(t, "SELECT `series_id`, `title`, `info`, `release_date`, `views`, `uploaded_user_id` FROM `series` VIEW `views_index` WHERE `views` >= $p0", query)
+		require.Equal(t, []any{driver.NamedValue{Name: "p0", Value: 1000}}, args)
+	})
+
+	t.Run("JOIN with VIEW on both tables", func(t *testing.T) {
+		d := Dialect(dialect.YDB)
+		series := d.Table("series").View("users_index").As("t1")
+		users := d.Table("users").View("name_index").As("t2")
+		
+		query, args := d.Select(series.C("series_id"), series.C("title")).
+			From(series).
+			Join(users).
+			On(series.C("uploaded_user_id"), users.C("user_id")).
+			Where(EQ(users.C("name"), "John Doe")).
+			Query()
+		
+		require.Equal(t, "SELECT `t1`.`series_id`, `t1`.`title` FROM `series` VIEW `users_index` AS `t1` JOIN `users` VIEW `name_index` AS `t2` ON `t1`.`uploaded_user_id` = `t2`.`user_id` WHERE `t2`.`name` = $p0", query)
+		require.Equal(t, []any{driver.NamedValue{Name: "p0", Value: "John Doe"}}, args)
+	})
+
+	t.Run("VIEW on non-YDB dialect should error", func(t *testing.T) {
+		d := Dialect(dialect.Postgres)
+		table := d.Table("users").View("idx_name")
+		
+		err := table.Err()
+		require.Error(t, err)
+	})
+}
+
 func TestSelector_SetOperatorWithRecursive(t *testing.T) {
 	t1, t2, t3 := Table("files"), Table("files"), Table("path")
 	n := Queries{