tinybird 0.0.1.dev306__py3-none-any.whl → 1.0.5__py3-none-any.whl

tinybird/tb/modules/connection.py

@@ -4,13 +4,13 @@
 # - But please, **do not** interleave utility functions and command definitions.
 
 import uuid
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import click
 from click import Context
-from confluent_kafka.admin import AdminClient
 
 from tinybird.tb.client import TinyB
+from tinybird.tb.modules.build_common import process as build_project
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import (
     DataConnectorType,
@@ -18,21 +18,21 @@ from tinybird.tb.modules.common import (
     echo_safe_humanfriendly_tables_format_smart_table,
     get_gcs_connection_name,
     get_gcs_svc_account_creds,
-    get_kafka_connection_name,
     get_s3_connection_name,
     production_aws_iamrole_only,
     run_aws_iamrole_connection_flow,
     run_gcp_svc_account_connection_flow,
-    validate_kafka_auto_offset_reset,
-    validate_kafka_bootstrap_servers,
-    validate_kafka_key,
-    validate_kafka_schema_registry_url,
-    validate_kafka_secret,
+)
+from tinybird.tb.modules.connection_kafka import (
+    connection_create_kafka,
+    echo_kafka_data,
+    select_connection,
+    select_group_id,
+    select_topic,
 )
 from tinybird.tb.modules.create import (
     generate_aws_iamrole_connection_file_with_secret,
     generate_gcs_connection_file_with_secrets,
-    generate_kafka_connection_with_secrets,
 )
 from tinybird.tb.modules.exceptions import CLIConnectionException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
@@ -121,7 +121,8 @@ def connection_ls(ctx: Context, service: Optional[DataConnectorType] = None) ->
     obj: Dict[str, Any] = ctx.ensure_object(dict)
     client: TinyB = obj["client"]
 
-    connections = client.connections(connector=service)
+    datasources = client.datasources()
+    connections = client.connections(connector=service, datasources=datasources)
     columns = []
     table = []
 
@@ -314,198 +315,128 @@ def connection_create_gcs(ctx: Context) -> None:
     )
 
 
-@connection_create.command(name="kafka", short_help="Creates a Kafka connection.")
+@connection_create.command(name="kafka", help="Create a Kafka connection.")
+@click.option("--connection-name", default=None, help="The name of the connection to identify it in Tinybird")
+@click.option("--bootstrap-servers", default=None, help="Kafka Bootstrap Server in form mykafka.mycloud.com:9092")
+@click.option("--key", default=None, help="Key/User")
+@click.option("--secret", default=None, help="Secret/Password")
+@click.option(
+    "--auto-offset-reset", default=None, help="Offset reset, can be 'latest' or 'earliest'. Defaults to 'latest'."
+)
+@click.option("--schema-registry-url", default=None, help="Avro Confluent Schema Registry URL")
+@click.option(
+    "--sasl-mechanism",
+    default="PLAIN",
+    help="Authentication method for connection-based protocols. Defaults to 'PLAIN'",
+)
+@click.option(
+    "--security-protocol",
+    default="SASL_SSL",
+    help="Security protocol for connection-based protocols. Defaults to 'SASL_SSL'",
+)
+@click.option("--ssl-ca-pem", default=None, help="Path or content of the CA Certificate file in PEM format")
 @click.pass_context
-def connection_create_kafka_cmd(ctx: Context) -> None:
-    """
-    Creates a Kafka connection.
-
-    \b
-    $ tb connection create kafka
-    """
-    connection_create_kafka(ctx)
-
-
-def connection_create_kafka(ctx: Context) -> Tuple[str, str, str, str, str, str, str, str, List[str]]:
-    obj: Dict[str, Any] = ctx.ensure_object(dict)
-    click.echo(FeedbackManager.gray(message="\n» Creating Kafka connection..."))
+def connection_create_kafka_cmd(
+    ctx: Context,
+    connection_name: Optional[str],
+    bootstrap_servers: Optional[str],
+    key: Optional[str],
+    secret: Optional[str],
+    auto_offset_reset: Optional[str],
+    schema_registry_url: Optional[str],
+    sasl_mechanism: Optional[str],
+    security_protocol: Optional[str],
+    ssl_ca_pem: Optional[str],
+) -> None:
+    env: str = ctx.ensure_object(dict)["env"]
     project: Project = ctx.ensure_object(dict)["project"]
-    name = get_kafka_connection_name(project.folder, "")
-    default_bootstrap_servers = "localhost:9092"
-    bootstrap_servers = click.prompt(
-        FeedbackManager.highlight(
-            message=f"? Bootstrap servers (comma-separated list of host:port pairs) [{default_bootstrap_servers}]"
-        ),
-        default=default_bootstrap_servers,
-        show_default=False,
-    )
-    if not bootstrap_servers:
-        bootstrap_servers = click.prompt("Bootstrap Server")
-
-    validate_kafka_bootstrap_servers(bootstrap_servers)
-    secret_required = click.confirm(
-        FeedbackManager.highlight(message="  ? Do you want to store the bootstrap server in a .env.local file? [Y/n]"),
-        default=True,
-        show_default=False,
-    )
-    tb_secret_bootstrap_servers: Optional[str] = None
-    tb_secret_key: Optional[str] = None
-    tb_secret_secret: Optional[str] = None
-
-    if secret_required:
-        tb_secret_bootstrap_servers = str(click.prompt(FeedbackManager.highlight(message="    ? Secret name")))
-        try:
-            save_secret_to_env_file(project=project, name=tb_secret_bootstrap_servers, value=bootstrap_servers)
-        except Exception as e:
-            raise CLIConnectionException(FeedbackManager.error(message=str(e)))
-
-    key = click.prompt(FeedbackManager.highlight(message="? Kafka key"))
-
-    validate_kafka_key(key)
-
-    secret_required = click.confirm(
-        FeedbackManager.highlight(message="  ? Do you want to store the Kafka key in a .env.local file? [Y/n]"),
-        default=True,
-        show_default=False,
-    )
-
-    if secret_required:
-        tb_secret_key = str(click.prompt(FeedbackManager.highlight(message="    ? Secret name")))
-        try:
-            save_secret_to_env_file(project=project, name=tb_secret_key, value=key)
-        except Exception as e:
-            raise CLIConnectionException(FeedbackManager.error(message=str(e)))
-
-    secret = click.prompt(FeedbackManager.highlight(message="? Kafka secret"), hide_input=True)
-
-    validate_kafka_secret(secret)
-
-    secret_required = click.confirm(
-        FeedbackManager.highlight(message="  ? Do you want to store the Kafka secret in a .env.local file? [Y/n]"),
-        default=True,
-        show_default=False,
-    )
-
-    if secret_required:
-        tb_secret_secret = str(click.prompt(FeedbackManager.highlight(message="    ? Secret name")))
-        try:
-            save_secret_to_env_file(project=project, name=tb_secret_secret, value=secret)
-        except Exception as e:
-            raise CLIConnectionException(FeedbackManager.error(message=str(e)))
-
-    security_protocol_options = ["SASL_SSL", "PLAINTEXT"]
-    security_protocol = click.prompt(
-        FeedbackManager.highlight(message="? Security Protocol (SASL_SSL, PLAINTEXT) [SASL_SSL]"),
-        type=click.Choice(security_protocol_options),
-        show_default=False,
-        show_choices=False,
-        default="SASL_SSL",
-    )
-
-    sasl_mechanism_options = ["PLAIN", "SCRAM-SHA-256", "SCRAM-SHA-512"]
-    sasl_mechanism = click.prompt(
-        FeedbackManager.highlight(message="? SASL Mechanism (PLAIN, SCRAM-SHA-256, SCRAM-SHA-512) [PLAIN]"),
-        type=click.Choice(sasl_mechanism_options),
-        show_default=False,
-        show_choices=False,
-        default="PLAIN",
-    )
-
-    create_in_cloud = (
-        click.confirm(
-            FeedbackManager.highlight(
-                message="? Would you like to create this connection in the cloud environment as well? [Y/n]"
-            ),
-            default=True,
+    client: TinyB = ctx.ensure_object(dict)["client"]
+    config: Dict[str, Any] = ctx.ensure_object(dict)["config"]
+    if env == "local":
+        should_build = click.confirm(
+            FeedbackManager.highlight(message="? Do you want to build the project before continue? [Y/n]"),
             show_default=False,
+            default=True,
         )
-        if obj["env"] == "local" and (tb_secret_bootstrap_servers or tb_secret_key or tb_secret_secret)
-        else False
-    )
-
-    if create_in_cloud:
-        click.echo(FeedbackManager.gray(message="» Creating Secrets in cloud environment..."))
-        prod_config = obj["config"]
-        host = prod_config["host"]
-        token = prod_config["token"]
-        prod_client = TinyB(
-            token=token,
-            host=host,
-            staging=False,
-        )
-        if tb_secret_bootstrap_servers:
-            prod_client.create_secret(name=tb_secret_bootstrap_servers, value=bootstrap_servers)
-        if tb_secret_key:
-            prod_client.create_secret(name=tb_secret_key, value=key)
-        if tb_secret_secret:
-            prod_client.create_secret(name=tb_secret_secret, value=secret)
-        click.echo(FeedbackManager.success(message="✓ Secrets created!"))
-
-    schema_registry_url = click.prompt(
-        FeedbackManager.highlight(message="? Schema Registry URL (optional)"),
-        default="",
-        show_default=False,
-    )
-    if schema_registry_url:
-        validate_kafka_schema_registry_url(schema_registry_url)
-
-    auto_offset_reset_options = ["latest", "earliest"]
-    auto_offset_reset = click.prompt(
-        FeedbackManager.highlight(message="? Auto offset reset (latest, earliest) [latest]"),
-        type=click.Choice(auto_offset_reset_options),
-        default="latest",
-        show_default=False,
-        show_choices=False,
-    )
-    validate_kafka_auto_offset_reset(auto_offset_reset)
-    click.echo(FeedbackManager.gray(message="» Validating connection..."))
-
-    topics = list_kafka_topics(bootstrap_servers, key, secret, security_protocol, sasl_mechanism)
-
-    if topics is None:
-        raise CLIConnectionException(FeedbackManager.error(message="Invalid Kafka connection"))
-
-    click.echo(FeedbackManager.success(message="✓ Connection is valid"))
-    generate_kafka_connection_with_secrets(
-        name=name,
+        if should_build:
+            click.echo(FeedbackManager.gray(message="» Building project..."))
+            build_project(project=project, tb_client=client, watch=False, config=config, silent=True)
+            click.echo(FeedbackManager.success(message="✓ Build completed"))
+        else:
+            click.echo(FeedbackManager.gray(message="Skipping build..."))
+
+    result = connection_create_kafka(
+        ctx,
+        connection_name=connection_name,
         bootstrap_servers=bootstrap_servers,
-        tb_secret_bootstrap_servers=tb_secret_bootstrap_servers,
         key=key,
-        tb_secret_key=tb_secret_key,
         secret=secret,
-        tb_secret_secret=tb_secret_secret,
-        security_protocol=security_protocol,
+        auto_offset_reset=auto_offset_reset,
+        schema_registry_url=schema_registry_url,
         sasl_mechanism=sasl_mechanism,
-        folder=project.folder,
-    )
-    click.echo(FeedbackManager.info_file_created(file=f"connections/{name}.connection"))
-    click.echo(FeedbackManager.success(message="✓ Connection created!"))
-    return (
-        name,
-        bootstrap_servers,
-        key,
-        secret,
-        schema_registry_url,
-        auto_offset_reset,
-        sasl_mechanism,
-        security_protocol,
-        topics,
+        security_protocol=security_protocol,
+        ssl_ca_pem=ssl_ca_pem,
     )
 
+    if not result["error"]:
+        yes = click.confirm(
+            FeedbackManager.highlight(message="? Would you like to preview data from the connection?"),
+            default=True,
+            show_default=False,
+        )
+        if yes:
+            connection_data(connection_name=result["name"], client=client)
+        else:
+            click.echo(FeedbackManager.gray(message="Skipping data preview."))
+
+
+@connection.command(name="data", help="Preview data from an existing connection.")
+@click.argument("connection_name", type=str, required=False)
+@click.option("--kafka-topic", type=str, help="The Kafka topic to preview")
+@click.option("--kafka-group-id", type=str, help="The Kafka group ID to use for preview")
+@click.pass_context
+def connection_data_cmd(
+    ctx: Context,
+    connection_name: Optional[str] = None,
+    kafka_topic: Optional[str] = None,
+    kafka_group_id: Optional[str] = None,
+) -> None:
+    project: Project = ctx.ensure_object(dict)["project"]
+    client: TinyB = ctx.ensure_object(dict)["client"]
+    env: str = ctx.ensure_object(dict)["env"]
+    config: Dict[str, Any] = ctx.ensure_object(dict)["config"]
+
+    if env == "local":
+        click.echo(FeedbackManager.gray(message="» Building project to access the latest connections..."))
+        build_project(project=project, tb_client=client, watch=False, config=config, silent=True)
+        click.echo(FeedbackManager.success(message="✓ Build completed"))
+
+    if connection_name is None:
+        connections = client.connections("kafka")
+        connection = select_connection(None, "kafka", connections, client)
+        connection_name = connection["name"]
+
+    connection_data(connection_name, client, kafka_topic, kafka_group_id)
+
+
+def connection_data(
+    connection_name: str,
+    client: TinyB,
+    kafka_topic: Optional[str] = None,
+    kafka_group_id: Optional[str] = None,
+) -> None:
+    connection = next((c for c in client.connections() if c["name"] == connection_name), None)
+    if not connection:
+        raise CLIConnectionException(FeedbackManager.error(message=f"Connection {connection_name} not found."))
+
+    if connection["service"] != DataConnectorType.KAFKA:
+        raise CLIConnectionException(
+            FeedbackManager.error(
+                message=f"{connection['service']} connections are not supported yet for previewing data. Supported connections: kafka"
+            )
+        )
 
-def list_kafka_topics(bootstrap_servers, sasl_username, sasl_password, security_protocol, sasl_mechanism):
-    conf = {
-        "bootstrap.servers": bootstrap_servers,
-        "security.protocol": security_protocol,
-        "sasl.mechanism": sasl_mechanism,
-        "sasl.username": sasl_username,
-        "sasl.password": sasl_password,
-        "log_level": 0,
-    }
-
-    try:
-        client = AdminClient(conf)
-        metadata = client.list_topics(timeout=5)
-        return list(metadata.topics.keys())
-    except Exception:
-        return None
+    connection_id = connection["id"]
+    topic = select_topic(kafka_topic, connection_id, client)
+    group_id = select_group_id(kafka_group_id, topic, connection_id, client)
+    echo_kafka_data(connection_id, connection_name, topic, group_id, client)