Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  package de.zalando.sprocwrapper.proxy;
  
  
  import java.sql.Statement;
  
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
 
 
 
 
 

Author(s):
jmussler
 
 class StoredProcedure {
 
     private static final int TRUNCATE_DEBUG_PARAMS_MAX_LENGTH = 1024;
     private static final String TRUNCATE_DEBUG_PARAMS_ELLIPSIS = " ...";
 
     private static final Logger LOG = LoggerFactory.getLogger(StoredProcedure.class);
 
 
     private final String name;
     private String query = null;
     private Class<?> returnType = null;
 
     // whether the result type is a collection (List)
     private boolean collectionResult = false;
     private final boolean runOnAllShards;
     private final boolean searchShards;
     private boolean autoPartition;
     private final boolean parallel;
     private final boolean readOnly;
     private final WriteTransaction writeTransaction;
 
     private Executor executor = null;
 
     private List<ShardKeyParametershardKeyParameters = null;
     private final RowMapper<?> resultMapper;
 
     private int[] types = null;
 
     private static final Executor MULTI_ROW_SIMPLE_TYPE_EXECUTOR = new MultiRowSimpleTypeExecutor();
     private static final Executor MULTI_ROW_TYPE_MAPPER_EXECUTOR = new MultiRowTypeMapperExecutor();
     private static final Executor SINGLE_ROW_SIMPLE_TYPE_EXECUTOR = new SingleRowSimpleTypeExecutor();
     private static final Executor SINGLE_ROW_TYPE_MAPPER_EXECUTOR = new SingleRowTypeMapperExecutor();
 
     private final long timeout;
     private final AdvisoryLock adivsoryLock;
 
     public StoredProcedure(final String namefinal java.lang.reflect.Type genericType,
             final VirtualShardKeyStrategy sStrategyfinal boolean runOnAllShardsfinal boolean searchShards,
             final boolean parallelfinal RowMapper<?> resultMapperfinal long timeout,
             final AdvisoryLock advisoryLockfinal boolean useValidationfinal boolean readOnly,
             final WriteTransaction writeTransactionthrows InstantiationExceptionIllegalAccessException {
        this. = name;
        this. = runOnAllShards;
        this. = searchShards;
        this. = parallel;
        this. = readOnly;
        this. = resultMapper;
        this. = writeTransaction;
        this. = advisoryLock;
        this. = timeout;
         = sStrategy;
        ValueTransformer<?, ?> valueTransformerForClass = null;
        if (genericType instanceof ParameterizedType) {
            final ParameterizedType pType = (ParameterizedTypegenericType;
            if (java.util.List.class.isAssignableFrom((Class<?>) pType.getRawType())
                    && pType.getActualTypeArguments().length > 0) {
                 = (Class<?>) pType.getActualTypeArguments()[0];
                // check if we have a value transformer (and initialize the registry):
                valueTransformerForClass = GlobalValueTransformerLoader.getValueTransformerForClass();
                if (valueTransformerForClass != null
                        || ..containsKey()) {
                     = ;
                } else {
                     = ;
                }
                 = true;
            } else {
                 = ;
                 = (Class<?>) pType.getRawType();
            }
        } else {
             = (Class<?>) genericType;
            // check if we have a value transformer (and initialize the registry):
            valueTransformerForClass = GlobalValueTransformerLoader.getValueTransformerForClass();
            if (valueTransformerForClass != null || ..containsKey()) {
                 = ;
            } else {
                if (resultMapper != null) {
                     = new SingleRowCustomMapperExecutor(resultMapper);
                } else {
                     = ;
                }
            }
        }
        if (this. > 0 || this. != .) {
            // Wrapper provides locking and changing of session settings functionality
            this. = new ExecutorWrapper(this.this.);
        }
        if (useValidation) {
            this. = new ValidationExecutorWrapper(this.);
        }
        if (valueTransformerForClass != null) {
            // we need to transform the return value by the global value transformer.
            // add the transformation to the as a transformerExecutor
            this. = new GlobalTransformerExecutorWrapper(this.);
        }
    }
    public void addParam(final StoredProcedureParameter p) {
        .add(p);
    }
    public void setVirtualShardKeyStrategy(final VirtualShardKeyStrategy s) {
         = s;
    }
    public void addShardKeyParameter(final int jpfinal Class<?> clazz) {
        if ( == null) {
             = new ArrayList<ShardKeyParameter>(1);
        }
        if (List.class.isAssignableFrom(clazz)) {
             = true;
        }
        .add(new ShardKeyParameter(jp));
    }
    public String getName() {
        return ;
    }
    public Object[] getParams(final Object[] origParamsfinal Connection connection) {
        final Object[] ps = new Object[.size()];
        int i = 0;
        for (final StoredProcedureParameter p : ) {
            try {
                ps[i] = p.mapParam(origParams[p.getJavaPos()], connection);
            } catch (final Exception e) {
                final String errorMessage = "Could not map input parameter for stored procedure " +  + " of type "
                        + p.getType() + " at position " + p.getJavaPos() + ": "
                        + (p.isSensitive() ? "<SENSITIVE>" : origParams[p.getJavaPos()]);
                .error(errorMessagee);
                throw new IllegalArgumentException(errorMessagee);
            }
            i++;
        }
        return ps;
    }
    public int[] getTypes() {
        if ( == null) {
             = new int[.size()];
            int i = 0;
            for (final StoredProcedureParameter p : ) {
                [i++] = p.getType();
            }
        }
        return ;
    }
    public int getShardId(final Object[] objs) {
        if ( == null) {
            return .getShardId(null);
        }
        final Object[] keys = new Object[.size()];
        int i = 0;
        Object obj;
        for (final ShardKeyParameter p : ) {
            obj = objs[p.javaPos];
            if (obj instanceof ShardedObject) {
                obj = ((ShardedObjectobj).getShardKey();
            }
            keys[i] = obj;
            i++;
        }
        return .getShardId(keys);
    }
    public String getSqlParameterList() {
        String s = "";
        boolean first = true;
        for (int i = 1; i <= .size(); ++i) {
            if (!first) {
                s += ",";
            }
            first = false;
            s += "?";
        }
        return s;
    }
    public void setQuery(final String sql) {
         = sql;
    }
    public String getQuery() {
        if ( == null) {
             = "SELECT * FROM " +  + " ( " + getSqlParameterList() + " )";
        }
        return ;
    }

    
build execution string like create_or_update_multiple_objects({"(a,b)","(c,d)" }).

Parameters:
args
Returns:
    private String getDebugLog(final Object[] args) {
        final StringBuilder sb = new StringBuilder();
        sb.append('(');
        int i = 0;
        for (final Object param : args) {
            if (i > 0) {
                sb.append(',');
            }
            if (param == null) {
                sb.append("NULL");
            } else if (.get(i).isSensitive()) {
                sb.append("<SENSITIVE>");
            } else {
                sb.append(param);
            }
            i++;
            if (sb.length() > ) {
                break;
            }
        }
        if (sb.length() > ) {
            // Truncate params for debug output
            return sb.substring(0, ) +  + ")";
        } else {
            sb.append(')');
            return sb.toString();
        }
    }

    
split arguments by shard.

Parameters:
dataSourceProvider
args the original argument list
Returns:
map of virtual shard ID to argument list (TreeMap with ordered keys: sorted by shard ID)
    @SuppressWarnings("unchecked")
    private Map<IntegerObject[]> partitionArguments(final DataSourceProvider dataSourceProvider,
            final Object[] args) {
        // use TreeMap here to maintain ordering by shard ID
        final Map<IntegerObject[]> argumentsByShardId = Maps.newTreeMap();
        // we need to partition by datasource instead of virtual shard ID (different virtual shard IDs are mapped to
        // the same datasource e.g. by VirtualShardMd5Strategy)
        final Map<DataSourceIntegershardIdByDataSource = Maps.newHashMap();
        // TODO: currently only implemented for single shardKey argument as first argument!
        final List<ObjectoriginalArgument = (List<Object>) args[0];
        if (originalArgument == null || originalArgument.isEmpty()) {
            throw new IllegalArgumentException("ShardKey (first argument) of sproc '" +  + "' not defined");
        }
        List<ObjectpartitionedArgument = null;
        Object[] partitionedArguments = null;
        int shardId;
        Integer existingShardId;
        DataSource dataSource;
        for (final Object key : originalArgument) {
            shardId = getShardId(new Object[] {key});
            dataSource = dataSourceProvider.getDataSource(shardId);
            existingShardId = shardIdByDataSource.get(dataSource);
            if (existingShardId != null) {
                // we already saw the same datasource => use the virtual shard ID of the first argument with the same
                // datasource
                shardId = existingShardId;
            } else {
                shardIdByDataSource.put(dataSourceshardId);
            }
            partitionedArguments = argumentsByShardId.get(shardId);
            if (partitionedArguments == null) {
                partitionedArgument = Lists.newArrayList();
                partitionedArguments = new Object[args.length];
                partitionedArguments[0] = partitionedArgument;
                if (args.length > 1) {
                    System.arraycopy(args, 1, partitionedArguments, 1, args.length - 1);
                }
                argumentsByShardId.put(shardIdpartitionedArguments);
            } else {
                partitionedArgument = (List<Object>) partitionedArguments[0];
            }
            partitionedArgument.add(key);
        }
        return argumentsByShardId;
    }
    private static class Call implements Callable<Object> {
        private final StoredProcedure sproc;
        private final DataSource shardDs;
        private final Object[] params;
        private final InvocationContext invocation;
        public Call(final StoredProcedure sprocfinal DataSource shardDsfinal Object[] params,
                final InvocationContext invocation) {
            this. = sproc;
            this. = shardDs;
            this. = params;
            this. = invocation;
        }
        @Override
        public Object call() throws Exception {
            return ..executeSProc(.getQuery(), .getTypes(), ,
                    .);
        }
    }
    private static ExecutorService parallelThreadPool = Executors.newCachedThreadPool();
    public Object execute(final DataSourceProvider dpfinal InvocationContext invocation) {
        List<IntegershardIds = null;
        Map<IntegerObject[]> partitionedArguments = null;
        if ( || ) {
            shardIds = dp.getDistinctShardIds();
        } else {
            if () {
                partitionedArguments = partitionArguments(dpinvocation.getArgs());
                shardIds = Lists.newArrayList(partitionedArguments.keySet());
            } else {
                shardIds = Lists.newArrayList(getShardId(invocation.getArgs()));
            }
        }
        if (partitionedArguments == null) {
            partitionedArguments = Maps.newHashMap();
            for (final int shardId : shardIds) {
                partitionedArguments.put(shardIdinvocation.getArgs());
            }
        }
        final DataSource firstDs = dp.getDataSource(shardIds.get(0));
        Connection connection = null;
        try {
            connection = firstDs.getConnection();
        } catch (final SQLException e) {
            throw new CannotGetJdbcConnectionException("Failed to acquire connection for virtual shard "
                    + shardIds.get(0) + " for " + e);
        }
        final List<Object[]> paramValues = Lists.newArrayList();
        try {
            for (final int shardId : shardIds) {
                paramValues.add(getParams(partitionedArguments.get(shardId), connection));
            }
        } finally {
            if (connection != null) {
                try {
                    connection.close();
                } catch (final Throwable t) {
                    .warn("Could not release connection"t);
                }
            }
        }
        if (shardIds.size() == 1 && !) {
            if (.isDebugEnabled()) {
                .debug(getDebugLog(paramValues.get(0)));
            }
            // most common case: only one shard and no argument partitioning
            return .executeSProc(firstDsgetQuery(), paramValues.get(0), getTypes(), invocation);
        } else {
            Map<IntegerSameConnectionDatasourcetransactionalDatasources = null;
            try {
                // we may need to start a transaction context
                transactionalDatasources = startTransaction(dpshardIds);
                final List<?> results = Lists.newArrayList();
                Object sprocResult = null;
                final long start = System.currentTimeMillis();
                if () {
                    sprocResult = executeInParallel(dpinvocationshardIdsparamValuestransactionalDatasources,
                            resultssprocResult);
                } else {
                    sprocResult = executeSequential(dpinvocationshardIdsparamValuestransactionalDatasources,
                            resultssprocResult);
                }
                if (.isTraceEnabled()) {
                    .trace("[{}] execution of [{}] on [{}] shards took [{}] ms",
                        new Object[] {
                             ? "parallel" : "serial"shardIds.size(), System.currentTimeMillis() - start
                        });
                }
                // no error - we may need to commit
                commitTransaction(transactionalDatasources);
                if () {
                    return results;
                } else {
                    // return last result
                    return sprocResult;
                }
            } catch (final RuntimeException runtimeException) {
                .trace("[{}] execution of [{}] on [{}] shards aborted by runtime exception [{}]",
                    new Object[] {
                         ? "parallel" : "serial"shardIds.size(), runtimeException.getMessage(),
                        runtimeException
                    });
                // error occured, we may need to rollback all transactions.
                rollbackTransaction(transactionalDatasources);
                // re-throw
                throw runtimeException;
            } catch (final Throwable throwable) {
                .trace("[{}] execution of [{}] on [{}] shards aborted by throwable exception [{}]",
                    new Object[] {
                         ? "parallel" : "serial"shardIds.size(), throwable.getMessage(), throwable
                    });
                // error occured, we may need to rollback all transactions.
                rollbackTransaction(transactionalDatasources);
                // throw runtime:
                throw new RuntimeException(throwable);
            }
        }
    }
    @SuppressWarnings({ "rawtypes""unchecked" })
    private Object executeSequential(final DataSourceProvider dpfinal InvocationContext invocation,
            final List<IntegershardIdsfinal List<Object[]> paramValues,
            final Map<IntegerSameConnectionDatasourcetransactionalDatasourcesfinal List<?> results,
            Object sprocResult) {
        DataSource shardDs;
        int i = 0;
        final List<Stringexceptions = Lists.newArrayList();
        final ImmutableMap.Builder<IntegerThrowablecauses = ImmutableMap.builder();
        for (final int shardId : shardIds) {
            shardDs = getShardDs(dptransactionalDatasourcesshardId);
            if (.isDebugEnabled()) {
                .debug(getDebugLog(paramValues.get(i)));
            }
            sprocResult = null;
            try {
                sprocResult = .executeSProc(shardDsgetQuery(), paramValues.get(i), getTypes(), invocation,
                        );
            } catch (final Exception e) {
                // remember all exceptions and go on
                exceptions.add("shardId: " + shardId + ", message: " + e.getMessage() + ", query: " + getQuery());
                causes.put(shardIde);
            }
            if (addResultsBreakWhenSharded(resultssprocResult)) {
                break;
            }
            i++;
        }
        if (!exceptions.isEmpty()) {
            throw new ShardedDataAccessException("Got exception(s) while executing sproc on shards: "
                    + Joiner.on(", ").join(exceptions), causes.build());
        }
        return sprocResult;
    }
    @SuppressWarnings({ "rawtypes""unchecked" })
    private Object executeInParallel(final DataSourceProvider dpfinal InvocationContext invocation,
            final List<IntegershardIdsfinal List<Object[]> paramValues,
            final Map<IntegerSameConnectionDatasourcetransactionalDatasourcesfinal List<?> results,
            Object sprocResult) {
        DataSource shardDs;
        final Map<IntegerFutureTask<Object>> tasks = Maps.newHashMapWithExpectedSize(shardIds.size());
        FutureTask<Objecttask;
        int i = 0;
        for (final int shardId : shardIds) {
            shardDs = getShardDs(dptransactionalDatasourcesshardId);
            if (.isDebugEnabled()) {
                .debug(getDebugLog(paramValues.get(i)));
            }
            task = new FutureTask<Object>(new Call(thisshardDsparamValues.get(i), invocation));
            tasks.put(shardIdtask);
            .execute(task);
            i++;
        }
        final List<Stringexceptions = Lists.newArrayList();
        final ImmutableMap.Builder<IntegerThrowablecauses = ImmutableMap.builder();
        for (final Entry<IntegerFutureTask<Object>> taskToFinish : tasks.entrySet()) {
            try {
                sprocResult = taskToFinish.getValue().get();
            } catch (final InterruptedException ex) {
                // remember all exceptions and go on
                exceptions.add("got sharding execution exception: " + ex.getMessage() + ", query: " + getQuery());
                causes.put(taskToFinish.getKey(), ex);
            } catch (final ExecutionException ex) {
                // remember all exceptions and go on
                exceptions.add("got sharding execution exception: " + ex.getCause().getMessage() + ", query: "
                        + getQuery());
                causes.put(taskToFinish.getKey(), ex.getCause());
            }
            if (addResultsBreakWhenSharded(resultssprocResult)) {
                break;
            }
        }
        if (!exceptions.isEmpty()) {
            throw new ShardedDataAccessException("Got exception(s) while executing sproc on shards: "
                    + Joiner.on(", ").join(exceptions), causes.build());
        }
        return sprocResult;
    }
    @SuppressWarnings({ "rawtypes""unchecked" })
    private boolean addResultsBreakWhenSharded(final Collection resultsfinal Object sprocResult) {
        boolean breakSearch = false;
        if ( && sprocResult != null && !((CollectionsprocResult).isEmpty()) {
            // Result is a non-empty collection
            results.addAll((CollectionsprocResult);
            // Break if shardedSearch
            breakSearch = ;
        } else if (! && sprocResult != null && ) {
            // Result is non-null, but not a collection
            // Break if shardedSearch
            breakSearch = true;
        }
        return breakSearch;
    }
    private DataSource getShardDs(final DataSourceProvider dp,
            final Map<IntegerSameConnectionDatasourcetransactionIdsfinal int shardId) {
        if (transactionIds.isEmpty()) {
            return dp.getDataSource(shardId);
        }
        return transactionIds.get(shardId);
    }
            final List<IntegershardIdsthrows SQLException {
        final Map<IntegerSameConnectionDatasourceret = Maps.newHashMap();
        if ( == false &&  != .) {
            for (final int shardId : shardIds) {
                final DataSource shardDs = dp.getDataSource(shardId);
                // we need to pin the calls to a single connection
                final SameConnectionDatasource sameConnDs = new SameConnectionDatasource(shardDs.getConnection());
                ret.put(shardIdsameConnDs);
                .trace("startTransaction on shard [{}]"shardId);
                final Statement st = sameConnDs.getConnection().createStatement();
                st.execute("BEGIN");
                st.close();
            }
        }
        return ret;
    }
    private void commitTransaction(final Map<IntegerSameConnectionDatasourcedatasources) {
        if ( == false &&  != .) {
            if ( == .) {
                for (final Entry<IntegerSameConnectionDatasourceshardEntry : datasources.entrySet()) {
                    try {
                        .trace("commitTransaction on shard [{}]"shardEntry.getKey());
                        final DataSource shardDs = shardEntry.getValue();
                        final Statement st = shardDs.getConnection().createStatement();
                        st.execute("COMMIT");
                        st.close();
                        shardEntry.getValue().close();
                    } catch (final Exception e) {
                        // do our best. we cannot rollback at this point.
                        // store other shards as much as possible.
                        .error(
                            "ERROR: could not commitTransaction on shard [{}] - this will produce inconsistent data.",
                            shardEntry.getKey(), e);
                    }
                }
            } else if ( == .) {
                boolean commitFailed = false;
                final String transactionId = "sprocwrapper_" + UUID.randomUUID();
                final String prepareTransactionStatement = "PREPARE TRANSACTION '" + transactionId + "'";
                for (final Entry<IntegerSameConnectionDatasourceshardEntry : datasources.entrySet()) {
                    try {
                        .trace("prepare transaction on shard [{}]"shardEntry.getKey());
                        final DataSource shardDs = shardEntry.getValue();
                        final Statement st = shardDs.getConnection().createStatement();
                        st.execute(prepareTransactionStatement);
                        st.close();
                    } catch (final Exception e) {
                        commitFailed = true;
                        // log, but go on, prepare other transactions - but they will be removed as well.
                        .debug("prepare transaction [{}] on shard [{}] failed!",
                            new Object[] {transactionIdshardEntry.getKey(), e});
                    }
                }
                if (commitFailed) {
                    rollbackPrepared(datasourcestransactionId);
                } else {
                    final String commitStatement = "COMMIT PREPARED '" + transactionId + "'";
                    for (final Entry<IntegerSameConnectionDatasourceshardEntry : datasources.entrySet()) {
                        try {
                            .trace("commit prepared transaction [{}] on shard [{}]"transactionId,
                                shardEntry.getKey());
                            final DataSource shardDs = shardEntry.getValue();
                            final Statement st = shardDs.getConnection().createStatement();
                            st.execute(commitStatement);
                            st.close();
                            shardEntry.getValue().close();
                        } catch (final Exception e) {
                            commitFailed = true;
                            // do our best. we cannot rollback at this point.
                            // store other shards as much as possible.
                            // the not yet stored transactions are visible in postgres prepared transactions
                            // a nagios check should detect them so that we can handle any errors
                            // that may be produced at this point.
                            .error(
                                "FAILED: could not commit prepared transaction [{}] on shard [{}] - this will produce inconsistent data.",
                                new Object[] {transactionIdshardEntry.getKey(), e});
                        }
                    }
                    // for all failed commits:
                    if (commitFailed) {
                        rollbackPrepared(datasourcestransactionId);
                    }
                }
            } else {
                throw new IllegalArgumentException("Unknown writeTransaction state: " + );
            }
        }
    }
    private void rollbackPrepared(final Map<IntegerSameConnectionDatasourcedatasources,
            final String transactionId) {
        final String rollbackQuery = "ROLLBACK PREPARED '" + transactionId + "'";
        for (final Entry<IntegerSameConnectionDatasourceshardEntry : datasources.entrySet()) {
            try {
                .error("rollback prepared transaction [{}] on shard [{}]"transactionIdshardEntry.getKey());
                final DataSource shardDs = shardEntry.getValue();
                final Statement st = shardDs.getConnection().createStatement();
                st.execute(rollbackQuery);
                st.close();
                shardEntry.getValue().close();
            } catch (final Exception e) {
                .error(
                    "FAILED: could not rollback prepared transaction [{}] on shard [{}] - this will produce inconsistent data.",
                    new Object[] {transactionIdshardEntry.getKey(), e});
            }
        }
    }
    private void rollbackTransaction(final Map<IntegerSameConnectionDatasourcedatasources) {
        if ( == false &&  != .) {
            for (final Entry<IntegerSameConnectionDatasourceshardEntry : datasources.entrySet()) {
                try {
                    .trace("rollbackTransaction on shard [{}]"shardEntry.getKey());
                    final DataSource shardDs = shardEntry.getValue();
                    final Statement st = shardDs.getConnection().createStatement();
                    st.execute("ROLLBACK");
                    st.close();
                    shardEntry.getValue().close();
                } catch (final Exception e) {
                    .error("ERROR: could not rollback on shard [{}] - this will produce inconsistent data.",
                        shardEntry.getKey());
                }
            }
        }
    }
    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder();
        sb.append('(');
        boolean f = true;
        for (final StoredProcedureParameter p : ) {
            if (!f) {
                sb.append(',');
            }
            f = false;
            sb.append(p.getType());
            if (!"".equals(p.getTypeName())) {
                sb.append("=>").append(p.getTypeName());
            }
        }
        sb.append(')');
        return sb.toString();
    }
New to GrepCode? Check out our FAQ X