@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); }
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 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 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 + "'", .createCollection(getDatabase().getSelfLink(), collectionDefinition, null) .getResource();
public <T> List<T> findAll(String collectionName, final Class<T> entityClass) { final List<DocumentCollection> collections = documentDbFactory.getDocumentClient(). queryCollections( getDatabaseLink(this.databaseName), new SqlQuerySpec("SELECT * FROM ROOT r WHERE r.id=@id", new SqlParameterCollection(new SqlParameter("@id", collectionName))), null) .getQueryIterable().toList(); if (collections.size() != 1) { throw new RuntimeException("expect only one collection: " + collectionName + " in database: " + this.databaseName + ", but found " + collections.size()); } final List<Document> results = documentDbFactory.getDocumentClient() .queryDocuments(collections.get(0).getSelfLink(), "SELECT * FROM c", null) .getQueryIterable().toList(); final List<T> entities = new ArrayList<>(); for (int i = 0; i < results.size(); i++) { final T entity = mappingDocumentDbConverter.read(entityClass, results.get(i)); entities.add(entity); } return entities; }
@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 FeedResponse<DocumentCollection> invoke() throws Exception { return client.queryCollections(databaseLink, query, options); } });
@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 <T> List<T> delete(Query query, Class<T> entityClass, String collectionName) { final SqlQuerySpec sqlQuerySpec = createSqlQuerySpec(query, entityClass); final Optional<Object> partitionKeyValue = getPartitionKeyValue(query, entityClass); final DocumentCollection collection = getDocCollection(collectionName); final FeedOptions feedOptions = new FeedOptions(); if (!partitionKeyValue.isPresent()) { feedOptions.setEnableCrossPartitionQuery(true); } final List<Document> results = documentDbFactory.getDocumentClient() .queryDocuments(collection.getSelfLink(), sqlQuerySpec, feedOptions).getQueryIterable().toList(); final RequestOptions options = new RequestOptions(); if (partitionKeyValue.isPresent()) { options.setPartitionKey(new PartitionKey(partitionKeyValue.get())); } final List<T> deletedResult = new ArrayList<>(); for (final Document document : results) { try { documentDbFactory.getDocumentClient().deleteDocument((document).getSelfLink(), options); deletedResult.add(getConverter().read(entityClass, document)); } catch (DocumentClientException e) { throw new DocumentDBAccessException( String.format("Failed to delete document [%s]", (document).getSelfLink()), e); } } return deletedResult; }
this.useHashQueryForScan, this.queryText); AzureCosmosClient.client = new DocumentClient(uri, primaryKey, connectionPolicy, consistencyLevel); LOGGER.info("Azure Cosmos connection created: {}", uri); } catch (IllegalArgumentException e) { AzureCosmosClient.client.readDatabase(getDatabaseLink(this.databaseName), new RequestOptions()); } catch (DocumentClientException e) { throw new DBException("Invalid database name (" + this.databaseName + ") or failed to read database.", e);
@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 FeedResponse<Database> invoke() throws Exception { return client.queryDatabases(querySpec, options); } });
@Override public ResourceResponse<DocumentCollection> invoke() throws Exception { return client.createCollection(databaseLink, collection, options); } });
.readDocument(getDocumentLink(this.databaseName, collectionName, id), null) .getResource(); documentDbFactory.getDocumentClient().replaceDocument( originalDoc.getSelfLink(), originalDoc,
@Override public ResourceResponse<DocumentCollection> invoke() throws Exception { return client.deleteCollection(collectionLink, options); } });
@Override public ResourceResponse<Database> invoke() throws Exception { return client.createDatabase(database, options); } });
private Database createDatabaseIfNotExists(String dbName) { try { final List<Database> dbList = documentDbFactory.getDocumentClient() .queryDatabases(new SqlQuerySpec("SELECT * FROM root r WHERE r.id=@id", new SqlParameterCollection(new SqlParameter("@id", dbName))), null) .getQueryIterable().toList(); if (!dbList.isEmpty()) { return dbList.get(0); } else { // create new database final Database db = new Database(); db.setId(dbName); if (LOGGER.isDebugEnabled()) { LOGGER.debug("execute createDatabase {}", dbName); } final Resource resource = documentDbFactory.getDocumentClient() .createDatabase(db, null).getResource(); if (resource instanceof Database) { return (Database) resource; } else { LOGGER.error("create database {} get unexpected result: {}" + resource.getSelfLink()); throw new RuntimeException("create database {} get unexpected result: " + resource.getSelfLink()); } } } catch (DocumentClientException ex) { throw new RuntimeException("createOrGetDatabase exception", ex); } }
public <T> List<T> findAll(String collectionName, final Class<T> entityClass) { final List<DocumentCollection> collections = documentDbFactory.getDocumentClient(). queryCollections( getDatabaseLink(this.databaseName), new SqlQuerySpec("SELECT * FROM ROOT r WHERE r.id=@id", new SqlParameterCollection(new SqlParameter("@id", collectionName))), null) .getQueryIterable().toList(); if (collections.size() != 1) { throw new IllegalCollectionException("expect only one collection: " + collectionName + " in database: " + this.databaseName + ", but found " + collections.size()); } final FeedOptions feedOptions = new FeedOptions(); feedOptions.setEnableCrossPartitionQuery(true); final SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("SELECT * FROM root c"); final List<Document> results = documentDbFactory.getDocumentClient() .queryDocuments(collections.get(0).getSelfLink(), sqlQuerySpec, feedOptions) .getQueryIterable().toList(); final List<T> entities = new ArrayList<>(); for (int i = 0; i < results.size(); i++) { final T entity = mappingDocumentDbConverter.read(entityClass, results.get(i)); entities.add(entity); } return entities; }
@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; }