move cli

http-api/cli/seeds.py

@@ -0,0 +1,109 @@
+from typing import Any, Generator
+
+import layercake.jsonl as jsonl
+from elasticsearch import Elasticsearch
+from layercake.dynamodb import deserialize
+from tqdm import tqdm
+
+from boto3clients import dynamodb_client
+
+elastic_client = Elasticsearch('http://127.0.0.1:9200')
+jsonl_files = (
+    'test-orders.jsonl',
+    'test-users.jsonl',
+    'test-enrollments.jsonl',
+    'test-courses.jsonl',
+)
+
+
+def put_item(item: dict, table_name: str, /, dynamodb_client) -> bool:
+    try:
+        dynamodb_client.put_item(
+            TableName=table_name,
+            Item=item,
+        )
+    except Exception:
+        return False
+    else:
+        return True
+
+
+def scan_table(table_name: str, /, dynamodb_client, **kwargs) -> Generator:
+    try:
+        r = dynamodb_client.scan(TableName=table_name, **kwargs)
+    except Exception:
+        yield from ()
+    else:
+        for item in r['Items']:
+            yield deserialize(item)
+
+        if 'LastEvaluatedKey' in r:
+            yield from scan_table(
+                table_name,
+                dynamodb_client=dynamodb_client,
+                ExclusiveStartKey=r['LastEvaluatedKey'],
+            )
+
+
+class Elastic:
+    def __init__(self, client: Elasticsearch) -> None:
+        self.client = client
+
+    def index_item(
+        self,
+        id: str,
+        index: str,
+        doc: dict,
+    ):
+        return self.client.index(
+            index=index,
+            id=id,
+            document=_serialize_python_type(doc),
+        )
+
+    def delete_index(self, index: str) -> bool:
+        try:
+            self.client.indices.delete(index=index)
+        except Exception:
+            return False
+        else:
+            return True
+
+
+def _serialize_python_type(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {k: _serialize_python_type(v) for k, v in value.items()}
+
+    if isinstance(value, set):
+        return list(value)
+
+    if isinstance(value, list):
+        return [_serialize_python_type(v) for v in value]
+
+    return value
+
+
+if __name__ == '__main__':
+    elastic = Elastic(elastic_client)
+
+    for file in tqdm(jsonl_files, desc='Processing files'):
+        with jsonl.readlines(f'seeds/{file}') as lines:
+            table_name = file.removesuffix('.jsonl')
+
+            for line in tqdm(lines, desc=f'Processing lines in {file}'):
+                put_item(line, table_name, dynamodb_client)
+
+    for file in tqdm(jsonl_files, desc='Scanning tables'):
+        table_name = file.removesuffix('.jsonl')
+        elastic.delete_index(table_name)
+
+        for record in tqdm(
+            scan_table(
+                table_name,
+                dynamodb_client,
+                FilterExpression='sk = :sk',
+                ExpressionAttributeValues={':sk': {'S': '0'}},
+            ),
+            desc=f'Indexing {table_name}',
+        ):
+            elastic.index_item(id=record['id'], index=table_name, doc=record)