From 2be57fbff61919244f1c9e37948b915beea96f4c Mon Sep 17 00:00:00 2001
From: zhiheng123 <903292776@qq.com>
Date: Thu, 29 Aug 2024 17:09:49 +0800
Subject: [PATCH] feat: add as expression in query builder

Signed-off-by: zhiheng123 <903292776@qq.com>
---
 opengemini/query_builder_test.go | 55 ++++++++++++++++++++++++++++++++
 opengemini/query_expression.go   | 16 ++++++++++
 2 files changed, 71 insertions(+)

diff --git a/opengemini/query_builder_test.go b/opengemini/query_builder_test.go
index 6ed3eed..28f2132 100644
--- a/opengemini/query_builder_test.go
+++ b/opengemini/query_builder_test.go
@@ -203,3 +203,58 @@ func TestQueryBuilderSelectWithWhereAndTimezone(t *testing.T) {
 
 	require.Equal(t, expectedQuery, query.Command)
 }
+
+func TestQueryBuilderSelectWithAsExpression(t *testing.T) {
+	qb := CreateQueryBuilder()
+
+	waterLevelField := NewFieldExpression("water_level")
+
+	locationCondition := NewComparisonCondition("location", Equals, "santa_monica")
+	startTimeCondition := NewComparisonCondition("time", GreaterThanOrEquals, "2019-08-18T00:00:00Z")
+	endTimeCondition := NewComparisonCondition("time", LessThanOrEquals, "2019-08-18T00:18:00Z")
+
+	finalCondition := NewCompositeCondition(And, locationCondition, startTimeCondition, endTimeCondition)
+
+	location, err := time.LoadLocation("America/Chicago")
+	require.NoError(t, err)
+
+        asWL := NewAsExpression("WL", waterLevelField)
+
+	query := qb.Select(asWL).
+		From("h2o_feet").
+		Where(finalCondition).
+		Timezone(location).
+		Build()
+
+	expectedQuery := `SELECT "water_level" AS "WL" FROM "h2o_feet" WHERE ("location" = 'santa_monica' AND "time" >= '2019-08-18T00:00:00Z' AND "time" <= '2019-08-18T00:18:00Z') TZ('America/Chicago')`
+
+	require.Equal(t, expectedQuery, query.Command)
+}
+
+func TestQueryBuilderSelectWithAggregate(t *testing.T) {
+	qb := CreateQueryBuilder()
+
+	waterLevelField := NewFieldExpression("water_level")
+	countWaterLevelField := NewFunctionExpression(FunctionCount, waterLevelField)
+
+	locationCondition := NewComparisonCondition("location", Equals, "santa_monica")
+	startTimeCondition := NewComparisonCondition("time", GreaterThanOrEquals, "2019-08-18T00:00:00Z")
+	endTimeCondition := NewComparisonCondition("time", LessThanOrEquals, "2019-08-18T00:18:00Z")
+
+	finalCondition := NewCompositeCondition(And, locationCondition, startTimeCondition, endTimeCondition)
+
+	location, err := time.LoadLocation("America/Chicago")
+	require.NoError(t, err)
+
+        asWL := NewAsExpression("WL", countWaterLevelField)
+
+	query := qb.Select(asWL).
+		From("h2o_feet").
+		Where(finalCondition).
+		Timezone(location).
+		Build()
+
+	expectedQuery := `SELECT COUNT("water_level") AS "WL" FROM "h2o_feet" WHERE ("location" = 'santa_monica' AND "time" >= '2019-08-18T00:00:00Z' AND "time" <= '2019-08-18T00:18:00Z') TZ('America/Chicago')`
+
+	require.Equal(t, expectedQuery, query.Command)
+}
diff --git a/opengemini/query_expression.go b/opengemini/query_expression.go
index f847293..5fa34c3 100644
--- a/opengemini/query_expression.go
+++ b/opengemini/query_expression.go
@@ -61,6 +61,22 @@ func NewFunctionExpression(function FunctionEnum, arguments ...Expression) *Func
 	}
 }
 
+type AsExpression struct {
+    Alias string
+    OriginExpr Expression
+}
+
+func (a *AsExpression) build() string {
+    return fmt.Sprintf("%s AS \"%s\"", a.OriginExpr.build(), a.Alias)
+}
+
+func NewAsExpression(alias string, expr Expression) *AsExpression {
+    return &AsExpression{
+        Alias : alias,
+        OriginExpr: expr,
+    }
+}
+
 type ArithmeticExpression struct {
 	Operator ArithmeticOperator
 	Operands []Expression