diff --git a/db.go b/db.go index b9fa812..511e297 100644 --- a/db.go +++ b/db.go @@ -5,6 +5,7 @@ import ( "context" "errors" "fmt" + "sync" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/dynamodb" @@ -17,6 +18,8 @@ import ( // DB is a DynamoDB client. type DB struct { client dynamodbiface.DynamoDBAPI + // table description cache for LEK inference + descs *sync.Map // table name → Description } // New creates a new client with the given configuration. @@ -32,6 +35,7 @@ func New(cfg aws.Config, options ...func(*dynamodb.Options)) *DB { func NewFromIface(client dynamodbiface.DynamoDBAPI) *DB { db := &DB{ client: client, + descs: new(sync.Map), } return db } @@ -41,6 +45,17 @@ func (db *DB) Client() dynamodbiface.DynamoDBAPI { return db.client } +func (db *DB) loadDesc(name string) (desc Description, ok bool) { + if descv, exists := db.descs.Load(name); exists { + desc, ok = descv.(Description) + } + return +} + +func (db *DB) storeDesc(desc Description) { + db.descs.Store(desc.Name, desc) +} + // ListTables is a request to list tables. // See: http://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_ListTables.html type ListTables struct { diff --git a/describetable.go b/describetable.go index 4c295ca..afbdacf 100644 --- a/describetable.go +++ b/describetable.go @@ -268,7 +268,7 @@ func (dt *DescribeTable) Run(ctx context.Context) (Description, error) { } desc := newDescription(result.Table) - dt.table.desc.Store(desc) + dt.table.db.storeDesc(desc) return desc, nil } diff --git a/scan_test.go b/scan_test.go index de2dcc9..c25622b 100644 --- a/scan_test.go +++ b/scan_test.go @@ -3,6 +3,7 @@ package dynamo import ( "context" "reflect" + "sync" "testing" "time" ) @@ -176,7 +177,12 @@ func TestScanMagicLEK(t *testing.T) { if testDB == nil { t.Skip(offlineSkipMsg) } - table := testDB.Table(testTableWidgets) + + testDB0 := *testDB + testDB0.descs = new(sync.Map) + freshTestDB := &testDB0 + + table := freshTestDB.Table(testTableWidgets) ctx := context.Background() widgets := []interface{}{ diff --git a/table.go b/table.go index 5faffdb..056e991 100644 --- a/table.go +++ b/table.go @@ -3,7 +3,6 @@ package dynamo import ( "context" "fmt" - "sync/atomic" "time" "github.com/aws/aws-sdk-go-v2/service/dynamodb" @@ -33,8 +32,6 @@ const ( type Table struct { name string db *DB - // desc is this table's cached description, used for inferring keys - desc *atomic.Value // Description } // Table returns a Table handle specified by name. @@ -42,7 +39,6 @@ func (db *DB) Table(name string) Table { return Table{ name: name, db: db, - desc: new(atomic.Value), } } @@ -119,7 +115,8 @@ func (table Table) primaryKeys(ctx context.Context, lek, esk Item, index string) // now we're forced to call DescribeTable // do we have a description cached? - if desc, ok := table.desc.Load().(Description); ok { + + if desc, ok := table.db.loadDesc(table.name); ok { keys := desc.keys(index) if keys != nil { return keys, nil