First commit

.github/actions/install-dynamodb/action.yml

@@ -0,0 +1,13 @@
+name: Install DynamoDB
+description: >
+  Install a local version of DynamoDB.
+  See https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.DownloadingAndRunning.html
+outputs:
+  dynamodb-dir:
+    description: Path to the directory containing the JAR file
+    value: ${{ github.action_path }}
+runs:
+  using: composite
+  steps:
+    - run: curl https://d1ni2b6xgvw0s0.cloudfront.net/v2.x/dynamodb_local_latest.tar.gz | tar -C ${{ github.action_path }} -x
+      shell: bash

.github/dependabot.yml

@@ -0,0 +1,12 @@
+version: 2
+
+updates:
+  - package-ecosystem: github-actions
+    directory: /
+    schedule:
+      interval: weekly
+
+  - package-ecosystem: pip
+    directory: /
+    schedule:
+      interval: weekly

.github/workflows/publish.yml

@@ -0,0 +1,12 @@
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - v*
+
+jobs:
+  publish:
+    uses: OpenAstronomy/github-actions-workflows/.github/workflows/publish_pure_python.yml@v1
+    secrets:
+      pypi_token: ${{ secrets.pypi_token }}

.github/workflows/pull_request.yml

@@ -0,0 +1,43 @@
+name: Pull Request
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  python:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Check for Lint
+        uses: chartboost/ruff-action@v1
+
+      - name: Check Formatting
+        uses: chartboost/ruff-action@v1
+        with:
+          args: format --check
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - id: install-dynamodb
+        name: Install DynamoDB
+        uses: ./.github/actions/install-dynamodb
+
+      - name: Pip Install
+        run: pip install -e .
+
+      - name: Pytest
+        run: pytest --dynamodb-dir=${{ steps.install-dynamodb.outputs.dynamodb-dir }}
+
+      - name: MyPy
+        run: mypy .
+
+      - name: Upload to Codecov.io
+        uses: codecov/codecov-action@v4

.gitignore

@@ -0,0 +1,8 @@
+*.egg-info
+__pycache__
+.coverage
+.tox
+.vscode
+build
+dist
+_version.py

.vscode/extensions.json

@@ -0,0 +1,5 @@
+{
+  "recommendations": [
+    "charliermarsh.ruff"
+  ]
+}

dynamodb_autoincrement/__init__.py

@@ -0,0 +1,189 @@
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+
+import boto3
+from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource
+
+from .types import DynamoDBItem
+
+
+TransactionCanceledException = boto3.client(
+    "dynamodb"
+).exceptions.TransactionCanceledException
+
+
+@dataclass(frozen=True)
+class BaseDynamoDBAutoIncrement:
+    dynamodb: DynamoDBServiceResource
+    counter_table_name: str
+    counter_table_key: DynamoDBItem
+    attribute_name: str
+    table_name: str
+    initial_value: int
+    dangerously: bool = False
+
+    def next(self, item: DynamoDBItem) -> tuple[Iterable[dict[str, Any]], str]:
+        raise NotImplementedError
+
+    def _put_item(self, *, TableName, **kwargs):
+        # FIXME: DynamoDB resource does not have put_item method; emulate it
+        self.dynamodb.Table(TableName).put_item(**kwargs)
+
+    def _get_item(self, *, TableName, **kwargs):
+        # FIXME: DynamoDB resource does not have get_item method; emulate it
+        return self.dynamodb.Table(TableName).get_item(**kwargs)
+
+    def _query(self, *, TableName, **kwargs):
+        return self.dynamodb.Table(TableName).query(**kwargs)
+
+    def put(self, item: DynamoDBItem):
+        while True:
+            puts, next_counter = self.next(item)
+            if self.dangerously:
+                for put in puts:
+                    self._put_item(**put)
+            else:
+                try:
+                    self.dynamodb.transact_write_items(
+                        TransactItems=[{"Put": put} for put in puts]
+                    )
+                except TransactionCanceledException:
+                    continue
+            return next_counter
+
+
+class DynamoDBAutoIncrement(BaseDynamoDBAutoIncrement):
+    def next(self, item):
+        counter = (
+            self._get_item(
+                AttributesToGet=[self.attribute_name],
+                Key=self.counter_table_key,
+                TableName=self.counter_table_name,
+            )
+            .get("Item", {})
+            .get(self.attribute_name)
+        )
+
+        if counter is None:
+            next_counter = self.initial_value
+            put_kwargs = {"ConditionExpression": "attribute_not_exists(#counter)"}
+        else:
+            next_counter = counter + 1
+            put_kwargs = {
+                "ConditionExpression": "#counter = :counter",
+                "ExpressionAttributeValues": {
+                    ":counter": counter,
+                },
+            }
+
+        puts = [
+            {
+                **put_kwargs,
+                "ExpressionAttributeNames": {
+                    "#counter": self.attribute_name,
+                },
+                "Item": {
+                    **self.counter_table_key,
+                    self.attribute_name: next_counter,
+                },
+                "TableName": self.counter_table_name,
+            },
+            {
+                "ConditionExpression": "attribute_not_exists(#counter)",
+                "ExpressionAttributeNames": {
+                    "#counter": self.attribute_name,
+                },
+                "Item": {self.attribute_name: next_counter, **item},
+                "TableName": self.table_name,
+            },
+        ]
+
+        return puts, next_counter
+
+
+class DynamoDBHistoryAutoIncrement(BaseDynamoDBAutoIncrement):
+    def list(self) -> list[int]:
+        result = self._query(
+            TableName=self.table_name,
+            ExpressionAttributeNames={
+                **{f"#{i}": key for i, key in enumerate(self.counter_table_key.keys())},
+                "#counter": self.attribute_name,
+            },
+            ExpressionAttributeValues={
+                f":{i}": value
+                for i, value in enumerate(self.counter_table_key.values())
+            },
+            KeyConditionExpression=" AND ".join(
+                f"#{i} = :{i}" for i in range(len(self.counter_table_key.keys()))
+            ),
+            ProjectionExpression="#counter",
+        )
+        return sorted(item[self.attribute_name] for item in result["Items"])
+
+    def get(self, version: Optional[int] = None) -> DynamoDBItem:
+        if version is None:
+            kwargs = {
+                "TableName": self.counter_table_name,
+                "Key": self.counter_table_key,
+            }
+        else:
+            kwargs = {
+                "TableName": self.table_name,
+                "Key": {**self.counter_table_key, self.attribute_name: version},
+            }
+        return self._get_item(**kwargs).get("Item")
+
+    def next(self, item):
+        existing_item = self._get_item(
+            TableName=self.counter_table_name,
+            Key=self.counter_table_key,
+        ).get("Item")
+
+        counter = (
+            None if existing_item is None else existing_item.get(self.attribute_name)
+        )
+
+        if counter is None:
+            next_counter = self.initial_value
+            put_kwargs = {"ConditionExpression": "attribute_not_exists(#counter)"}
+        else:
+            next_counter = counter + 1
+            put_kwargs = {
+                "ConditionExpression": "#counter = :counter",
+                "ExpressionAttributeValues": {
+                    ":counter": counter,
+                },
+            }
+
+        if existing_item is not None and counter is not None:
+            existing_item[self.attribute_name] = next_counter
+            next_counter += 1
+
+        puts = [
+            {
+                **put_kwargs,
+                "ExpressionAttributeNames": {
+                    "#counter": self.attribute_name,
+                },
+                "Item": {
+                    **item,
+                    **self.counter_table_key,
+                    self.attribute_name: next_counter,
+                },
+                "TableName": self.counter_table_name,
+            },
+        ]
+
+        if existing_item is not None:
+            puts.append(
+                {
+                    "ConditionExpression": "attribute_not_exists(#counter)",
+                    "ExpressionAttributeNames": {
+                        "#counter": self.attribute_name,
+                    },
+                    "Item": existing_item,
+                    "TableName": self.table_name,
+                }
+            )
+
+        return puts, next_counter

dynamodb_autoincrement/tests/conftest.py

@@ -0,0 +1,36 @@
+import pytest
+
+
+@pytest.fixture
+def create_tables(dynamodb):
+    for kwargs in [
+        {
+            "AttributeDefinitions": [
+                {"AttributeName": "tableName", "AttributeType": "S"}
+            ],
+            "BillingMode": "PAY_PER_REQUEST",
+            "KeySchema": [{"AttributeName": "tableName", "KeyType": "HASH"}],
+            "TableName": "autoincrement",
+        },
+        {
+            "BillingMode": "PAY_PER_REQUEST",
+            "AttributeDefinitions": [
+                {"AttributeName": "widgetID", "AttributeType": "N"}
+            ],
+            "KeySchema": [{"AttributeName": "widgetID", "KeyType": "HASH"}],
+            "TableName": "widgets",
+        },
+        {
+            "BillingMode": "PAY_PER_REQUEST",
+            "AttributeDefinitions": [
+                {"AttributeName": "widgetID", "AttributeType": "N"},
+                {"AttributeName": "version", "AttributeType": "N"},
+            ],
+            "KeySchema": [
+                {"AttributeName": "widgetID", "KeyType": "HASH"},
+                {"AttributeName": "version", "KeyType": "RANGE"},
+            ],
+            "TableName": "widgetHistory",
+        },
+    ]:
+        dynamodb.create_table(**kwargs)

dynamodb_autoincrement/tests/test_autoincrement.py

@@ -0,0 +1,18 @@
+from .. import DynamoDBAutoIncrement
+
+
+N = 20
+
+
+def test_dangerously_handles_many_serial_puts(dynamodb, create_tables):
+    autoincrement = DynamoDBAutoIncrement(
+        counter_table_name="autoincrement",
+        counter_table_key={"tableName": "widgets"},
+        table_name="widgets",
+        attribute_name="widgetID",
+        initial_value=1,
+        dynamodb=dynamodb,
+    )
+    ids = list(range(1, N + 1))
+    results = [autoincrement.put({"widgetName": id}) for id in ids]
+    assert sorted(results) == ids

dynamodb_autoincrement/types.py

@@ -0,0 +1,10 @@
+from decimal import Decimal
+from typing import Mapping, Optional, Sequence, Union
+
+PrimitiveDynamoDBValues = Optional[Union[str, int, float, Decimal, bool]]
+DynamoDBValues = Union[
+    PrimitiveDynamoDBValues,
+    Mapping[str, PrimitiveDynamoDBValues],
+    Sequence[PrimitiveDynamoDBValues],
+]
+DynamoDBItem = Mapping[str, DynamoDBValues]