Search in sources :

Example 1 with PartitionKey

use of com.azure.cosmos.models.PartitionKey in project YCSB by brianfrankcooper.

the class AzureCosmosClient method delete.

@Override
public Status delete(String table, String key) {
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("Delete key {} from table {}", key, table);
    }
    try {
        CosmosContainer container = AzureCosmosClient.containerCache.get(table);
        if (container == null) {
            container = AzureCosmosClient.database.getContainer(table);
            AzureCosmosClient.containerCache.put(table, container);
        }
        container.deleteItem(key, new PartitionKey(key), new CosmosItemRequestOptions());
        return Status.OK;
    } catch (Exception e) {
        if (!AzureCosmosClient.includeExceptionStackInLog) {
            e = null;
        }
        LOGGER.error("Failed to delete key {} in collection {}", key, table, e);
    }
    return Status.ERROR;
}
Also used : CosmosItemRequestOptions(com.azure.cosmos.models.CosmosItemRequestOptions) CosmosContainer(com.azure.cosmos.CosmosContainer) PartitionKey(com.azure.cosmos.models.PartitionKey) CosmosException(com.azure.cosmos.CosmosException) DBException(site.ycsb.DBException)

Example 2 with PartitionKey

use of com.azure.cosmos.models.PartitionKey in project YCSB by brianfrankcooper.

the class AzureCosmosClient method update.

/**
 * Update a record in the database. Any field/value pairs in the specified
 * values HashMap will be written into the record with the specified record key,
 * overwriting any existing values with the same field name.
 *
 * @param table  The name of the table
 * @param key    The record key of the record to write.
 * @param values A HashMap of field/value pairs to update in the record
 * @return Zero on success, a non-zero error code on error
 */
@Override
public Status update(String table, String key, Map<String, ByteIterator> values) {
    String readEtag = "";
    // that will be a future improvement.
    for (int attempt = 0; attempt < NUM_UPDATE_ATTEMPTS; attempt++) {
        try {
            CosmosContainer container = AzureCosmosClient.containerCache.get(table);
            if (container == null) {
                container = AzureCosmosClient.database.getContainer(table);
                AzureCosmosClient.containerCache.put(table, container);
            }
            CosmosItemResponse<ObjectNode> response = container.readItem(key, new PartitionKey(key), ObjectNode.class);
            readEtag = response.getETag();
            ObjectNode node = response.getItem();
            for (Entry<String, ByteIterator> pair : values.entrySet()) {
                node.put(pair.getKey(), pair.getValue().toString());
            }
            CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions();
            requestOptions.setIfMatchETag(readEtag);
            PartitionKey pk = new PartitionKey(key);
            container.replaceItem(node, key, pk, requestOptions);
            return Status.OK;
        } catch (CosmosException e) {
            if (!AzureCosmosClient.includeExceptionStackInLog) {
                e = null;
            }
            LOGGER.error("Failed to update key {} to collection {} in database {} on attempt {}", key, table, AzureCosmosClient.databaseName, attempt, e);
        }
    }
    return Status.ERROR;
}
Also used : ObjectNode(com.fasterxml.jackson.databind.node.ObjectNode) ByteIterator(site.ycsb.ByteIterator) StringByteIterator(site.ycsb.StringByteIterator) CosmosItemRequestOptions(com.azure.cosmos.models.CosmosItemRequestOptions) CosmosContainer(com.azure.cosmos.CosmosContainer) PartitionKey(com.azure.cosmos.models.PartitionKey) CosmosException(com.azure.cosmos.CosmosException)

Example 3 with PartitionKey

use of com.azure.cosmos.models.PartitionKey in project YCSB by brianfrankcooper.

the class AzureCosmosClient method insert.

/**
 * Insert a record in the database. Any field/value pairs in the specified
 * values HashMap will be written into the record with the specified record key.
 *
 * @param table  The name of the table
 * @param key    The record key of the record to insert.
 * @param values A HashMap of field/value pairs to insert in the record
 * @return Zero on success, a non-zero error code on error
 */
@Override
public Status insert(String table, String key, Map<String, ByteIterator> values) {
    if (LOGGER.isDebugEnabled()) {
        LOGGER.debug("Insert key: {} into table: {}", key, table);
    }
    try {
        CosmosContainer container = AzureCosmosClient.containerCache.get(table);
        if (container == null) {
            container = AzureCosmosClient.database.getContainer(table);
            AzureCosmosClient.containerCache.put(table, container);
        }
        PartitionKey pk = new PartitionKey(key);
        ObjectNode node = OBJECT_MAPPER.createObjectNode();
        node.put("id", key);
        for (Map.Entry<String, ByteIterator> pair : values.entrySet()) {
            node.put(pair.getKey(), pair.getValue().toString());
        }
        if (AzureCosmosClient.useUpsert) {
            container.upsertItem(node, pk, new CosmosItemRequestOptions());
        } else {
            container.createItem(node, pk, new CosmosItemRequestOptions());
        }
        return Status.OK;
    } catch (CosmosException e) {
        if (!AzureCosmosClient.includeExceptionStackInLog) {
            e = null;
        }
        LOGGER.error("Failed to insert key {} to collection {} in database {}", key, table, AzureCosmosClient.databaseName, e);
    }
    return Status.ERROR;
}
Also used : ObjectNode(com.fasterxml.jackson.databind.node.ObjectNode) ByteIterator(site.ycsb.ByteIterator) StringByteIterator(site.ycsb.StringByteIterator) CosmosItemRequestOptions(com.azure.cosmos.models.CosmosItemRequestOptions) CosmosContainer(com.azure.cosmos.CosmosContainer) PartitionKey(com.azure.cosmos.models.PartitionKey) CosmosException(com.azure.cosmos.CosmosException) HashMap(java.util.HashMap) Map(java.util.Map) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap)

Example 4 with PartitionKey

use of com.azure.cosmos.models.PartitionKey in project YCSB by brianfrankcooper.

the class AzureCosmosClient method read.

/**
 * Read a record from the database. Each field/value pair from the result will
 * be stored in a HashMap.
 *
 * @param table  The name of the table
 * @param key    The record key of the record to read.
 * @param fields The list of fields to read, or null for all of them
 * @param result A HashMap of field/value pairs for the result
 * @return Zero on success, a non-zero error code on error
 */
@Override
public Status read(String table, String key, Set<String> fields, Map<String, ByteIterator> result) {
    try {
        CosmosContainer container = AzureCosmosClient.containerCache.get(table);
        if (container == null) {
            container = AzureCosmosClient.database.getContainer(table);
            AzureCosmosClient.containerCache.put(table, container);
        }
        CosmosItemResponse<ObjectNode> response = container.readItem(key, new PartitionKey(key), ObjectNode.class);
        ObjectNode node = response.getItem();
        Map<String, String> stringResults = new HashMap<>(node.size());
        if (fields == null) {
            Iterator<Map.Entry<String, JsonNode>> iter = node.fields();
            while (iter.hasNext()) {
                Entry<String, JsonNode> pair = iter.next();
                stringResults.put(pair.getKey().toString(), pair.getValue().toString());
            }
            StringByteIterator.putAllAsByteIterators(result, stringResults);
        } else {
            Iterator<Map.Entry<String, JsonNode>> iter = node.fields();
            while (iter.hasNext()) {
                Entry<String, JsonNode> pair = iter.next();
                if (fields.contains(pair.getKey())) {
                    stringResults.put(pair.getKey().toString(), pair.getValue().toString());
                }
            }
            StringByteIterator.putAllAsByteIterators(result, stringResults);
        }
        return Status.OK;
    } catch (CosmosException e) {
        LOGGER.error("Failed to read key {} in collection {} in database {}", key, table, AzureCosmosClient.databaseName, e);
        return Status.NOT_FOUND;
    }
}
Also used : Entry(java.util.Map.Entry) ObjectNode(com.fasterxml.jackson.databind.node.ObjectNode) HashMap(java.util.HashMap) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) CosmosContainer(com.azure.cosmos.CosmosContainer) PartitionKey(com.azure.cosmos.models.PartitionKey) JsonNode(com.fasterxml.jackson.databind.JsonNode) CosmosException(com.azure.cosmos.CosmosException)

Example 5 with PartitionKey

use of com.azure.cosmos.models.PartitionKey in project cas by apereo.

the class CosmosDbServiceRegistry method findServiceById.

@Override
public RegisteredService findServiceById(final long id) {
    try {
        val key = String.valueOf(id);
        LOGGER.debug("Reading registered services with id [{}] from [{}]", key, container.getId());
        val doc = container.readItem(key, new PartitionKey(key), CosmosDbDocument.class).getItem();
        return getRegisteredServiceFromDocumentBody(doc);
    } catch (final CosmosException e) {
        if (e.getStatusCode() == HttpStatus.NOT_FOUND.value()) {
            LOGGER.debug("Unable to locate registered service with id [{}]", id);
            LOGGER.trace(e.getMessage(), e);
        } else {
            LoggingUtils.warn(LOGGER, e);
        }
    }
    return null;
}
Also used : lombok.val(lombok.val) PartitionKey(com.azure.cosmos.models.PartitionKey) CosmosDbDocument(org.apereo.cas.cosmosdb.CosmosDbDocument) CosmosException(com.azure.cosmos.CosmosException)

Aggregations

CosmosException (com.azure.cosmos.CosmosException)5 PartitionKey (com.azure.cosmos.models.PartitionKey)5 CosmosContainer (com.azure.cosmos.CosmosContainer)4 CosmosItemRequestOptions (com.azure.cosmos.models.CosmosItemRequestOptions)3 ObjectNode (com.fasterxml.jackson.databind.node.ObjectNode)3 HashMap (java.util.HashMap)2 ConcurrentHashMap (java.util.concurrent.ConcurrentHashMap)2 ByteIterator (site.ycsb.ByteIterator)2 StringByteIterator (site.ycsb.StringByteIterator)2 JsonNode (com.fasterxml.jackson.databind.JsonNode)1 Map (java.util.Map)1 Entry (java.util.Map.Entry)1 lombok.val (lombok.val)1 CosmosDbDocument (org.apereo.cas.cosmosdb.CosmosDbDocument)1 DBException (site.ycsb.DBException)1