From 962839ed89c03e3e893ac3ffe9820c9da7861da4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ask=20Bj=C3=B8rn=20Hansen?= <ask@develooper.com>
Date: Sun, 23 Feb 2025 09:28:30 -0800
Subject: [PATCH] add dns query count endpoint

---
 chdb/geoqueries.go | 54 ++++++++++++++++++++++++++++++++++++++++++++++
 server/server.go   | 17 ++++++++++++++-
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/chdb/geoqueries.go b/chdb/geoqueries.go
index e324e6a..d9dab76 100644
--- a/chdb/geoqueries.go
+++ b/chdb/geoqueries.go
@@ -24,9 +24,11 @@ type UserCountry []flatAPI
 func (s UserCountry) Len() int {
 	return len(s)
 }
+
 func (s UserCountry) Swap(i, j int) {
 	s[i], s[j] = s[j], s[i]
 }
+
 func (s UserCountry) Less(i, j int) bool {
 	return s[i].IPv4 > s[j].IPv4
 }
@@ -183,3 +185,55 @@ func (d *ClickHouse) UserCountryData(ctx context.Context) (*UserCountry, error)
 
 	return nil, nil
 }
+
+type DNSQueryCounts struct {
+	T   uint32  `json:"t"`
+	Avg float64 `json:"avg"`
+	Max uint64  `json:"max"`
+}
+
+func (d *ClickHouse) DNSQueries(ctx context.Context) ([]DNSQueryCounts, error) {
+	log := logger.Setup()
+	ctx, span := tracing.Tracer().Start(ctx, "DNSQueries")
+	defer span.End()
+
+	startUnix := time.Now().Add(-90 * time.Minute).Unix()
+	startUnix -= startUnix % (60 * 5)
+
+	log.InfoContext(ctx, "start time", "start", startUnix)
+
+	rows, err := d.Logs.Query(clickhouse.Context(ctx, clickhouse.WithSpan(span.SpanContext())),
+		`
+	select toUnixTimestamp(toStartOfFiveMinute(t)) as t,
+  sum(q)/300 as avg, max(q) as max
+from (
+ select window as t, sumSimpleState(queries) as q
+ from geodns.by_origin_1s
+ where
+   window > FROM_UNIXTIME(?)
+   and Origin IN ('pool.ntp.org', 'g.ntpns.org')
+ group by t order by t
+)
+group by t order by t
+`, startUnix)
+	if err != nil {
+		log.ErrorContext(ctx, "query error", "err", err)
+		return nil, fmt.Errorf("database error")
+	}
+
+	var t uint32
+	var avg float64
+	var max uint64
+
+	r := []DNSQueryCounts{}
+
+	for rows.Next() {
+		if err := rows.Scan(&t, &avg, &max); err != nil {
+			return nil, err
+		}
+		log.InfoContext(ctx, "data", "t", t, "avg", avg, "max", max)
+		r = append(r, DNSQueryCounts{t, avg, max})
+	}
+
+	return r, nil
+}
diff --git a/server/server.go b/server/server.go
index e5a18bd..19d4797 100644
--- a/server/server.go
+++ b/server/server.go
@@ -207,6 +207,7 @@ func (srv *Server) Run() error {
 	e.GET("/api/usercc", srv.userCountryData)
 	e.GET("/api/server/dns/answers/:server", srv.dnsAnswers)
 	e.GET("/api/server/scores/:server/:mode", srv.history)
+	e.GET("/api/dns/counts", srv.dnsQueryCounts)
 
 	if len(ntpconf.WebHostname()) > 0 {
 		e.POST("/api/server/scores/:server/:mode", func(c echo.Context) error {
@@ -261,7 +262,7 @@ func (srv *Server) userCountryData(c echo.Context) error {
 		log.InfoContext(ctx, "didn't get zoneStats")
 	}
 
-	data, err := srv.ch.UserCountryData(c.Request().Context())
+	data, err := srv.ch.UserCountryData(ctx)
 	if err != nil {
 		log.ErrorContext(ctx, "UserCountryData", "err", err)
 		return c.String(http.StatusInternalServerError, err.Error())
@@ -276,6 +277,20 @@ func (srv *Server) userCountryData(c echo.Context) error {
 	})
 }
 
+func (srv *Server) dnsQueryCounts(c echo.Context) error {
+	log := logger.Setup()
+	ctx, span := tracing.Tracer().Start(c.Request().Context(), "dnsQueryCounts")
+	defer span.End()
+
+	data, err := srv.ch.DNSQueries(ctx)
+	if err != nil {
+		log.ErrorContext(ctx, "dnsQueryCounts", "err", err)
+		return c.String(http.StatusInternalServerError, err.Error())
+	}
+
+	return c.JSON(http.StatusOK, data)
+}
+
 func healthHandler(srv *Server, log *slog.Logger) func(w http.ResponseWriter, req *http.Request) {
 	return func(w http.ResponseWriter, req *http.Request) {
 		ctx := req.Context()