Rebase + fixes

Author: Tommy Healy

src/tinybird_sdk/__init__.py

@@ -32,7 +32,10 @@
     "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"),
+    "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

@@ -92,9 +92,13 @@ def generate_connection(connection: ConnectionDefinition) -> GeneratedConnection
             name=connection._name, content=_generate_s3_connection(connection)
         )
     if isinstance(connection, GCSConnectionDefinition):
-        return GeneratedConnection(name=connection._name, content=_generate_gcs_connection(connection))
+        return GeneratedConnection(
+            name=connection._name, content=_generate_gcs_connection(connection)
+        )
     if isinstance(connection, DynamoDBConnectionDefinition):
-        return GeneratedConnection(name=connection._name, content=_generate_dynamodb_connection(connection))
+        return GeneratedConnection(
+            name=connection._name, content=_generate_dynamodb_connection(connection)
+        )
     raise ValueError(f"Unsupported connection type: {connection._connectionType}")
 
 

src/tinybird_sdk/migrate/emit_ts.py

@@ -348,7 +348,9 @@ def _emit_gcs_connection(connection: GCSConnectionModel) -> str:
 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"{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("})")
@@ -357,7 +359,10 @@ def _emit_dynamodb_connection(connection: DynamoDBConnectionModel) -> str:
 
 
 def _emit_connection(
-    connection: KafkaConnectionModel | S3ConnectionModel | GCSConnectionModel | DynamoDBConnectionModel,
+    connection: KafkaConnectionModel
+    | S3ConnectionModel
+    | GCSConnectionModel
+    | DynamoDBConnectionModel,
 ) -> str:
     if isinstance(connection, S3ConnectionModel):
         return _emit_s3_connection(connection)

src/tinybird_sdk/migrate/parse_connection.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from .parser_utils import MigrationParseError, is_blank, parse_directive_line, parse_quoted_value, read_directive_block, split_lines
+from .parser_utils import (
+    MigrationParseError,
+    is_blank,
+    parse_directive_line,
+    parse_quoted_value,
+    read_directive_block,
+    split_lines,
+)
 from .types import (
     ConnectionModel,
     DynamoDBConnectionModel,

src/tinybird_sdk/migrate/runner.py

@@ -194,7 +194,9 @@ def run_migrate(options: MigrateOptions | dict[str, Any]) -> MigrationResult:
                 )
 
         if datasource.dynamodb:
-            dynamodb_connection_type = parsed_connection_type_by_name.get(datasource.dynamodb.connection_name)
+            dynamodb_connection_type = parsed_connection_type_by_name.get(
+                datasource.dynamodb.connection_name
+            )
             if dynamodb_connection_type != "dynamodb":
                 errors.append(
                     MigrationError(

src/tinybird_sdk/migrate/types.py

@@ -227,7 +227,9 @@ class DynamoDBConnectionModel:
     arn: str
 
 
-ConnectionModel = KafkaConnectionModel | S3ConnectionModel | GCSConnectionModel | DynamoDBConnectionModel
+ConnectionModel = (
+    KafkaConnectionModel | S3ConnectionModel | GCSConnectionModel | DynamoDBConnectionModel
+)
 
 ParsedResource = (
     DatasourceModel

src/tinybird_sdk/schema/connection.py

@@ -80,7 +80,10 @@ class DynamoDBConnectionDefinition:
 
 
 ConnectionDefinition = (
-    KafkaConnectionDefinition | S3ConnectionDefinition | GCSConnectionDefinition | DynamoDBConnectionDefinition
+    KafkaConnectionDefinition
+    | S3ConnectionDefinition
+    | GCSConnectionDefinition
+    | DynamoDBConnectionDefinition
 )
 
 
@@ -133,7 +136,11 @@ def define_dynamodb_connection(
     name: str, options: dict[str, Any] | DynamoDBConnectionOptions
 ) -> DynamoDBConnectionDefinition:
     _validate_connection_name(name)
-    normalized = options if isinstance(options, DynamoDBConnectionOptions) else DynamoDBConnectionOptions(**options)
+    normalized = (
+        options
+        if isinstance(options, DynamoDBConnectionOptions)
+        else DynamoDBConnectionOptions(**options)
+    )
 
     if not normalized.region.strip():
         raise ValueError("DynamoDB connection `region` is required.")
@@ -147,7 +154,12 @@ def define_dynamodb_connection(
 def is_connection_definition(value: Any) -> bool:
     return isinstance(
         value,
-        (KafkaConnectionDefinition, S3ConnectionDefinition, GCSConnectionDefinition, DynamoDBConnectionDefinition),
+        (
+            KafkaConnectionDefinition,
+            S3ConnectionDefinition,
+            GCSConnectionDefinition,
+            DynamoDBConnectionDefinition,
+        ),
     )
 
 

src/tinybird_sdk/schema/datasource.py

@@ -134,7 +134,9 @@ def define_datasource(
         gcs_cfg = options.get("gcs")
         gcs = GCSConfig(**gcs_cfg) if isinstance(gcs_cfg, dict) else gcs_cfg
         dynamodb_cfg = options.get("dynamodb")
-        dynamodb = DynamoDBConfig(**dynamodb_cfg) if isinstance(dynamodb_cfg, dict) else dynamodb_cfg
+        dynamodb = (
+            DynamoDBConfig(**dynamodb_cfg) if isinstance(dynamodb_cfg, dict) else dynamodb_cfg
+        )
         normalized = DatasourceOptions(
             description=options.get("description"),
             schema=options["schema"],
@@ -152,10 +154,14 @@ def define_datasource(
         )
 
     ingestion_count = sum(
-        1 for x in [normalized.kafka, normalized.s3, normalized.gcs, normalized.dynamodb] if x is not None
+        1
+        for x in [normalized.kafka, normalized.s3, normalized.gcs, normalized.dynamodb]
+        if x is not None
     )
     if ingestion_count > 1:
-        raise ValueError("Datasource can only define one ingestion option: `kafka`, `s3`, `gcs`, or `dynamodb`.")
+        raise ValueError(
+            "Datasource can only define one ingestion option: `kafka`, `s3`, `gcs`, or `dynamodb`."
+        )
 
     if normalized.backfill not in {None, "skip"}:
         raise ValueError('Invalid datasource backfill value: only "skip" is supported.')

tests/test_dynamodb_connector.py

@@ -205,7 +205,7 @@ def test_emit_ts_round_trip_for_dynamodb() -> None:
 
     output = emit_migration_file_content([connection, datasource])
     assert "define_dynamodb_connection" in output
-    assert "events_dynamodb = define_dynamodb_connection(\"events_dynamodb\"" in output
+    assert 'events_dynamodb = define_dynamodb_connection("events_dynamodb"' in output
     assert "'dynamodb': {" in output
     assert "'table_arn': \"arn:aws:dynamodb:us-east-1:1:table/orders\"" in output
     assert "'export_bucket': \"s3://export\"" in output