package edb.client; import com.carrotsearch.hppc.ByteScatterSet; import com.google.protobuf.ByteString; import edb.common.*; import edb.common.Split; import edb.rpc.*; import io.grpc.ManagedChannel; import io.grpc.ManagedChannelBuilder; import io.grpc.StatusRuntimeException; import java.util.ArrayList; import java.util.List; public class DBClient implements IExampleDB { private final ManagedChannel _channel; private final ExampleDBGrpc.ExampleDBBlockingStub _blockingStub; public DBClient(String host, int port) { this(ManagedChannelBuilder.forAddress(host, port) // Channels are secure by default (via SSL/TLS). For the example we disable TLS to avoid // needing certificates. .usePlaintext(true) .build()); } public DBClient(ManagedChannel channel) { _channel = channel; _blockingStub = ExampleDBGrpc.newBlockingStub(channel); } public String ping(String id) { PingRequest.Builder builder = PingRequest.newBuilder(); builder.setId(id); PingRequest request = builder.build(); PingResponse response; try { response = _blockingStub.ping(request); } catch (StatusRuntimeException e) { e.printStackTrace(); return null; } return response.getId(); } public List<String> listTables() { ListTablesRequest.Builder builder = ListTablesRequest.newBuilder(); ListTablesRequest request = builder.build(); ListTablesResponse response; try { response = _blockingStub.listTables(request); } catch (StatusRuntimeException e) { e.printStackTrace(); return null; } return response.getTableNamesList(); } public void createTable(String name, Schema schema) throws ExistingTableException { CreateTableRequest.Builder builder = CreateTableRequest.newBuilder(); builder.setName(name); TableSchema.Builder schemaBuilder = TableSchema.newBuilder(); schema.build(schemaBuilder); builder.setSchema(schemaBuilder.build()); CreateTableRequest request = builder.build(); CreateTableResponse response; try { response = _blockingStub.createTable(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (!response.getResult()) { throw new ExistingTableException(name); } } public void createTable(String name, Schema schema, String clusterColumn) throws ExistingTableException { CreateTableRequest.Builder builder = CreateTableRequest.newBuilder(); builder.setName(name); TableSchema.Builder schemaBuilder = TableSchema.newBuilder(); schema.build(schemaBuilder); builder.setSchema(schemaBuilder.build()); builder.setClusterColumn(clusterColumn); CreateTableRequest request = builder.build(); CreateTableResponse response; try { response = _blockingStub.createTable(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (!response.getResult()) { throw new ExistingTableException(name); } } public synchronized String createTemporaryTable(Schema schema) { CreateTemporaryTableRequest.Builder builder = CreateTemporaryTableRequest.newBuilder(); TableSchema.Builder schemaBuilder = TableSchema.newBuilder(); schema.build(schemaBuilder); builder.setSchema(schemaBuilder.build()); CreateTemporaryTableRequest request = builder.build(); CreateTemporaryTableResponse response; try { response = _blockingStub.createTemporaryTable(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } return response.getName(); } public Schema getTableSchema(String name) throws UnknownTableException { GetTableSchemaRequest.Builder builder = GetTableSchemaRequest.newBuilder(); builder.setName(name); GetTableSchemaRequest request = builder.build(); GetTableSchemaResponse response; try { response = _blockingStub.getTableSchema(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (response.getResult()) { Schema schema = new Schema(response.getSchema()); return schema; } else { throw new UnknownTableException(name); } } public String getTableClusteredIndexColumn(String tableName) throws UnknownTableException { GetTableClusteredIndexColumnRequest.Builder builder = GetTableClusteredIndexColumnRequest.newBuilder(); builder.setTableName(tableName); GetTableClusteredIndexColumnRequest request = builder.build(); GetTableClusteredIndexColumnResponse response; try { response = _blockingStub.getTableClusteredIndexColumn(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (response.getResult()) { if (response.hasColumnName()) { return response.getColumnName(); } else { return null; } } else { throw new UnknownTableException(tableName); } } public void bulkInsert(String name, List<Row> rows) throws UnknownTableException { BulkInsertRequest.Builder builder = BulkInsertRequest.newBuilder(); builder.setName(name); List<DataRow> rpcRows = new ArrayList<>(); for (Row row : rows) { DataRow.Builder rowBuilder = DataRow.newBuilder(); row.build(rowBuilder); rpcRows.add(rowBuilder.build()); } builder.addAllRow(rpcRows); BulkInsertRequest request = builder.build(); BulkInsertResponse response; try { response = _blockingStub.bulkInsert(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (!response.getResult()) { throw new UnknownTableException(name); } } public synchronized void bulkInsertFromTables(String destination, boolean truncateDestination, List<String> sourceTables) throws UnknownTableException { BulkInsertFromTablesRequest.Builder builder = BulkInsertFromTablesRequest.newBuilder(); builder.setDestination(destination); builder.setTruncateDestination(truncateDestination); builder.addAllSource(sourceTables); BulkInsertFromTablesRequest request = builder.build(); BulkInsertFromTablesResponse response; try { response = _blockingStub.bulkInsertFromTables(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (response.hasUnknown()) { throw new UnknownTableException(response.getUnknown()); } } public List<Row> getAllRows(String name) throws UnknownTableException { return getAllRows(name, null, null); } public List<Row> getAllRows(String name, List<String> columns) throws UnknownTableException { return getAllRows(name, null, columns); } public List<Row> getAllRows(String name, Split split) throws UnknownTableException { return getAllRows(name, split, null); } public List<Row> getAllRows(String name, Split split, List<String> columns) throws UnknownTableException { GetAllRowsRequest.Builder builder = GetAllRowsRequest.newBuilder(); builder.setName(name); if (split != null) { edb.rpc.Split.Builder splitBuilder = edb.rpc.Split.newBuilder(); splitBuilder.setOpaque(ByteString.copyFrom(split.serialize())); builder.setSplit(splitBuilder.build()); } if (columns != null) { builder.addAllColumns(columns); } GetAllRowsRequest request = builder.build(); GetAllRowsResponse response; try { response = _blockingStub.getAllRows(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (response.getResult()) { List<Row> rows = new ArrayList<>(); List<DataRow> rpcRows = response.getRowList(); for (DataRow rpcRow : rpcRows) { rows.add(new Row(rpcRow)); } return rows; } else { throw new UnknownTableException(name); } } public List<Split> getSplits(String table) throws UnknownTableException { return getSplits(table, 0); } public List<Split> getSplits(String table, int count) throws UnknownTableException { GetSplitsRequest.Builder builder = GetSplitsRequest.newBuilder(); builder.setName(table); if (count != 0) { builder.setCount(count); } GetSplitsRequest request = builder.build(); GetSplitsResponse response; try { response = _blockingStub.getSplits(request); } catch (StatusRuntimeException e) { e.printStackTrace(); throw e; } if (response.getResult()) { List<Split> splits = new ArrayList<>(); List<edb.rpc.Split> rpcSplits = response.getSplitsList(); for (edb.rpc.Split rpcSplit : rpcSplits) { byte[] splitBytes = rpcSplit.getOpaque().toByteArray(); Split split = Split.deserialize(splitBytes); splits.add(split); } return splits; } else { throw new UnknownTableException(table); } } }