diff --git a/README.md b/README.md
index cedd727..6af9466 100644
--- a/README.md
+++ b/README.md
@@ -16,8 +16,8 @@ export interface UserDoc {
   name: string;
 }
 
-export class UserDataSource extends CosmosDataSource<UserDoc, ApolloContext> {}
-export class PostDataSource extends CosmosDataSource<PostDoc, ApolloContext> {}
+export class UserDataSource extends CosmosDataSource<UserDoc> {}
+export class PostDataSource extends CosmosDataSource<PostDoc> {}
 ```
 
 `server.ts`
@@ -45,6 +45,20 @@ const server = new ApolloServer({
 });
 ```
 
+## Context
+
+It is often useful to define a context. See [Apollo docs on context](https://www.apollographql.com/docs/apollo-server/data/context/) To make this strongly typed, there is a second type paramater on the CosmosDbDataSource:
+
+```typescript
+interface MyQueryContext {
+  currentUserId: string
+}
+
+/////
+
+const userDataSource extends CosmosDataSource<UserDoc, MyQueryContext> {}
+```
+
 ## Custom Queries
 
 CosmosDataSource exposes a `findManyByQuery` method that accepts a ComosDB SQL query either as a string or a `SqlQuerySpec` object containing the query and a parameter collection. This can be used directly in the resolvers, but probably better to create wrappers that hide the query details.
diff --git a/src/datasource.ts b/src/datasource.ts
index c46a150..dafaf09 100644
--- a/src/datasource.ts
+++ b/src/datasource.ts
@@ -7,6 +7,7 @@ import { Logger } from "./helpers";
 import { isCosmosDbContainer } from "./helpers";
 import { createCachingMethods, CachedMethods, FindArgs } from "./cache";
 
+
 export interface CosmosDataSourceOptions {
   logger?: Logger;
 }
@@ -24,7 +25,7 @@ export interface CosmosQueryDbArgs {
 
 export type QueryFindArgs = FindArgs & CosmosQueryDbArgs;
 
-export class CosmosDataSource<TData extends { id: string }, TContext>
+export class CosmosDataSource<TData extends { id: string }, TContext = null>
   extends DataSource<TContext>
   implements CachedMethods<TData> {
   container: Container;
diff --git a/tests/schema.test.ts b/tests/schema.test.ts
new file mode 100644
index 0000000..72a0255
--- /dev/null
+++ b/tests/schema.test.ts
@@ -0,0 +1,19 @@
+/* eslint @typescript-eslint/no-unused-vars: "off" */
+
+import { CosmosDataSource } from "../src/datasource";
+
+interface UserDoc {
+  id: string;
+  email: string;
+  partitionKey?: string;
+}
+
+interface Context {
+  something: string;
+}
+
+// normal context specified
+class UserDataSource1 extends CosmosDataSource<UserDoc, Context> {}
+
+// no context specified
+class UserDataSource2 extends CosmosDataSource<UserDoc> {}
\ No newline at end of file