io.prestosql.plugin.accumulo.AccumuloMetadata Maven / Gradle / Ivy

Go to download
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.prestosql.plugin.accumulo;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import io.airlift.json.JsonCodec;
import io.airlift.json.JsonCodecFactory;
import io.airlift.json.ObjectMapperProvider;
import io.airlift.slice.Slice;
import io.prestosql.plugin.accumulo.metadata.AccumuloTable;
import io.prestosql.plugin.accumulo.model.AccumuloColumnHandle;
import io.prestosql.plugin.accumulo.model.AccumuloTableHandle;
import io.prestosql.spi.PrestoException;
import io.prestosql.spi.connector.ColumnHandle;
import io.prestosql.spi.connector.ColumnMetadata;
import io.prestosql.spi.connector.ConnectorInsertTableHandle;
import io.prestosql.spi.connector.ConnectorMetadata;
import io.prestosql.spi.connector.ConnectorNewTableLayout;
import io.prestosql.spi.connector.ConnectorOutputMetadata;
import io.prestosql.spi.connector.ConnectorOutputTableHandle;
import io.prestosql.spi.connector.ConnectorSession;
import io.prestosql.spi.connector.ConnectorTableHandle;
import io.prestosql.spi.connector.ConnectorTableMetadata;
import io.prestosql.spi.connector.ConnectorTableProperties;
import io.prestosql.spi.connector.ConnectorViewDefinition;
import io.prestosql.spi.connector.Constraint;
import io.prestosql.spi.connector.ConstraintApplicationResult;
import io.prestosql.spi.connector.SchemaTableName;
import io.prestosql.spi.connector.SchemaTablePrefix;
import io.prestosql.spi.connector.TableNotFoundException;
import io.prestosql.spi.predicate.TupleDomain;
import io.prestosql.spi.statistics.ComputedStatistics;

import javax.inject.Inject;

import java.util.Collection;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

import static com.google.common.base.Preconditions.checkState;
import static io.prestosql.plugin.accumulo.AccumuloErrorCode.ACCUMULO_TABLE_EXISTS;
import static io.prestosql.spi.StandardErrorCode.NOT_SUPPORTED;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;

/**
 * Presto metadata provider for Accumulo.
 * Responsible for creating/dropping/listing tables, schemas, columns, all sorts of goodness. Heavily leverages {@link AccumuloClient}.
 */
public class AccumuloMetadata
        implements ConnectorMetadata
{
    private static final JsonCodec VIEW_CODEC =
            new JsonCodecFactory(new ObjectMapperProvider()).jsonCodec(ConnectorViewDefinition.class);

    private final AccumuloClient client;
    private final AtomicReference rollbackAction = new AtomicReference<>();

    @Inject
    public AccumuloMetadata(AccumuloClient client)
    {
        this.client = requireNonNull(client, "client is null");
    }

    @Override
    public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Optional layout)
    {
        checkNoRollback();

        SchemaTableName tableName = tableMetadata.getTable();
        AccumuloTable table = client.createTable(tableMetadata);

        AccumuloTableHandle handle = new AccumuloTableHandle(
                tableName.getSchemaName(),
                tableName.getTableName(),
                table.getRowId(),
                table.isExternal(),
                table.getSerializerClassName(),
                table.getScanAuthorizations());

        setRollback(() -> rollbackCreateTable(table));

        return handle;
    }

    @Override
    public Optional finishCreateTable(ConnectorSession session, ConnectorOutputTableHandle tableHandle, Collection fragments, Collection computedStatistics)
    {
        clearRollback();
        return Optional.empty();
    }

    private void rollbackCreateTable(AccumuloTable table)
    {
        client.dropTable(table);
    }

    @Override
    public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean ignoreExisting)
    {
        client.createTable(tableMetadata);
    }

    @Override
    public void dropTable(ConnectorSession session, ConnectorTableHandle tableHandle)
    {
        AccumuloTableHandle handle = (AccumuloTableHandle) tableHandle;
        AccumuloTable table = client.getTable(handle.toSchemaTableName());
        if (table != null) {
            client.dropTable(table);
        }
    }

    @Override
    public void renameTable(ConnectorSession session, ConnectorTableHandle tableHandle,
            SchemaTableName newTableName)
    {
        if (client.getTable(newTableName) != null) {
            throw new PrestoException(ACCUMULO_TABLE_EXISTS, "Table " + newTableName + " already exists");
        }

        AccumuloTableHandle handle = (AccumuloTableHandle) tableHandle;
        client.renameTable(handle.toSchemaTableName(), newTableName);
    }

    @Override
    public void createView(ConnectorSession session, SchemaTableName viewName, ConnectorViewDefinition definition, boolean replace)
    {
        String viewData = VIEW_CODEC.toJson(definition);
        if (replace) {
            client.createOrReplaceView(viewName, viewData);
        }
        else {
            client.createView(viewName, viewData);
        }
    }

    @Override
    public void dropView(ConnectorSession session, SchemaTableName viewName)
    {
        client.dropView(viewName);
    }

    @Override
    public Optional getView(ConnectorSession session, SchemaTableName viewName)
    {
        return Optional.ofNullable(client.getView(viewName))
                .map(view -> VIEW_CODEC.fromJson(view.getData()));
    }

    @Override
    public List listViews(ConnectorSession session, Optional schemaName)
    {
        return listViews(schemaName);
    }

    /**
     * Gets all views in the given schema, or all schemas if null.
     *
     * @param filterSchema Schema to filter the views, or absent to list all schemas
     * @return List of views
     */
    private List listViews(Optional filterSchema)
    {
        ImmutableList.Builder builder = ImmutableList.builder();
        if (filterSchema.isPresent()) {
            for (String view : client.getViewNames(filterSchema.get())) {
                builder.add(new SchemaTableName(filterSchema.get(), view));
            }
        }
        else {
            for (String schemaName : client.getSchemaNames()) {
                for (String view : client.getViewNames(schemaName)) {
                    builder.add(new SchemaTableName(schemaName, view));
                }
            }
        }

        return builder.build();
    }

    @Override
    public ConnectorInsertTableHandle beginInsert(ConnectorSession session, ConnectorTableHandle tableHandle)
    {
        checkNoRollback();
        AccumuloTableHandle handle = (AccumuloTableHandle) tableHandle;
        setRollback(() -> rollbackInsert(handle));
        return handle;
    }

    @Override
    public Optional finishInsert(ConnectorSession session, ConnectorInsertTableHandle insertHandle, Collection fragments, Collection computedStatistics)
    {
        clearRollback();
        return Optional.empty();
    }

    private static void rollbackInsert(ConnectorInsertTableHandle insertHandle)
    {
        // Rollbacks for inserts are off the table when it comes to data in Accumulo.
        // When a batch of Mutations fails to be inserted, the general strategy
        // is to run the insert operation again until it is successful
        // Any mutations that were successfully written will be overwritten
        // with the same values, so that isn't a problem.
        AccumuloTableHandle handle = (AccumuloTableHandle) insertHandle;
        throw new PrestoException(NOT_SUPPORTED, format("Unable to rollback insert for table %s.%s. Some rows may have been written. Please run your insert again.", handle.getSchema(), handle.getTable()));
    }

    @Override
    public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTableName tableName)
    {
        if (!listSchemaNames(session).contains(tableName.getSchemaName().toLowerCase(Locale.ENGLISH))) {
            return null;
        }

        // Need to validate that SchemaTableName is a table
        if (!this.listViews(session, Optional.of(tableName.getSchemaName())).contains(tableName)) {
            AccumuloTable table = client.getTable(tableName);
            if (table == null) {
                return null;
            }

            return new AccumuloTableHandle(
                    table.getSchema(),
                    table.getTable(),
                    table.getRowId(),
                    table.isExternal(),
                    table.getSerializerClassName(),
                    table.getScanAuthorizations());
        }

        return null;
    }

    @Override
    public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle table)
    {
        AccumuloTableHandle handle = (AccumuloTableHandle) table;
        SchemaTableName tableName = new SchemaTableName(handle.getSchema(), handle.getTable());
        ConnectorTableMetadata metadata = getTableMetadata(tableName);
        if (metadata == null) {
            throw new TableNotFoundException(tableName);
        }
        return metadata;
    }

    @Override
    public Map getColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle)
    {
        AccumuloTableHandle handle = (AccumuloTableHandle) tableHandle;

        AccumuloTable table = client.getTable(handle.toSchemaTableName());
        if (table == null) {
            throw new TableNotFoundException(handle.toSchemaTableName());
        }

        ImmutableMap.Builder columnHandles = ImmutableMap.builder();
        for (AccumuloColumnHandle column : table.getColumns()) {
            columnHandles.put(column.getName(), column);
        }
        return columnHandles.build();
    }

    @Override
    public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle)
    {
        return ((AccumuloColumnHandle) columnHandle).getColumnMetadata();
    }

    @Override
    public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle source, String target)
    {
        AccumuloTableHandle handle = (AccumuloTableHandle) tableHandle;
        AccumuloColumnHandle columnHandle = (AccumuloColumnHandle) source;
        AccumuloTable table = client.getTable(handle.toSchemaTableName());
        if (table == null) {
            throw new TableNotFoundException(new SchemaTableName(handle.getSchema(), handle.getTable()));
        }

        client.renameColumn(table, columnHandle.getName(), target);
    }

    @Override
    public List listSchemaNames(ConnectorSession session)
    {
        return ImmutableList.copyOf(client.getSchemaNames());
    }

    @Override
    public List listTables(ConnectorSession session, Optional filterSchema)
    {
        Set schemaNames = filterSchema.>map(ImmutableSet::of)
                .orElseGet(client::getSchemaNames);

        ImmutableList.Builder builder = ImmutableList.builder();
        for (String schemaName : schemaNames) {
            for (String tableName : client.getTableNames(schemaName)) {
                builder.add(new SchemaTableName(schemaName, tableName));
            }
        }
        return builder.build();
    }

    @Override
    public Map> listTableColumns(ConnectorSession session, SchemaTablePrefix prefix)
    {
        requireNonNull(prefix, "prefix is null");
        ImmutableMap.Builder> columns = ImmutableMap.builder();
        for (SchemaTableName tableName : listTables(session, prefix)) {
            ConnectorTableMetadata tableMetadata = getTableMetadata(tableName);
            // table can disappear during listing operation
            if (tableMetadata != null) {
                columns.put(tableName, tableMetadata.getColumns());
            }
        }
        return columns.build();
    }

    @Override
    public boolean usesLegacyTableLayouts()
    {
        return false;
    }

    @Override
    public Optional> applyFilter(ConnectorSession session, ConnectorTableHandle table, Constraint constraint)
    {
        AccumuloTableHandle handle = (AccumuloTableHandle) table;

        TupleDomain oldDomain = handle.getConstraint();
        TupleDomain newDomain = oldDomain.intersect(constraint.getSummary());
        if (oldDomain.equals(newDomain)) {
            return Optional.empty();
        }

        handle = new AccumuloTableHandle(
                handle.getSchema(),
                handle.getTable(),
                handle.getRowId(),
                newDomain,
                handle.isExternal(),
                handle.getSerializerClassName(),
                handle.getScanAuthorizations());

        return Optional.of(new ConstraintApplicationResult<>(handle, constraint.getSummary()));
    }

    @Override
    public ConnectorTableProperties getTableProperties(ConnectorSession session, ConnectorTableHandle handle)
    {
        return new ConnectorTableProperties();
    }

    private void checkNoRollback()
    {
        checkState(rollbackAction.get() == null, "Cannot begin a new write while in an existing one");
    }

    private void setRollback(Runnable action)
    {
        checkState(rollbackAction.compareAndSet(null, action), "Should not have to override existing rollback action");
    }

    private void clearRollback()
    {
        rollbackAction.set(null);
    }

    public void rollback()
    {
        Runnable rollbackAction = this.rollbackAction.getAndSet(null);
        if (rollbackAction != null) {
            rollbackAction.run();
        }
    }

    private ConnectorTableMetadata getTableMetadata(SchemaTableName tableName)
    {
        if (!client.getSchemaNames().contains(tableName.getSchemaName())) {
            return null;
        }

        // Need to validate that SchemaTableName is a table
        if (!this.listViews(Optional.ofNullable(tableName.getSchemaName())).contains(tableName)) {
            AccumuloTable table = client.getTable(tableName);
            if (table == null) {
                return null;
            }

            return new ConnectorTableMetadata(tableName, table.getColumnsMetadata());
        }

        return null;
    }

    private List listTables(ConnectorSession session, SchemaTablePrefix prefix)
    {
        // List all tables if schema or table is null
        if (!prefix.getTable().isPresent()) {
            return listTables(session, prefix.getSchema());
        }

        // Make sure requested table exists, returning the single table of it does
        SchemaTableName table = prefix.toSchemaTableName();
        if (getTableHandle(session, table) != null) {
            return ImmutableList.of(table);
        }

        // Else, return empty list
        return ImmutableList.of();
    }
}