org.raystack.depot.bigtable.client.BigTableClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of depot Show documentation
A sink connector library containing multiple sink implementations
The newest version!
package org.raystack.depot.bigtable.client;

import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient;
import com.google.cloud.bigtable.admin.v2.BigtableTableAdminSettings;
import com.google.cloud.bigtable.admin.v2.models.ColumnFamily;
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
import com.google.cloud.bigtable.data.v2.models.MutateRowsException;
import org.raystack.depot.bigtable.exception.BigTableInvalidSchemaException;
import org.raystack.depot.bigtable.model.BigTableRecord;
import org.raystack.depot.bigtable.model.BigTableSchema;
import org.raystack.depot.bigtable.response.BigTableResponse;
import org.raystack.depot.config.BigTableSinkConfig;
import org.raystack.depot.metrics.BigTableMetrics;
import org.raystack.depot.metrics.Instrumentation;

import java.io.FileInputStream;
import java.io.IOException;
import java.time.Instant;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

public class BigTableClient {
    private final BigtableTableAdminClient bigtableTableAdminClient;
    private final BigtableDataClient bigtableDataClient;
    private final BigTableSinkConfig sinkConfig;
    private final BigTableSchema bigtableSchema;
    private final BigTableMetrics bigtableMetrics;
    private final Instrumentation instrumentation;

    public BigTableClient(BigTableSinkConfig sinkConfig, BigTableSchema bigtableSchema, BigTableMetrics bigtableMetrics,
            Instrumentation instrumentation) throws IOException {
        this(sinkConfig, getBigTableDataClient(sinkConfig), getBigTableAdminClient(sinkConfig), bigtableSchema,
                bigtableMetrics, instrumentation);
    }

    public BigTableClient(BigTableSinkConfig sinkConfig, BigtableDataClient bigtableDataClient,
            BigtableTableAdminClient bigtableTableAdminClient, BigTableSchema bigtableSchema,
            BigTableMetrics bigtableMetrics, Instrumentation instrumentation) {
        this.sinkConfig = sinkConfig;
        this.bigtableDataClient = bigtableDataClient;
        this.bigtableTableAdminClient = bigtableTableAdminClient;
        this.bigtableSchema = bigtableSchema;
        this.bigtableMetrics = bigtableMetrics;
        this.instrumentation = instrumentation;
    }

    private static BigtableDataClient getBigTableDataClient(BigTableSinkConfig sinkConfig) throws IOException {
        BigtableDataSettings settings = BigtableDataSettings.newBuilder()
                .setProjectId(sinkConfig.getGCloudProjectID())
                .setInstanceId(sinkConfig.getInstanceId())
                .setCredentialsProvider(FixedCredentialsProvider
                        .create(GoogleCredentials.fromStream(new FileInputStream(sinkConfig.getCredentialPath()))))
                .build();
        return BigtableDataClient.create(settings);
    }

    private static BigtableTableAdminClient getBigTableAdminClient(BigTableSinkConfig sinkConfig) throws IOException {
        BigtableTableAdminSettings settings = BigtableTableAdminSettings.newBuilder()
                .setProjectId(sinkConfig.getGCloudProjectID())
                .setInstanceId(sinkConfig.getInstanceId())
                .setCredentialsProvider(FixedCredentialsProvider
                        .create(GoogleCredentials.fromStream(new FileInputStream(sinkConfig.getCredentialPath()))))
                .build();
        return BigtableTableAdminClient.create(settings);
    }

    public BigTableResponse send(List records) {
        BigTableResponse bigTableResponse = null;
        BulkMutation batch = BulkMutation.create(sinkConfig.getTableId());
        records.forEach(record -> batch.add(record.getRowMutationEntry()));
        try {
            Instant startTime = Instant.now();
            bigtableDataClient.bulkMutateRows(batch);
            instrument(startTime, batch.getEntryCount());
        } catch (MutateRowsException e) {
            bigTableResponse = new BigTableResponse(e);
            instrumentation.logError("Some entries failed to be applied. {}", e.getCause());
        }
        return bigTableResponse;
    }

    private void instrument(Instant startTime, long entryCount) {
        instrumentation.captureDurationSince(
                bigtableMetrics.getBigtableOperationLatencyMetric(),
                startTime,
                String.format(BigTableMetrics.BIGTABLE_INSTANCE_TAG, sinkConfig.getInstanceId()),
                String.format(BigTableMetrics.BIGTABLE_TABLE_TAG, sinkConfig.getTableId()));
        instrumentation.captureCount(
                bigtableMetrics.getBigtableOperationTotalMetric(),
                entryCount,
                String.format(BigTableMetrics.BIGTABLE_INSTANCE_TAG, sinkConfig.getInstanceId()),
                String.format(BigTableMetrics.BIGTABLE_TABLE_TAG, sinkConfig.getTableId()));
    }

    public void validateBigTableSchema() throws BigTableInvalidSchemaException {
        String tableId = sinkConfig.getTableId();
        instrumentation.logDebug(String.format("Validating schema for table: %s...", tableId));
        checkIfTableExists(tableId);
        checkIfColumnFamiliesExist(tableId);
        instrumentation.logDebug("Validation complete, Schema is valid.");
    }

    private void checkIfTableExists(String tableId) throws BigTableInvalidSchemaException {
        if (!bigtableTableAdminClient.exists(tableId)) {
            throw new BigTableInvalidSchemaException(String.format(
                    "Table not found on the path: projects/%s/instances/%s/tables/%s",
                    bigtableTableAdminClient.getProjectId(), bigtableTableAdminClient.getInstanceId(), tableId));
        }
    }

    private void checkIfColumnFamiliesExist(String tableId) throws BigTableInvalidSchemaException {
        Set existingColumnFamilies = bigtableTableAdminClient.getTable(tableId)
                .getColumnFamilies()
                .stream()
                .map(ColumnFamily::getId)
                .collect(Collectors.toSet());
        Set missingColumnFamilies = bigtableSchema.getMissingColumnFamilies(existingColumnFamilies);
        if (missingColumnFamilies.size() > 0) {
            throw new BigTableInvalidSchemaException(
                    String.format("Column families %s do not exist in table %s!", missingColumnFamilies, tableId));
        }
    }
}