3
3
4
4
using System ;
5
5
using System . Threading . Tasks ;
6
- using Microsoft . Azure . Documents ;
7
- using Microsoft . Azure . Documents . Client ;
6
+ using Microsoft . Azure . Cosmos ;
8
7
9
8
namespace Azure . Functions . PowerShell . Tests . E2E
10
9
{
11
- public class TestDocument
12
- {
13
- public string id { get ; set ; }
14
- public string name { get ; set ; }
15
- }
16
10
17
11
public static class CosmosDBHelpers
18
12
{
19
- private static DocumentClient _docDbClient ;
20
- private static Uri inputCollectionsUri = UriFactory . CreateDocumentCollectionUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName ) ;
21
- private static Uri outputCollectionsUri = UriFactory . CreateDocumentCollectionUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName ) ;
22
- private static Uri leasesCollectionsUri = UriFactory . CreateDocumentCollectionUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . LeaseCollectionName ) ;
13
+ private static CosmosClient _cosmosDbClient ;
14
+
15
+ private class Document
16
+ {
17
+ // lower-case because Cosmos expects a field with this name
18
+ public string id { get ; set ; }
19
+ }
23
20
24
21
static CosmosDBHelpers ( )
25
22
{
26
23
var builder = new System . Data . Common . DbConnectionStringBuilder ( ) ;
27
24
builder . ConnectionString = Constants . CosmosDB . CosmosDBConnectionStringSetting ;
28
- var serviceUri = new Uri ( builder [ "AccountEndpoint" ] . ToString ( ) ) ;
29
- _docDbClient = new DocumentClient ( serviceUri , builder [ "AccountKey" ] . ToString ( ) ) ;
25
+ var serviceUri = builder [ "AccountEndpoint" ] . ToString ( ) ;
26
+ _cosmosDbClient = new CosmosClient ( serviceUri , builder [ "AccountKey" ] . ToString ( ) ) ;
30
27
}
31
28
32
29
// keep
33
30
public async static Task CreateDocument ( string docId )
34
31
{
35
32
Document documentToTest = new Document ( )
36
33
{
37
- Id = docId
34
+ id = docId
38
35
} ;
39
36
40
- Document insertedDoc = await _docDbClient . CreateDocumentAsync ( UriFactory . CreateDocumentCollectionUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName ) , documentToTest ) ;
41
- }
37
+ Container _inputContainer = _cosmosDbClient . GetContainer ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName ) ;
42
38
43
- public async static Task CreateDocument ( TestDocument testDocument )
44
- {
45
- Document insertedDoc = await _docDbClient . CreateDocumentAsync ( UriFactory . CreateDocumentCollectionUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName ) , testDocument ) ;
39
+ Document insertedDoc = await _inputContainer . CreateItemAsync < Document > ( documentToTest , new PartitionKey ( documentToTest . id ) ) ;
46
40
}
47
41
48
42
// keep
49
43
public async static Task < string > ReadDocument ( string docId )
50
44
{
51
- var docUri = UriFactory . CreateDocumentUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName , docId ) ;
52
45
Document retrievedDocument = null ;
53
46
await Utilities . RetryAsync ( async ( ) =>
54
47
{
55
48
try
56
49
{
57
- retrievedDocument = await _docDbClient . ReadDocumentAsync ( docUri ) ;
50
+ Container container = _cosmosDbClient . GetContainer ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName ) ;
51
+
52
+ retrievedDocument = await container . ReadItemAsync < Document > ( docId , new PartitionKey ( docId ) ) ;
58
53
return true ;
59
54
}
60
- catch ( DocumentClientException ex ) when ( ex . Error . Code == "NotFound" || ex . Error . Code == "Not Found" )
55
+ catch ( CosmosException ex ) when ( ex . StatusCode == System . Net . HttpStatusCode . NotFound )
61
56
{
62
57
return false ;
63
58
}
64
59
} , 120000 , 4000 ) ;
65
- return retrievedDocument . Id ;
60
+ return retrievedDocument . id ;
66
61
}
67
62
68
63
// keep
69
64
public async static Task DeleteTestDocuments ( string docId )
70
65
{
71
- var inputDocUri = UriFactory . CreateDocumentUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName , docId ) ;
72
- await DeleteDocument ( inputDocUri ) ;
73
- var outputDocUri = UriFactory . CreateDocumentUri ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName , docId ) ;
74
- await DeleteDocument ( outputDocUri ) ;
66
+ await DeleteDocument ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName , docId ) ;
67
+ await DeleteDocument ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName , docId ) ;
75
68
}
76
69
77
- private async static Task DeleteDocument ( Uri docUri )
70
+ private async static Task DeleteDocument ( string dbName , string collectionName , string docId )
78
71
{
79
72
try
80
73
{
81
- await _docDbClient . DeleteDocumentAsync ( docUri ) ;
74
+ Container container = _cosmosDbClient . GetContainer ( dbName , collectionName ) ;
75
+ await container . DeleteItemAsync < Document > ( docId , new PartitionKey ( docId ) ) ;
82
76
}
83
77
catch ( Exception )
84
78
{
@@ -89,41 +83,55 @@ private async static Task DeleteDocument(Uri docUri)
89
83
// keep
90
84
public async static Task CreateDocumentCollections ( )
91
85
{
92
- Database db = await _docDbClient . CreateDatabaseIfNotExistsAsync ( new Database { Id = Constants . CosmosDB . DbName } ) ;
93
- Uri dbUri = UriFactory . CreateDatabaseUri ( db . Id ) ;
94
-
95
- await CreateCollection ( dbUri , Constants . CosmosDB . InputCollectionName ) ;
96
- await CreateCollection ( dbUri , Constants . CosmosDB . OutputCollectionName ) ;
97
- await CreateCollection ( dbUri , Constants . CosmosDB . LeaseCollectionName ) ;
86
+ Database db = await _cosmosDbClient . CreateDatabaseIfNotExistsAsync ( Constants . CosmosDB . DbName ) ;
98
87
88
+ await CreateCollection ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName , "/id" ) ;
89
+ await CreateCollection ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName , "/id" ) ;
90
+ // While using extensions v2-3, the leases may not have a partition key, but the new SDK requires
91
+ // one to manually create a collection. This comment may be removed and this line uncommented when
92
+ // extension bundles for tests are updated.
93
+ //await CreateCollection(Constants.CosmosDB.DbName, Constants.CosmosDB.LeaseCollectionName, "/id");
99
94
}
100
95
public async static Task DeleteDocumentCollections ( )
101
96
{
102
- await DeleteCollection ( inputCollectionsUri ) ;
103
- await DeleteCollection ( outputCollectionsUri ) ;
104
- await DeleteCollection ( leasesCollectionsUri ) ;
97
+ await DeleteCollection ( Constants . CosmosDB . DbName , Constants . CosmosDB . InputCollectionName ) ;
98
+ await DeleteCollection ( Constants . CosmosDB . DbName , Constants . CosmosDB . OutputCollectionName ) ;
99
+ await DeleteCollection ( Constants . CosmosDB . DbName , Constants . CosmosDB . LeaseCollectionName ) ;
105
100
}
106
101
107
- private async static Task DeleteCollection ( Uri collectionUri )
102
+ private async static Task DeleteCollection ( string dbName , string collectionName )
108
103
{
109
104
try
110
105
{
111
- await _docDbClient . DeleteDocumentCollectionAsync ( collectionUri ) ;
106
+ Database database = _cosmosDbClient . GetDatabase ( dbName ) ;
107
+ await database . GetContainer ( collectionName ) . DeleteContainerAsync ( ) ;
112
108
}
113
109
catch ( Exception )
114
110
{
115
111
//Ignore
116
112
}
117
113
}
118
114
119
- private async static Task CreateCollection ( Uri dbUri , string collectioName )
115
+ private async static Task CreateCollection ( string dbName , string collectionName , string partitionKey )
120
116
{
121
- DocumentCollection collection = new DocumentCollection ( ) { Id = collectioName } ;
122
- await _docDbClient . CreateDocumentCollectionIfNotExistsAsync ( dbUri , collection ,
123
- new RequestOptions ( )
117
+ Database database = _cosmosDbClient . GetDatabase ( dbName ) ;
118
+ IndexingPolicy indexingPolicy = new IndexingPolicy
119
+ {
120
+ IndexingMode = IndexingMode . Consistent ,
121
+ Automatic = true ,
122
+ IncludedPaths =
124
123
{
125
- OfferThroughput = 400
126
- } ) ;
124
+ new IncludedPath
125
+ {
126
+ Path = "/*"
127
+ }
128
+ }
129
+ } ;
130
+ var containerProperties = new ContainerProperties ( collectionName , partitionKey )
131
+ {
132
+ IndexingPolicy = indexingPolicy
133
+ } ;
134
+ await database . CreateContainerIfNotExistsAsync ( containerProperties , 400 ) ;
127
135
}
128
136
}
129
137
}
0 commit comments