Merge pull request #19 from tinybirdco/add-dynamodb-connector-to-sdk

Add DynamoDB connector to Python SDK

Author: sandshoes

README.md

@@ -516,7 +516,13 @@ Use a local Tinybird container for development without affecting cloud workspace
 ### Connections
 
 ```python
-from tinybird_sdk import define_gcs_connection, define_kafka_connection, define_s3_connection, secret
+from tinybird_sdk import (
+    define_dynamodb_connection,
+    define_gcs_connection,
+    define_kafka_connection,
+    define_s3_connection,
+    secret,
+)
 
 events_kafka = define_kafka_connection(
     "events_kafka",
@@ -543,6 +549,42 @@ landing_gcs = define_gcs_connection(
         "service_account_credentials_json": secret("GCS_SERVICE_ACCOUNT_CREDENTIALS_JSON"),
     },
 )
+
+events_dynamodb = define_dynamodb_connection(
+    "events_dynamodb",
+    {
+        "region": "us-east-1",
+        "arn": secret("DYNAMODB_ROLE_ARN"),
+    },
+)
+```
+
+The DynamoDB connector uses an IAM role to authenticate. Reference the connection from
+a datasource to import a DynamoDB table:
+
+```python
+from tinybird_sdk import column, define_datasource, engine, t
+
+orders = define_datasource(
+    "orders",
+    {
+        "schema": {
+            "id": column(t.string(), {"json_path": "$.Item.id"}),
+            "_record": column(t.string(), {"json_path": "$.NewImage"}),
+            "_timestamp": column(t.date_time64(3), {"json_path": "$.ApproximateCreationDateTime"}),
+            "_event_name": column(t.string().low_cardinality(), {"json_path": "$.eventName"}),
+            "_is_deleted": column(t.uint8(), {"json_path": "$._is_deleted"}),
+        },
+        "engine": engine.replacing_merge_tree(
+            {"sorting_key": ["id"], "ver": "_timestamp", "is_deleted": "_is_deleted"}
+        ),
+        "dynamodb": {
+            "connection": events_dynamodb,
+            "table_arn": "arn:aws:dynamodb:us-east-1:123456789012:table/orders",
+            "export_bucket": "s3://my-tinybird-dynamodb-exports",
+        },
+    },
+)
 ```
 
 ### Datasources

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tinybird-sdk"
-version = "0.1.9"
+version = "0.2.0"
 description = "Python SDK for Tinybird Forward"
 readme = "README.md"
 license = "MIT"
@@ -9,7 +9,7 @@ authors = [
 ]
 requires-python = ">=3.11"
 dependencies = [
-    "tinybird==4.5.0",
+    "tinybird==4.6.0",
 ]
 classifiers = [
     "Development Status :: 4 - Beta",

src/tinybird_sdk/__init__.py

@@ -26,11 +26,16 @@
     "define_kafka_connection": ("tinybird_sdk.schema", "define_kafka_connection"),
     "define_s3_connection": ("tinybird_sdk.schema", "define_s3_connection"),
     "define_gcs_connection": ("tinybird_sdk.schema", "define_gcs_connection"),
+    "define_dynamodb_connection": ("tinybird_sdk.schema", "define_dynamodb_connection"),
     "get_connection_type": ("tinybird_sdk.schema", "get_connection_type"),
     "is_connection_definition": ("tinybird_sdk.schema", "is_connection_definition"),
     "is_kafka_connection_definition": ("tinybird_sdk.schema", "is_kafka_connection_definition"),
     "is_s3_connection_definition": ("tinybird_sdk.schema", "is_s3_connection_definition"),
     "is_gcs_connection_definition": ("tinybird_sdk.schema", "is_gcs_connection_definition"),
+    "is_dynamodb_connection_definition": (
+        "tinybird_sdk.schema",
+        "is_dynamodb_connection_definition",
+    ),
     "secret": ("tinybird_sdk.schema", "secret"),
     "define_token": ("tinybird_sdk.schema", "define_token"),
     "is_token_definition": ("tinybird_sdk.schema", "is_token_definition"),

src/tinybird_sdk/generator/connection.py

@@ -4,6 +4,7 @@
 
 from ..schema.connection import (
     ConnectionDefinition,
+    DynamoDBConnectionDefinition,
     GCSConnectionDefinition,
     KafkaConnectionDefinition,
     S3ConnectionDefinition,
@@ -70,6 +71,17 @@ def _generate_gcs_connection(connection: GCSConnectionDefinition) -> str:
     return "\n".join(parts)
 
 
+def _generate_dynamodb_connection(connection: DynamoDBConnectionDefinition) -> str:
+    options = connection.options
+    parts = [
+        "TYPE dynamodb",
+        f"DYNAMODB_ARN {options.arn}",
+        f"DYNAMODB_REGION {options.region}",
+    ]
+
+    return "\n".join(parts)
+
+
 def generate_connection(connection: ConnectionDefinition) -> GeneratedConnection:
     if isinstance(connection, KafkaConnectionDefinition):
         return GeneratedConnection(
@@ -83,6 +95,10 @@ def generate_connection(connection: ConnectionDefinition) -> GeneratedConnection
         return GeneratedConnection(
             name=connection._name, content=_generate_gcs_connection(connection)
         )
+    if isinstance(connection, DynamoDBConnectionDefinition):
+        return GeneratedConnection(
+            name=connection._name, content=_generate_dynamodb_connection(connection)
+        )
     raise ValueError(f"Unsupported connection type: {connection._connectionType}")
 
 

src/tinybird_sdk/generator/datasource.py

@@ -124,6 +124,15 @@ def _generate_import_config(import_config: Any) -> str:
     return "\n".join(lines)
 
 
+def _generate_dynamodb_config(dynamodb: Any) -> str:
+    lines = [
+        f"IMPORT_CONNECTION_NAME {dynamodb.connection._name}",
+        f"IMPORT_TABLE_ARN {dynamodb.table_arn}",
+        f"IMPORT_EXPORT_BUCKET {dynamodb.export_bucket}",
+    ]
+    return "\n".join(lines)
+
+
 def _generate_forward_query(forward_query: str | None) -> str | None:
     if not forward_query or not forward_query.strip():
         return None
@@ -193,6 +202,9 @@ def generate_datasource(datasource: DatasourceDefinition) -> GeneratedDatasource
     if datasource.options.gcs:
         parts.extend(["", _generate_import_config(datasource.options.gcs)])
 
+    if datasource.options.dynamodb:
+        parts.extend(["", _generate_dynamodb_config(datasource.options.dynamodb)])
+
     forward_query = _generate_forward_query(datasource.options.forward_query)
     if forward_query:
         parts.extend(["", forward_query])

src/tinybird_sdk/migrate/emit_ts.py

@@ -8,6 +8,7 @@
 from .parser_utils import parse_literal_from_datafile
 from .types import (
     DatasourceModel,
+    DynamoDBConnectionModel,
     GCSConnectionModel,
     KafkaConnectionModel,
     ParsedResource,
@@ -261,6 +262,14 @@ def _emit_datasource(ds: DatasourceModel) -> str:
             lines.append(f"        'from_timestamp': {_escape_string(ds.gcs.from_timestamp)},")
         lines.append("    },")
 
+    if ds.dynamodb:
+        connection_var = to_snake_case(ds.dynamodb.connection_name)
+        lines.append("    'dynamodb': {")
+        lines.append(f"        'connection': {connection_var},")
+        lines.append(f"        'table_arn': {_escape_string(ds.dynamodb.table_arn)},")
+        lines.append(f"        'export_bucket': {_escape_string(ds.dynamodb.export_bucket)},")
+        lines.append("    },")
+
     if ds.forward_query:
         lines.append("    'forward_query': '''")
         lines.append(ds.forward_query)
@@ -336,13 +345,31 @@ def _emit_gcs_connection(connection: GCSConnectionModel) -> str:
     return "\n".join(lines)
 
 
+def _emit_dynamodb_connection(connection: DynamoDBConnectionModel) -> str:
+    variable_name = to_snake_case(connection.name)
+    lines: list[str] = []
+    lines.append(
+        f"{variable_name} = define_dynamodb_connection({_escape_string(connection.name)}, {{"
+    )
+    lines.append(f"    'region': {_escape_string(connection.region)},")
+    lines.append(f"    'arn': {_escape_string(connection.arn)},")
+    lines.append("})")
+    lines.append("")
+    return "\n".join(lines)
+
+
 def _emit_connection(
-    connection: KafkaConnectionModel | S3ConnectionModel | GCSConnectionModel,
+    connection: KafkaConnectionModel
+    | S3ConnectionModel
+    | GCSConnectionModel
+    | DynamoDBConnectionModel,
 ) -> str:
     if isinstance(connection, S3ConnectionModel):
         return _emit_s3_connection(connection)
     if isinstance(connection, GCSConnectionModel):
         return _emit_gcs_connection(connection)
+    if isinstance(connection, DynamoDBConnectionModel):
+        return _emit_dynamodb_connection(connection)
     return _emit_kafka_connection(connection)
 
 
@@ -474,6 +501,8 @@ def emit_migration_file_content(resources: list[ParsedResource]) -> str:
             imports.add("define_s3_connection")
         elif isinstance(conn, GCSConnectionModel):
             imports.add("define_gcs_connection")
+        elif isinstance(conn, DynamoDBConnectionModel):
+            imports.add("define_dynamodb_connection")
     if needs_column:
         imports.add("column")
     if needs_params:

src/tinybird_sdk/migrate/parse_connection.py

@@ -10,6 +10,7 @@
 )
 from .types import (
     ConnectionModel,
+    DynamoDBConnectionModel,
     GCSConnectionModel,
     KafkaConnectionModel,
     ResourceFile,
@@ -30,6 +31,8 @@
     "S3_ACCESS_KEY",
     "S3_SECRET",
     "GCS_SERVICE_ACCOUNT_CREDENTIALS_JSON",
+    "DYNAMODB_ARN",
+    "DYNAMODB_REGION",
 }
 
 
@@ -58,6 +61,9 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
     access_secret: str | None = None
     service_account_credentials_json: str | None = None
 
+    dynamodb_arn: str | None = None
+    dynamodb_region: str | None = None
+
     i = 0
     while i < len(lines):
         raw_line = lines[i]
@@ -119,6 +125,10 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
             access_secret = parse_quoted_value(value)
         elif name == "GCS_SERVICE_ACCOUNT_CREDENTIALS_JSON":
             service_account_credentials_json = parse_quoted_value(value)
+        elif name == "DYNAMODB_ARN":
+            dynamodb_arn = parse_quoted_value(value)
+        elif name == "DYNAMODB_REGION":
+            dynamodb_region = parse_quoted_value(value)
         else:
             raise MigrationParseError(
                 resource.file_path,
@@ -135,12 +145,20 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
         )
 
     if connection_type == "kafka":
-        if region or arn or access_key or access_secret or service_account_credentials_json:
+        if (
+            region
+            or arn
+            or access_key
+            or access_secret
+            or service_account_credentials_json
+            or dynamodb_arn
+            or dynamodb_region
+        ):
             raise MigrationParseError(
                 resource.file_path,
                 "connection",
                 resource.name,
-                "S3/GCS directives are not valid for kafka connections.",
+                "S3/GCS/DynamoDB directives are not valid for kafka connections.",
             )
 
         if not bootstrap_servers:
@@ -175,12 +193,14 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
             or schema_registry_url
             or ssl_ca_pem
             or service_account_credentials_json
+            or dynamodb_arn
+            or dynamodb_region
         ):
             raise MigrationParseError(
                 resource.file_path,
                 "connection",
                 resource.name,
-                "Kafka/GCS directives are not valid for s3 connections.",
+                "Kafka/GCS/DynamoDB directives are not valid for s3 connections.",
             )
 
         if not region:
@@ -231,12 +251,14 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
             or arn
             or access_key
             or access_secret
+            or dynamodb_arn
+            or dynamodb_region
         ):
             raise MigrationParseError(
                 resource.file_path,
                 "connection",
                 resource.name,
-                "Kafka/S3 directives are not valid for gcs connections.",
+                "Kafka/S3/DynamoDB directives are not valid for gcs connections.",
             )
 
         if not service_account_credentials_json:
@@ -255,6 +277,53 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
             service_account_credentials_json=service_account_credentials_json,
         )
 
+    if connection_type == "dynamodb":
+        if (
+            bootstrap_servers
+            or security_protocol
+            or sasl_mechanism
+            or key
+            or secret
+            or schema_registry_url
+            or ssl_ca_pem
+            or region
+            or arn
+            or access_key
+            or access_secret
+            or service_account_credentials_json
+        ):
+            raise MigrationParseError(
+                resource.file_path,
+                "connection",
+                resource.name,
+                "Kafka/S3/GCS directives are not valid for dynamodb connections.",
+            )
+
+        if not dynamodb_arn:
+            raise MigrationParseError(
+                resource.file_path,
+                "connection",
+                resource.name,
+                "DYNAMODB_ARN is required for dynamodb connections.",
+            )
+
+        if not dynamodb_region:
+            raise MigrationParseError(
+                resource.file_path,
+                "connection",
+                resource.name,
+                "DYNAMODB_REGION is required for dynamodb connections.",
+            )
+
+        return DynamoDBConnectionModel(
+            kind="connection",
+            name=resource.name,
+            file_path=resource.file_path,
+            connection_type="dynamodb",
+            region=dynamodb_region,
+            arn=dynamodb_arn,
+        )
+
     raise MigrationParseError(
         resource.file_path,
         "connection",