Update sample code to use new API

Also remove use of TypedDatastore