From 0c37c95c22754fc75884289f2b12d9d18af145d3 Mon Sep 17 00:00:00 2001
From: Frank Chen <frank.chen021@outlook.com>
Date: Tue, 13 Sep 2022 18:28:36 +0800
Subject: [PATCH] Add test cases

Signed-off-by: Frank Chen <frank.chen021@outlook.com>
---
 ...simple_queries_for_opentelemetry.reference |  4 +
 .../02421_simple_queries_for_opentelemetry.sh | 83 +++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.reference
 create mode 100755 tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.sh

diff --git a/tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.reference b/tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.reference
new file mode 100644
index 00000000000..6b614bd6d82
--- /dev/null
+++ b/tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.reference
@@ -0,0 +1,4 @@
+{"query":"show processlist format Null\n  "}
+{"query":"show databases format Null\n  "}
+{"query":"insert into opentelemetry_test values","read_rows":"3","read_bytes":"24","written_rows":"3","written_bytes":"24"}
+{"query":"select * from opentelemetry_test format Null\n  ","read_rows":"3","read_bytes":"24","written_rows":"","written_bytes":""}
diff --git a/tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.sh b/tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.sh
new file mode 100755
index 00000000000..c6f8bddf571
--- /dev/null
+++ b/tests/queries/0_stateless/02421_simple_queries_for_opentelemetry.sh
@@ -0,0 +1,83 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../shell_config.sh
+
+# This function takes 2 arguments:
+# $1 - query
+# $2 - trace id
+function executeQuery()
+{
+  ${CLICKHOUSE_CLIENT} --database=${CLICKHOUSE_DATABASE} --opentelemetry_start_trace_probability=1 --query_id $1 -nq "
+      ${2}
+  "
+}
+
+# For some quries, it's not able to know how many bytes/rows are read when tests are executed on CI,
+# so we only to check the db.statement only
+function check_query_span_query_only()
+{
+${CLICKHOUSE_CLIENT} -nq "
+    SYSTEM FLUSH LOGS;
+    SELECT attribute['db.statement']       as query
+    FROM system.opentelemetry_span_log
+    WHERE finish_date                      >= yesterday()
+    AND   operation_name                   = 'query'
+    AND   attribute['clickhouse.query_id'] = '${1}'
+    Format JSONEachRow
+    ;"
+}
+
+function check_query_span()
+{
+${CLICKHOUSE_CLIENT} -nq "
+    SYSTEM FLUSH LOGS;
+    SELECT attribute['db.statement']             as query,
+           attribute['clickhouse.read_rows']     as read_rows,
+           attribute['clickhouse.read_bytes']    as read_bytes,
+           attribute['clickhouse.written_rows']  as written_rows,
+           attribute['clickhouse.written_bytes'] as written_bytes
+    FROM system.opentelemetry_span_log
+    WHERE finish_date                      >= yesterday()
+    AND   operation_name                   = 'query'
+    AND   attribute['clickhouse.query_id'] = '${1}'
+    Format JSONEachRow
+    ;"
+}
+
+#
+# Set up
+#
+${CLICKHOUSE_CLIENT} --database=${CLICKHOUSE_DATABASE} -nq "
+DROP   TABLE IF EXISTS opentelemetry_test;
+CREATE TABLE opentelemetry_test (id UInt64) Engine=MergeTree Order By id;
+"
+
+# test 1, a query that has special path in the code
+# Format Null is used to make sure no output is generated so that it won't pollute the reference file
+query_id=$(${CLICKHOUSE_CLIENT} -q "select generateUUIDv4()");
+executeQuery $query_id 'show processlist format Null'
+check_query_span_query_only "$query_id"
+
+# test 2, a normal show command
+query_id=$(${CLICKHOUSE_CLIENT} -q "select generateUUIDv4()");
+executeQuery $query_id 'show databases format Null'
+check_query_span_query_only "$query_id"
+
+# test 3, a normal insert query on local table
+query_id=$(${CLICKHOUSE_CLIENT} -q "select generateUUIDv4()");
+executeQuery $query_id 'insert into opentelemetry_test values(1)(2)(3)'
+check_query_span "$query_id"
+
+# test 4, a normal select query
+query_id=$(${CLICKHOUSE_CLIENT} -q "select generateUUIDv4()");
+executeQuery $query_id 'select * from opentelemetry_test format Null'
+check_query_span $query_id
+
+
+#
+# Tear down
+#
+${CLICKHOUSE_CLIENT} --database=${CLICKHOUSE_DATABASE} -q "
+DROP   TABLE IF EXISTS opentelemetry_test;
+"
\ No newline at end of file