topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-11 03:58:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

(feature)(Headless) arrow flight sql endpoint (#634) (#1136)

Browse Source
This commit is contained in:
jipeli
2024-06-12 15:02:55 +08:00
committed by GitHub
parent 91a6308d9e
commit abf440c4d4
19 changed files with 697 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
launchers/standalone/src/test/java/com/tencent/supersonic/headless/FlightSqlTest.java Normal file
View File

@@ -0,0 +1,101 @@
package com.tencent.supersonic.headless;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import com.tencent.supersonic.auth.api.authentication.utils.UserHolder;
import com.tencent.supersonic.auth.authentication.strategy.FakeUserStrategy;
import com.tencent.supersonic.headless.server.listener.FlightSqlListener;
import org.apache.arrow.flight.CallHeaders;
import org.apache.arrow.flight.FlightCallHeaders;
import org.apache.arrow.flight.FlightClient;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.HeaderCallOption;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.sql.FlightSqlClient;
import org.apache.arrow.memory.RootAllocator;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
public class FlightSqlTest extends BaseTest {
@Autowired
private FlightSqlListener flightSqlListener;
@Autowired
private FakeUserStrategy fakeUserStrategy;
@Test
void test01() throws Exception {
String host = flightSqlListener.getHost();
Integer port = flightSqlListener.getPort();
UserHolder.setStrategy(fakeUserStrategy);
flightSqlListener.startServer();
FlightSqlClient sqlClient = new FlightSqlClient(
FlightClient.builder(new RootAllocator(Integer.MAX_VALUE), Location.forGrpcInsecure(host, port))
.build());
CallHeaders headers = new FlightCallHeaders();
headers.insert("dataSetId", "1");
headers.insert("name", "admin");
headers.insert("password", "admin");
HeaderCallOption headerOption = new HeaderCallOption(headers);
try (final FlightSqlClient.PreparedStatement preparedStatement = sqlClient.prepare(
"SELECT 部门, SUM(访问次数) AS 访问次数 FROM 超音数PVUV统计 GROUP BY 部门",
headerOption)) {
final FlightInfo info = preparedStatement.execute();
FlightStream stream = sqlClient.getStream(info
.getEndpoints()
.get(0).getTicket());
int rowCnt = 0;
int colCnt = 0;
while (stream.next()) {
if (stream.getRoot().getRowCount() > 0) {
colCnt = stream.getRoot().getFieldVectors().size();
rowCnt += stream.getRoot().getRowCount();
}
}
assertEquals(2, colCnt);
assertTrue(rowCnt > 0);
} catch (Exception e) {
e.printStackTrace();
}
}
@Test
void test02() throws Exception {
String host = flightSqlListener.getHost();
Integer port = flightSqlListener.getPort();
UserHolder.setStrategy(fakeUserStrategy);
flightSqlListener.startServer();
FlightSqlClient sqlClient = new FlightSqlClient(
FlightClient.builder(new RootAllocator(Integer.MAX_VALUE), Location.forGrpcInsecure(host, port))
.build());
CallHeaders headers = new FlightCallHeaders();
headers.insert("dataSetId", "1");
headers.insert("name", "admin");
headers.insert("password", "admin");
HeaderCallOption headerOption = new HeaderCallOption(headers);
try {
FlightInfo flightInfo = sqlClient.execute("SELECT 部门, SUM(访问次数) AS 访问次数 FROM 超音数PVUV统计 GROUP BY 部门",
headerOption);
FlightStream stream = sqlClient.getStream(flightInfo
.getEndpoints()
.get(0).getTicket());
int rowCnt = 0;
int colCnt = 0;
while (stream.next()) {
if (stream.getRoot().getRowCount() > 0) {
colCnt = stream.getRoot().getFieldVectors().size();
rowCnt += stream.getRoot().getRowCount();
}
}
assertEquals(2, colCnt);
assertTrue(rowCnt > 0);
} catch (Exception e) {
e.printStackTrace();
}
}
}

2
launchers/standalone/src/test/resources/db/data-h2.sql
View File

@@ -1,5 +1,5 @@
-- sample user
MERGE INTO s2_user (id, `name`, password, display_name, email, is_admin) values (1, 'admin','admin','admin','admin@xx.com', 1);
MERGE INTO s2_user (id, `name`, password, salt, display_name, email, is_admin) values (1, 'admin','c3VwZXJzb25pY0BiaWNvbTD12g9wGXESwL7+o7xUW90=','jGl25bVBBBW96Qi9Te4V3w==','admin','admin@xx.com', 1);
MERGE INTO s2_user (id, `name`, password, display_name, email) values (2, 'jack','123456','jack','jack@xx.com');
MERGE INTO s2_user (id, `name`, password, display_name, email) values (3, 'tom','123456','tom','tom@xx.com');
MERGE INTO s2_user (id, `name`, password, display_name, email, is_admin) values (4, 'lucy','123456','lucy','lucy@xx.com', 1);

1
launchers/standalone/src/test/resources/db/schema-h2.sql
View File

@@ -91,6 +91,7 @@ create table IF NOT EXISTS s2_user
password varchar(100) null,
email varchar(100) null,
is_admin INT null,
salt varchar(100) null,
PRIMARY KEY (`id`)
);
COMMENT ON TABLE s2_user IS 'user information table';