Scala DynamoDB类代码示例

//设置package包名称以及导入依赖的类
package org.zalando.react.nakadi.commit.handlers.aws

import com.amazonaws.regions.Regions
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
import com.amazonaws.services.dynamodbv2.document.DynamoDB
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient
import org.zalando.react.nakadi.properties.CommitProperties


trait Provider {
  def client: DynamoDB
  def leaseProperties: CommitProperties
}

class ClientProvider(override val leaseProperties: CommitProperties) extends Provider {

  private val credentialsProviderChain = new DefaultAWSCredentialsProviderChain()
  private val region = Regions.fromName(leaseProperties.awsCommitRegion)

  override val client: DynamoDB = {
    val c = new AmazonDynamoDBClient(credentialsProviderChain)
    c.configureRegion(region)
    new DynamoDB(c)
  }
} 

//设置package包名称以及导入依赖的类
package common

import com.amazonaws.services.dynamodbv2.document.{ DynamoDB, Table }
import com.amazonaws.services.dynamodbv2.model._
import com.amazonaws.services.dynamodbv2.{ AmazonDynamoDB, AmazonDynamoDBClient }

import scala.collection.JavaConversions._


trait DynamoDBTableHelper {

  def createTableIfNotExist(describe: DescribeTableRequest, create: CreateTableRequest): Unit = try {
    dynamoDB.describeTable(describe)
  } catch {
    case e: ResourceNotFoundException => // ???????????????
      dynamoDB.createTable(create)
  }

  def deleteContents(scan: ScanRequest, deleteItem: Map[String, AttributeValue] => DeleteItemRequest): Unit =
    dynamoDB.scan(scan).getItems.toList.foreach { kv => // ??????????????
      dynamoDB.deleteItem(deleteItem(kv.toMap))
    }

  def deleteTable(): Unit = table.delete()

  val db = new DynamoDB(dynamoDB)

  def table: Table

  private def dynamoDB: AmazonDynamoDB = {
    val client = new AmazonDynamoDBClient()
    client.setEndpoint("http://localhost:8000")
    client
  }
} 

//设置package包名称以及导入依赖的类
package dynamite

import com.amazonaws.services.dynamodbv2.model.DeleteTableRequest
import com.amazonaws.services.dynamodbv2.util.TableUtils
import com.amazonaws.services.dynamodbv2.document.DynamoDB
import org.scalatest._
import scala.concurrent.duration._
import com.amazonaws.auth.{ BasicAWSCredentials, AWSStaticCredentialsProvider }

trait DynamoTestClient {
  val dynamoPortKey = "dynamodb.local.port"

  val dynamoPort = sys.props.get(dynamoPortKey).getOrElse {
    throw new Exception(s"Failed to find $dynamoPortKey")
  }

  val credentials = new AWSStaticCredentialsProvider(new BasicAWSCredentials("", ""))
  lazy val client = Repl.dynamoClient(Some(s"http://127.0.0.1:$dynamoPort"), Some(credentials))
}

trait DynamoSpec
    extends BeforeAndAfterAll
    with BeforeAndAfterEach
    with DynamoTestClient { self: Suite =>

  def tableNames: Seq[String]

  lazy val dynamo = new DynamoDB(client)

  override def afterEach() = {
    super.afterEach()
    tableNames.foreach { tableName =>
      TableUtils.deleteTableIfExists(client, new DeleteTableRequest().withTableName(tableName))
    }
  }

  override def afterAll() = {
    super.afterAll()
    client.shutdown()
  }

} 

//设置package包名称以及导入依赖的类
package infrastructures

import com.amazonaws.regions.Regions
import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec
import com.amazonaws.services.dynamodbv2.document.utils.{ NameMap, ValueMap }
import com.amazonaws.services.dynamodbv2.document.{ DynamoDB, Item, PrimaryKey }
import domains.{ Hero, Heroes }
import infrastructures.HeroDBClient.{ AttributeName, HeroConverter }

import scala.collection.JavaConversions._

class HeroDBClient {

  private val db = new DynamoDB(Regions.AP_NORTHEAST_1)
  private val table = db.getTable("hero")

  def findAll: Heroes =
    Heroes(table.scan().toList.map(HeroConverter.toHero)).sort

  def create(hero: Hero): Unit =
    table.putItem(
      new Item()
          .withPrimaryKey(AttributeName.Id, hero.id)
          .withString(AttributeName.Name, hero.name)
    )

  def update(hero: Hero): Unit = {
    val updateSpec = new UpdateItemSpec()
        .withPrimaryKey(AttributeName.Id, hero.id)
        .withUpdateExpression(s"SET #${ AttributeName.Name } = :v_name")
        .withNameMap(new NameMap().`with`("#" + { AttributeName.Name }, AttributeName.Name))
        .withValueMap(new ValueMap()
            .withString(":v_name", hero.name)
        )
    table.updateItem(updateSpec)
  }

  def delete(id: Int): Unit =
    table.deleteItem(new PrimaryKey("id", Int.box(id)))
}

object HeroDBClient {

  object AttributeName {
    val Id = "id"
    val Name = "name"
  }

  object HeroConverter {
    def toHero(item: Item): Hero =
      Hero(id = item.getInt(AttributeName.Id),
        name = item.getString(AttributeName.Name)
      )
  }

} 

//设置package包名称以及导入依赖的类
package org.zalando.znap.persistence.dynamo

import com.amazonaws.services.dynamodbv2.document.{DynamoDB, TableWriteItems}
import org.slf4j.LoggerFactory

trait DynamoDBWriter {
  private val logger = LoggerFactory.getLogger(classOf[DynamoDBWriter])

  protected val dynamoDB: DynamoDB

  protected def writeWithRetries(writeItems: TableWriteItems): Unit = {
    var outcome = dynamoDB.batchWriteItem(writeItems)

    do {
      // Check for unprocessed keys which could happen if you exceed provisioned throughput

      val unprocessedItems = outcome.getUnprocessedItems

      if (outcome.getUnprocessedItems.size() == 0) {
        //        println("No unprocessed items found")
      } else {
        logger.debug("Retrieving the unprocessed items")
        outcome = dynamoDB.batchWriteItemUnprocessed(unprocessedItems)
      }
    } while (outcome.getUnprocessedItems.size() > 0)
  }
} 

//设置package包名称以及导入依赖的类
package com.gu.cm

import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
import com.amazonaws.regions.RegionUtils
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient
import com.amazonaws.services.dynamodbv2.document.{PrimaryKey, DynamoDB}
import com.typesafe.config.{ConfigFactory, Config}
import scala.util.{Failure, Success, Try}

class DynamoDbConfigurationSource(dynamoDb: DynamoDB, identity: Identity, prefix: String) extends ConfigurationSource {

  val tableName = s"$prefix${identity.stack}"

  override def load: Config = {
    val config = for {
      table <- Try(dynamoDb.getTable(tableName))
      item <- Try(table.getItem(new PrimaryKey("App", identity.app, "Stage", identity.stage)))
    } yield {
      ConfigFactory.parseMap(item.getMap("Config"), s"Dynamo DB table $tableName [App=${identity.app}, Stage=${identity.stage}]")
    }

    config match {
      case Success(theConfig) => theConfig
      case Failure(theFailure) => ConfigFactory.empty(s"no DynamoDB config (or failed to load) for $tableName [App=${identity.app}, Stage=${identity.stage}], exception=[$theFailure]")
    }
  }
}

object DynamoDbConfigurationSource {
  def apply(identity: Identity, prefix: String = "config-"): DynamoDbConfigurationSource = {
    val dynamoDb = {
      val client = new AmazonDynamoDBClient(new DefaultAWSCredentialsProviderChain())
      client.setRegion(RegionUtils.getRegion(identity.region))
      new DynamoDB(client)
    }
    new DynamoDbConfigurationSource(dynamoDb, identity, prefix)
  }
} 

//设置package包名称以及导入依赖的类
package com.gu.cm

import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient
import com.amazonaws.services.dynamodbv2.document.{Item, DynamoDB}
import com.amazonaws.services.dynamodbv2.model._
import com.typesafe.config.ConfigFactory
import org.specs2.mutable.Specification
import org.specs2.specification.Scope

import scala.collection.JavaConverters._

class DynamoDbConfigurationSourceSpec extends Specification {
  "a dynamo DB configuration Source" should {
    "load a config from dynamodb" in new DynamoDbScope {
      val conf = dynamoDbConfigurationSource.load

      conf.getInt("foo.a") shouldEqual 42
    }
  }

  trait DynamoDbScope extends Scope {
    val dynamoDb = {
      val client = new AmazonDynamoDBClient(new DefaultAWSCredentialsProviderChain())
      client.setEndpoint("http://localhost:8000")
      new DynamoDB(client)
    }

    val identity = AwsApplication(
      stack = "test-stack",
      app = "configuration-magic",
      stage = "test",
      region = "eu-west-1"
    )

    val dynamoDbConfigurationSource = new DynamoDbConfigurationSource(dynamoDb, identity, "prefix-")

    def initTable() = {
      val req = new CreateTableRequest(
        "prefix-test-stack",
        List(
          new KeySchemaElement("App", KeyType.HASH),
          new KeySchemaElement("Stage", KeyType.RANGE)
        ).asJava).withAttributeDefinitions(
        new AttributeDefinition("App", "S"),
        new AttributeDefinition("Stage", "S")
      ).withProvisionedThroughput(new ProvisionedThroughput(1000L, 1000L))

      val table = dynamoDb.createTable(req)

      val confBlob = ConfigFactory.parseResourcesAnySyntax("sample").root.unwrapped()

      val item = new Item()
        .withString("App", "configuration-magic")
        .withString("Stage", "test")
        .withMap("Config", confBlob)
      table.putItem(item)
    }
    initTable()
  }
} 

//设置package包名称以及导入依赖的类
package publisher

import scala.util.Try
import com.amazonaws.services.dynamodbv2.document.DynamoDB
import com.amazonaws.services.dynamodbv2.document.Expected
import com.amazonaws.services.dynamodbv2.document.Item
import com.amazonaws.services.dynamodbv2.document.AttributeUpdate
import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec
import scala.collection.JavaConversions._
import publisher.QueueSource._

@Deprecated
class DynamoDBSessionRepository(dynamoDB: DynamoDB) extends SessionRepository {
  
  val table = dynamoDB.getTable(TableName)
  
  override def createSession(session: ChannelSession): Try[_] = {
    val createNew = Try(table.putItem(new Item()
        .withPrimaryKey(TableAttrs.channel, session.channel)
        .withStringSet(TableAttrs.sessions, session.id), 
          new Expected(TableAttrs.channel).ne(session.channel)))
    if(createNew.isFailure){
       Try(table.updateItem(new UpdateItemSpec()
            .withPrimaryKey(TableAttrs.channel, session.channel)
            .withAttributeUpdate(new AttributeUpdate(TableAttrs.sessions)
            .addElements(session.id))))
    }else createNew
  }
  
  override def deleteSession(session: ChannelSession): Unit = {
    table.updateItem(new UpdateItemSpec()
        .withPrimaryKey(TableAttrs.channel, session.channel)
        .withAttributeUpdate(new AttributeUpdate(TableAttrs.sessions)
        .removeElements(session.id)))
  }
  
  override def getSessions(channel: String): Seq[ChannelSession] = {
    Option(table.getItem(TableAttrs.channel, channel)).flatMap { item =>
      Option(item.getStringSet(TableAttrs.sessions)).map { sessions =>
        sessions.map{ sessionId => ChannelSession(channel, sessionId) }.toSeq
      }
    }.getOrElse(Nil)
  }
  
} 

//设置package包名称以及导入依赖的类
import com.amazonaws.auth.{AWSCredentials, AWSCredentialsProvider, BasicAWSCredentials}
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient
import com.amazonaws.services.dynamodbv2.document.DynamoDB
import com.amazonaws.services.dynamodbv2.model._

object DynamoDBTableLoader {

  val awsCredProvider = new AWSCredentialsProvider {
    override def refresh(): Unit = ()

    override def getCredentials: AWSCredentials = new BasicAWSCredentials("accessKey", "secretKey")
  }

  val dynamoDBClient = new AmazonDynamoDBClient(awsCredProvider.getCredentials)
  dynamoDBClient.setEndpoint("http://localhost:8000")
  val dynamoDB = new DynamoDB(dynamoDBClient)
  dynamoDB.shutdown()


  //Helper functions
  private val keySchemaType = (name: String, keyType: KeyType) => new KeySchemaElement().withAttributeName(name).withKeyType(keyType)
  private val keySchema = (name: String) => keySchemaType(name, KeyType.HASH)
  private val attributeDefType = (name: String, attrType: ScalarAttributeType) => new AttributeDefinition().withAttributeName(name).withAttributeType(attrType)
  private val attributeDef = (name: String) => attributeDefType(name, ScalarAttributeType.S)
  private val provisionedThroughput = (throughput: Long) => new ProvisionedThroughput().withReadCapacityUnits(throughput).withWriteCapacityUnits(throughput)
  private val globalSecondaryIndex = (indexName: String, keySchema: Seq[KeySchemaElement], throughput: ProvisionedThroughput) =>
    new GlobalSecondaryIndex()
      .withIndexName(indexName)
      .withKeySchema(keySchema: _*)
      .withProjection(new Projection().withProjectionType(ProjectionType.ALL))
      .withProvisionedThroughput(throughput)

  def initTestTable() = {
    new CreateTableRequest()
        .withTableName("test_table")
        .withKeySchema(keySchema("my_key"))

  }
} 

//设置package包名称以及导入依赖的类
package jp.ne.khddmks

import com.amazonaws.regions.Regions
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder
import com.amazonaws.services.dynamodbv2.document.{DynamoDB}
import com.amazonaws.services.lambda.runtime.Context


class Main {
  def handler(param: String, context: Context): String = {
    def createDynamoDB = {
      val dbClient = AmazonDynamoDBClientBuilder.standard
        .withRegion(Regions.AP_NORTHEAST_1)
        .build
      val db = new DynamoDB(dbClient)
      db
    }
    def readFromDynamoDb(db: DynamoDB) = {
      val table = db.getTable("sample001")
      val json = table.getItem("key", "config").getJSON("values")
      json
    }

    readFromDynamoDb(createDynamoDB)
  }
} 

//设置package包名称以及导入依赖的类
import com.amazonaws.services.dynamodbv2.document.{Item, Table, DynamoDB}
import com.amazonaws.services.dynamodbv2.{AmazonDynamoDB, AmazonDynamoDBClient}
import org.junit.runner.RunWith
import org.specs2.mutable.Specification
import org.specs2.runner.JUnitRunner
import play.api.{Logger, Play}
import play.api.test.WithApplication

import scala.xml.XML

@RunWith(classOf[JUnitRunner])
class DataLoaderSpec extends Specification{

  "DataLoader" should {
    "load data" in new WithApplication {
      val client = new AmazonDynamoDBClient()
        .withEndpoint("http://localhost:8000").asInstanceOf[AmazonDynamoDB]

      val dynamoDB = new DynamoDB(client)

      val table = (dynamoDB.getTable("SegmentEvents")).asInstanceOf[Table]


      val pid = "p1234564"
      val version_pid = "b1234564"
      val segment_pid = "p5432104"


      val doc = Play.getFile("public/segment_event/segment_event_test.xml")

      val xml = XML.loadFile(doc)
      val document = xml.toString()
      //val document = "<segment_event/>"
      Logger.info(document)
      val item = new Item()
        .withPrimaryKey("version_pid", version_pid, "pid", pid)
        .withString("segment_id", segment_pid)
        .withString("document", document)

      try {
        Logger.info("Attempting to load data; please wait...")
        table.putItem(item)

        Logger.info("Success.... ")

      } catch {
        case e: Exception => {
          Logger.info("Unable to add document: ")
          Logger.info(e.getMessage())
        }

      }

    }
  }

}