private RequestOptions getRequestOptions(String key) { RequestOptions requestOptions = new RequestOptions(); requestOptions.setPartitionKey(new PartitionKey(key)); return requestOptions; }
private FeedOptions getFeedOptions(String key) { if (useSinglePartitionCollection) { return null; } FeedOptions feedOptions = new FeedOptions(); if (this.useHashQueryForScan) { feedOptions.setEnableCrossPartitionQuery(false); feedOptions.setPartitionKey(new PartitionKey(key)); } else { feedOptions.setEnableCrossPartitionQuery(true); feedOptions.setMaxDegreeOfParallelism(this.maxDegreeOfParallelismForQuery); } return feedOptions; }
@Override public Status delete(String table, String key) { Document record = getDocumentById(table, key); try { // Delete the document by self link. documentClient.deleteDocument(record.getSelfLink(), null); } catch (DocumentClientException e) { e.printStackTrace(); return Status.ERROR; } return Status.OK; }
ConsistencyLevel consistencyLevel = ConsistencyLevel.valueOf(this.getStringProperty( "azurecosmos.consistencyLevel", DEFAULT_CONSISTENCY_LEVEL)); ConnectionPolicy connectionPolicy = new ConnectionPolicy(); connectionPolicy.setEnableEndpointDiscovery(false); connectionPolicy.setConnectionMode(ConnectionMode.valueOf(connectionModeString)); connectionPolicy.setMaxPoolSize(this.getIntProperty("azurecosmos.maxConnectionPoolSize", connectionPolicy.getMaxPoolSize())); connectionPolicy.setIdleConnectionTimeout(this.getIntProperty("azurecosmos.idleConnectionTimeout", connectionPolicy.getIdleConnectionTimeout())); RetryOptions retryOptions = new RetryOptions(); retryOptions.setMaxRetryAttemptsOnThrottledRequests(this.getIntProperty( "azurecosmos.maxRetryAttemptsOnThrottledRequests", retryOptions.getMaxRetryAttemptsOnThrottledRequests())); retryOptions.setMaxRetryWaitTimeInSeconds(this.getIntProperty( "azurecosmos.maxRetryWaitTimeInSeconds", retryOptions.getMaxRetryWaitTimeInSeconds())); connectionPolicy.setRetryOptions(retryOptions); + "queryText={}", uri, connectionPolicy.getConnectionMode(), consistencyLevel.toString(), connectionPolicy.getRetryOptions().getMaxRetryAttemptsOnThrottledRequests(), connectionPolicy.getRetryOptions().getMaxRetryWaitTimeInSeconds(), this.useSinglePartitionCollection, this.useUpsert,
private Database getDatabase() { if (database == null) { // Get the database if it exists List<Database> databaseList = documentClient .queryDatabases( "SELECT * FROM root r WHERE r.id='" + databaseId + "'", null) .getQueryIterable() .toList(); if (databaseList.size() > 0) { // Cache the database object so we won't have to query for it // later to retrieve the selfLink. database = databaseList.get(0); } else { // Create the database if it doesn't exist. try { Database databaseDefinition = new Database(); databaseDefinition.setId(databaseId); database = documentClient.createDatabase(databaseDefinition, null) .getResource(); } catch (DocumentClientException e) { // TODO: Something has gone terribly wrong - the app wasn't // able to query or create the collection. // Verify your connection, endpoint, and key. e.printStackTrace(System.err); } } } return database; }
.queryCollections(getDatabase().getSelfLink(), "SELECT * FROM root r WHERE r.id='" + collectionId + "'", null) .getQueryIterable() .toList(); DocumentCollection collectionDefinition = new DocumentCollection(); collectionDefinition.setId(collectionId); .createCollection(getDatabase().getSelfLink(), collectionDefinition, null) .getResource(); } catch (DocumentClientException e) { e.printStackTrace(System.err); throw e;
FeedResponse<Document> feedResponse = null; try { feedResponse = AzureCosmosClient.client.queryDocuments(getDocumentCollectionLink(this.databaseName, table), new SqlQuerySpec(queryText, new SqlParameterCollection(new SqlParameter("@recordcount", recordcount), new SqlParameter("@startkey", startkey))), getFeedOptions(startkey)); documents = feedResponse.getQueryIterable().toList(); } catch (Exception e) { if (!this.includeExceptionStackInLog) { LOGGER.debug("Queried {} records for key {} in {}us - ActivityID: {}", recordcount, startkey, elapsed, feedResponse != null ? feedResponse.getActivityId() : NA_STRING);
private Document getDocumentById(String collectionId, String id) { if (collection == null) { return null; } // Retrieve the document using the DocumentClient. List<Document> documentList = documentClient .queryDocuments(collection.getSelfLink(), "SELECT * FROM root r WHERE r.id='" + id + "'", feedOptions) .getQueryIterable() .toList(); if (documentList.size() > 0) { return documentList.get(0); } return null; } }
@Override public Status insert(String table, String key, Map<String, ByteIterator> values) { Document record = new Document(); record.set("id", key); for (Entry<String, ByteIterator> val : values.entrySet()) { record.set(val.getKey(), val.getValue().toString()); } try { documentClient.createDocument(collection.getSelfLink(), record, null, false); } catch (DocumentClientException e) { e.printStackTrace(System.err); return Status.ERROR; } return Status.OK; }
@Override public void init() throws DBException { host = getProperties().getProperty("documentdb.host", null); masterKey = getProperties().getProperty("documentdb.masterKey", null); if (host == null) { System.err.println("ERROR: 'documentdb.host' must be set!"); System.exit(1); } if (masterKey == null) { System.err.println("ERROR: 'documentdb.masterKey' must be set!"); System.exit(1); } databaseId = getProperties().getProperty("documentdb.databaseId", "ycsb"); String collectionId = getProperties().getProperty("documentdb.collectionId", "usertable"); documentClient = new DocumentClient(host, masterKey, ConnectionPolicy.GetDefault(), ConsistencyLevel.Session); try { // Initialize test database and collection. collection = getCollection(collectionId); } catch (DocumentClientException e) { throw new DBException("Initialze collection failed", e); } feedOptions = new FeedOptions(); feedOptions.setEmitVerboseTracesInQuery(false); }
@Override public Status update(String table, String key, Map<String, ByteIterator> values) { String documentLink = getDocumentLink(this.databaseName, table, key); Document document = getDocumentDefinition(key, values); RequestOptions reqOptions = getRequestOptions(key); if (reqOptions == null) { reqOptions = new RequestOptions(); } AccessCondition accessCondition = new AccessCondition(); accessCondition.setCondition(document.getETag()); accessCondition.setType(AccessConditionType.IfMatch); reqOptions.setAccessCondition(accessCondition); ResourceResponse<Document> updatedResource = null; long startTime = System.nanoTime(); try { updatedResource = AzureCosmosClient.client.replaceDocument(documentLink, document, reqOptions); } catch (DocumentClientException e) { if (!this.includeExceptionStackInLog) { e = null; } LOGGER.error("Failed to update key {}", key, e); return Status.ERROR; } finally { long elapsed = (System.nanoTime() - startTime) / NS_IN_US; LOGGER.debug("Updated key {} in {}us - ActivityID: {}", key, elapsed, updatedResource != null ? updatedResource.getActivityId() : NA_STRING); } return Status.OK; }
@Override public Status read(String table, String key, Set<String> fields, Map<String, ByteIterator> result) { String documentLink = getDocumentLink(this.databaseName, table, key); ResourceResponse<Document> readResource = null; Document document; long startTime = System.nanoTime(); try { readResource = AzureCosmosClient.client.readDocument(documentLink, getRequestOptions(key)); document = readResource.getResource(); } catch (DocumentClientException e) { if (!this.includeExceptionStackInLog) { e = null; } LOGGER.error("Failed to read key {} in collection {} in database {}", key, table, this.databaseName, e); return Status.ERROR; } finally { long elapsed = (System.nanoTime() - startTime) / NS_IN_US; LOGGER.debug("Read key {} in {}us - ActivityID: {}", key, elapsed, readResource != null ? readResource.getActivityId() : NA_STRING); } if (null != document) { result.putAll(extractResult(document)); LOGGER.trace("Read result: {}", document); } return Status.OK; }
@Override public Status insert(String table, String key, Map<String, ByteIterator> values) { Document documentDefinition = getDocumentDefinition(key, values); ResourceResponse<Document> resourceResponse = null; long startTime = System.nanoTime(); try { if (this.useUpsert) { resourceResponse = AzureCosmosClient.client.upsertDocument(getDocumentCollectionLink(this.databaseName, table), documentDefinition, getRequestOptions(key), true); } else { resourceResponse = AzureCosmosClient.client.createDocument(getDocumentCollectionLink(this.databaseName, table), documentDefinition, getRequestOptions(key), true); } } catch (DocumentClientException e) { if (!this.includeExceptionStackInLog) { e = null; } LOGGER.error("Failed to insert key {} to collection {} in database {}", key, table, this.databaseName, e); return Status.ERROR; } finally { long elapsed = (System.nanoTime() - startTime) / NS_IN_US; LOGGER.debug("Inserted key {} in {}us - ActivityID: {}", key, elapsed, resourceResponse != null ? resourceResponse.getActivityId() : NA_STRING); } return Status.OK; }
private Document getDocumentDefinition(String key, Map<String, ByteIterator> values) { Document document = new Document(); document.set("id", key); if (!this.useHashQueryForScan) { // This field is only needed for range scans. // Even if this field is present in the document you // should still partition on id for simplicity of config. document.set("myid", key); } for (Entry<String, ByteIterator> entry : values.entrySet()) { document.set(entry.getKey(), entry.getValue().toString()); } return document; } }
@Override public Status update(String table, String key, Map<String, ByteIterator> values) { Document record = getDocumentById(table, key); if (record == null) { return Status.ERROR; } // Update each field. for (Entry<String, ByteIterator> val : values.entrySet()) { record.set(val.getKey(), val.getValue().toString()); } // Replace the document. try { documentClient.replaceDocument(record, null); } catch (DocumentClientException e) { e.printStackTrace(System.err); return Status.ERROR; } return Status.OK; }
@Override public Status delete(String table, String key) { ResourceResponse<Document> deletedResource = null; long startTime = System.nanoTime(); try { deletedResource = AzureCosmosClient.client.deleteDocument(getDocumentLink(this.databaseName, table, key), getRequestOptions(key)); } catch (DocumentClientException e) { if (!this.includeExceptionStackInLog) { e = null; } LOGGER.error("Failed to delete key {} in collection {} in database {}", key, table, this.databaseName, e); return Status.ERROR; } finally { long elapsed = (System.nanoTime() - startTime) / NS_IN_US; LOGGER.debug("Deleted key {} in {}us - ActivityID: {}", key, elapsed, deletedResource != null ? deletedResource.getActivityId() : NA_STRING); } return Status.OK; }
@Override public Status read(String table, String key, Set<String> fields, Map<String, ByteIterator> result) { Document record = getDocumentById(table, key); if (record != null) { Set<String> fieldsToReturn = (fields == null ? record.getHashMap().keySet() : fields); for (String field : fieldsToReturn) { if (field.startsWith("_")) { continue; } result.put(field, new StringByteIterator(record.getString(field))); } return Status.OK; } // Unable to find the specidifed document. return Status.ERROR; }
private HashMap<String, ByteIterator> extractResult(Document item) { if (null == item) { return null; } HashMap<String, ByteIterator> rItems = new HashMap<>(item.getHashMap().size()); for (Entry<String, Object> attr : item.getHashMap().entrySet()) { LOGGER.trace("Result- key: {}, value: {}", attr.getKey(), attr.getValue().toString()); rItems.put(attr.getKey(), new StringByteIterator(attr.getValue().toString())); } return rItems; }
/** * Cleanup any state for this DB. Called once per DB instance; there is one DB * instance per client thread. */ @Override public void cleanup() throws DBException { if (INIT_COUNT.decrementAndGet() == 0) { try { AzureCosmosClient.client.close(); } catch (Exception e) { if (!this.includeExceptionStackInLog) { e = null; } LOGGER.error("Could not close DocumentClient", e); } finally { client = null; } } }