Merge pull request #11 from tinybirdco/support-multiline-connection-directives

Support multiline directives in connection file parser

Author: rmorehig

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tinybird-sdk"
-version = "0.1.3"
+version = "0.1.4"
 description = "Python SDK for Tinybird Forward"
 readme = "README.md"
 authors = [

src/tinybird_sdk/generator/connection.py

@@ -34,7 +34,11 @@ def _generate_kafka_connection(connection: KafkaConnectionDefinition) -> str:
     if options.schema_registry_url:
         parts.append(f"KAFKA_SCHEMA_REGISTRY_URL {options.schema_registry_url}")
     if options.ssl_ca_pem:
-        parts.append(f"KAFKA_SSL_CA_PEM {options.ssl_ca_pem}")
+        if "\n" in options.ssl_ca_pem:
+            indented = "\n".join(f"    {line}" for line in options.ssl_ca_pem.split("\n"))
+            parts.append(f"KAFKA_SSL_CA_PEM >\n{indented}")
+        else:
+            parts.append(f"KAFKA_SSL_CA_PEM {options.ssl_ca_pem}")
 
     return "\n".join(parts)
 

src/tinybird_sdk/migrate/parse_connection.py

@@ -1,8 +1,31 @@
 from __future__ import annotations
 
-from .parser_utils import MigrationParseError, is_blank, parse_directive_line, parse_quoted_value, split_lines
+from .parser_utils import MigrationParseError, is_blank, parse_directive_line, parse_quoted_value, read_directive_block, split_lines
 from .types import ConnectionModel, GCSConnectionModel, KafkaConnectionModel, ResourceFile, S3ConnectionModel
 
+CONNECTION_DIRECTIVES = {
+    "TYPE",
+    "KAFKA_BOOTSTRAP_SERVERS",
+    "KAFKA_SECURITY_PROTOCOL",
+    "KAFKA_SASL_MECHANISM",
+    "KAFKA_KEY",
+    "KAFKA_SECRET",
+    "KAFKA_SCHEMA_REGISTRY_URL",
+    "KAFKA_SSL_CA_PEM",
+    "S3_REGION",
+    "S3_ARN",
+    "S3_ACCESS_KEY",
+    "S3_SECRET",
+    "GCS_SERVICE_ACCOUNT_CREDENTIALS_JSON",
+}
+
+
+def _is_connection_directive_line(line: str) -> bool:
+    if not line:
+        return False
+    directive = parse_directive_line(line)
+    return directive["key"] in CONNECTION_DIRECTIVES
+
 
 def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
     lines = split_lines(resource.content)
@@ -22,9 +45,19 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
     access_secret: str | None = None
     service_account_credentials_json: str | None = None
 
-    for raw_line in lines:
+    i = 0
+    while i < len(lines):
+        raw_line = lines[i]
         line = raw_line.strip()
+
         if is_blank(line) or line.startswith("#"):
+            i += 1
+            continue
+
+        if line == "KAFKA_SSL_CA_PEM >":
+            block, next_index = read_directive_block(lines, i + 1, _is_connection_directive_line)
+            ssl_ca_pem = "\n".join(block)
+            i = next_index
             continue
 
         directive = parse_directive_line(line)
@@ -81,6 +114,8 @@ def parse_connection_file(resource: ResourceFile) -> ConnectionModel:
                 f'Unsupported connection directive in strict mode: "{line}"',
             )
 
+        i += 1
+
     if not connection_type:
         raise MigrationParseError(resource.file_path, "connection", resource.name, "TYPE directive is required.")
 

tests/test_phase1_schema_generator_parity.py

@@ -20,6 +20,36 @@ def test_generate_connection_includes_kafka_schema_registry_url() -> None:
     assert "KAFKA_SCHEMA_REGISTRY_URL https://registry.example.com" in generated.content
 
 
+def test_generate_connection_emits_multiline_ssl_ca_pem() -> None:
+    pem = "-----BEGIN CERTIFICATE-----\nMIIDXTCCAkWgAwIBAgIJAM\n-----END CERTIFICATE-----"
+    kafka = define_kafka_connection(
+        "broker",
+        {
+            "bootstrap_servers": "localhost:9092",
+            "ssl_ca_pem": pem,
+        },
+    )
+
+    generated = generate_connection(kafka)
+    assert "KAFKA_SSL_CA_PEM >" in generated.content
+    assert "    -----BEGIN CERTIFICATE-----" in generated.content
+    assert "    MIIDXTCCAkWgAwIBAgIJAM" in generated.content
+    assert "    -----END CERTIFICATE-----" in generated.content
+
+
+def test_generate_connection_emits_single_line_ssl_ca_pem() -> None:
+    kafka = define_kafka_connection(
+        "broker",
+        {
+            "bootstrap_servers": "localhost:9092",
+            "ssl_ca_pem": "{{ tb_secret('KAFKA_SSL_CA_PEM') }}",
+        },
+    )
+
+    generated = generate_connection(kafka)
+    assert "KAFKA_SSL_CA_PEM {{ tb_secret('KAFKA_SSL_CA_PEM') }}" in generated.content
+
+
 def test_generate_datasource_includes_indexes_and_store_raw_value() -> None:
     kafka = define_kafka_connection("broker", {"bootstrap_servers": "localhost:9092"})
     datasource = define_datasource(

tests/test_phase3_migrate_parser_parity.py

@@ -44,6 +44,72 @@ def test_parse_connection_supports_single_quoted_type_and_schema_registry() -> N
     assert parsed.schema_registry_url == "https://registry.example.com"
 
 
+def test_parse_connection_supports_multiline_ssl_ca_pem() -> None:
+    parsed = parse_connection_file(
+        _resource(
+            "connection",
+            "broker",
+            "\n".join(
+                [
+                    "TYPE kafka",
+                    "KAFKA_BOOTSTRAP_SERVERS localhost:9092",
+                    "KAFKA_SECURITY_PROTOCOL SASL_SSL",
+                    "KAFKA_SSL_CA_PEM >",
+                    "    -----BEGIN CERTIFICATE-----",
+                    "    MIIDXTCCAkWgAwIBAgIJAM",
+                    "    -----END CERTIFICATE-----",
+                ]
+            ),
+        )
+    )
+
+    assert parsed.connection_type == "kafka"
+    assert parsed.ssl_ca_pem == "-----BEGIN CERTIFICATE-----\nMIIDXTCCAkWgAwIBAgIJAM\n-----END CERTIFICATE-----"
+
+
+def test_parse_connection_supports_multiline_ssl_ca_pem_with_following_directives() -> None:
+    parsed = parse_connection_file(
+        _resource(
+            "connection",
+            "broker",
+            "\n".join(
+                [
+                    "TYPE kafka",
+                    "KAFKA_BOOTSTRAP_SERVERS localhost:9092",
+                    "KAFKA_SSL_CA_PEM >",
+                    "    -----BEGIN CERTIFICATE-----",
+                    "    MIIDXTCCAkWgAwIBAgIJAM",
+                    "    -----END CERTIFICATE-----",
+                    "KAFKA_SECURITY_PROTOCOL SASL_SSL",
+                    "KAFKA_KEY mykey",
+                ]
+            ),
+        )
+    )
+
+    assert parsed.ssl_ca_pem == "-----BEGIN CERTIFICATE-----\nMIIDXTCCAkWgAwIBAgIJAM\n-----END CERTIFICATE-----"
+    assert parsed.security_protocol == "SASL_SSL"
+    assert parsed.key == "mykey"
+
+
+def test_parse_connection_supports_single_line_ssl_ca_pem() -> None:
+    parsed = parse_connection_file(
+        _resource(
+            "connection",
+            "broker",
+            "\n".join(
+                [
+                    "TYPE kafka",
+                    "KAFKA_BOOTSTRAP_SERVERS localhost:9092",
+                    "KAFKA_SSL_CA_PEM {{ tb_secret('KAFKA_SSL_CA_PEM') }}",
+                ]
+            ),
+        )
+    )
+
+    assert parsed.ssl_ca_pem == "{{ tb_secret('KAFKA_SSL_CA_PEM') }}"
+
+
 def test_parse_datasource_supports_engine_is_deleted_and_kafka_store_raw_value() -> None:
     parsed = parse_datasource_file(
         _resource(