tinybird 0.0.1.dev196__tar.gz → 0.0.1.dev198__tar.gz

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tinybird
-Version: 0.0.1.dev196
+Version: 0.0.1.dev198
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/forward/commands
 Author: Tinybird
@@ -13,6 +13,7 @@ Requires-Dist: boto3
 Requires-Dist: click<8.2,>=8.1.6
 Requires-Dist: clickhouse-toolset==0.34.dev0
 Requires-Dist: colorama==0.4.6
+Requires-Dist: confluent-kafka==2.8.0
 Requires-Dist: cryptography~=41.0.0
 Requires-Dist: croniter==1.3.15
 Requires-Dist: docker==7.1.0

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird/sql_toolset.py

@@ -227,25 +227,56 @@ def sql_get_used_tables(
 
 
 class ReplacementsDict(dict):
+    def __init__(self, *args, enabled_table_functions=None, **kwargs):
+        self.enabled_table_functions = enabled_table_functions
+        super().__init__(*args, **kwargs)
+
     def __getitem__(self, key):
         v = super().__getitem__(key)
         if isinstance(v, tuple):
             k, r = v
             if callable(r):
-                r = r()
+                r = update_callable_signature(r)(self.enabled_table_functions)
                 super().__setitem__(key, (k, r))
             return k, r
         if callable(v):
-            v = v()
+            v = update_callable_signature(v)(self.enabled_table_functions)
             super().__setitem__(key, v)
         return v
 
 
-def tables_or_sql(replacement: dict, table_functions=False) -> set:
+def update_callable_signature(func):
+    """
+    Utility function to provide backward compatibility for callable functions
+    that don't accept the enabled_table_functions parameter.
+    """
+    if callable(func):
+
+        def wrapper(enabled_table_functions=None):
+            # Check if the function accepts the enabled_table_functions parameter
+            import inspect
+
+            sig = inspect.signature(func)
+            if len(sig.parameters) == 0:
+                # Old-style callable with no parameters
+                return func()
+            else:
+                # New-style callable that accepts enabled_table_functions
+                return func(enabled_table_functions)
+
+        return wrapper
+    return func
+
+
+def tables_or_sql(replacement: dict, table_functions=False, function_allow_list=None) -> set:
     try:
         return set(
             sql_get_used_tables(
-                replacement[1], default_database=replacement[0], raising=True, table_functions=table_functions
+                replacement[1],
+                default_database=replacement[0],
+                raising=True,
+                table_functions=table_functions,
+                function_allow_list=frozenset(function_allow_list),
             )
         )
     except Exception as e:
@@ -342,6 +373,7 @@ def replace_tables(
         _enabled_table_functions = ENABLED_TABLE_FUNCTIONS
     else:
         _enabled_table_functions = ENABLED_TABLE_FUNCTIONS.union(set(function_allow_list))
+    _replacements.enabled_table_functions = frozenset(_enabled_table_functions)
     while _tables:
         table = _tables.pop()
         if len(table) == 3:
@@ -355,7 +387,9 @@ def replace_tables(
         seen_tables.add(table)
         if table in _replacements:
             replacement = _replacements[table]
-            dependent_tables = tables_or_sql(replacement, table_functions=check_functions)
+            dependent_tables = tables_or_sql(
+                replacement, table_functions=check_functions, function_allow_list=_enabled_table_functions
+            )
             deps[table] |= {(d[0], d[1]) for d in dependent_tables}
             for dependent_table in list(dependent_tables):
                 if len(dependent_table) == 3:

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird/tb/__cli__.py

@@ -4,5 +4,5 @@ __description__ = 'Tinybird Command Line Tool'
 __url__ = 'https://www.tinybird.co/docs/forward/commands'
 __author__ = 'Tinybird'
 __author_email__ = 'support@tinybird.co'
-__version__ = '0.0.1.dev196'
-__revision__ = 'a82d1e4'
+__version__ = '0.0.1.dev198'
+__revision__ = '7aafa44'

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird/tb/modules/common.py

@@ -1899,10 +1899,13 @@ def get_connection_name(project_folder: str, connection_type: str, connection_na
 
     while not connection_name:
         short_id = str(uuid.uuid4())[:4]
+        default_name = f"{connection_type.lower()}_{short_id}"
         connection_name = click.prompt(
-            "Enter a name (only alphanumeric characters and underscores)",
-            show_default=True,
-            default=f"{connection_type.lower()}_{short_id}",
+            FeedbackManager.highlight(
+                message=f"? Connection name (only alphanumeric characters and underscores) [{default_name}]"
+            ),
+            show_default=False,
+            default=default_name,
         )
         assert isinstance(connection_name, str)
 

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird/tb/modules/connection.py

@@ -4,10 +4,11 @@
 # - But please, **do not** interleave utility functions and command definitions.
 
 import uuid
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import click
 from click import Context
+from confluent_kafka.admin import AdminClient
 
 from tinybird.tb.client import TinyB
 from tinybird.tb.modules.cli import cli
@@ -23,12 +24,18 @@ from tinybird.tb.modules.common import (
     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.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
 from tinybird.tb.modules.project import Project
 
@@ -280,18 +287,235 @@ async def connection_create_gcs(ctx: Context) -> None:
 
 
 @connection_create.command(name="kafka", short_help="Creates a Kafka connection.")
-@click.option("--name", help="The name of the connection")
+@click.option(
+    "--name",
+    default=None,
+    help="The name of your Kafka connection. If not provided, it's set as the bootstrap server",
+)
+@click.option("--bootstrap-servers", help="Kafka Bootstrap Server in form mykafka.mycloud.com:9092")
+@click.option("--key", help="Key")
+@click.option("--secret", help="Secret")
+@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=None,
+    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("--skip-secrets", is_flag=True, help="Skip Tinybird Secret creation")
 @click.pass_context
 @coro
-async def connection_create_kafka(ctx: Context, name: Optional[str] = None) -> None:
+async def connection_create_kafka_cmd(
+    ctx: Context,
+    name: str,
+    bootstrap_servers: str,
+    key: str,
+    secret: str,
+    schema_registry_url: str,
+    auto_offset_reset: str,
+    sasl_mechanism: str,
+    security_protocol: str,
+    skip_secrets: bool,
+) -> None:
     """
     Creates a Kafka connection.
 
     \b
     $ tb connection create kafka
     """
-    click.echo(FeedbackManager.highlight(message="» Creating Kafka connection..."))
+    await connection_create_kafka(
+        ctx,
+        name,
+        bootstrap_servers,
+        key,
+        secret,
+        schema_registry_url,
+        auto_offset_reset,
+        sasl_mechanism,
+        security_protocol,
+        skip_secrets,
+    )
+
+
+async def connection_create_kafka(
+    ctx: Context,
+    name="",
+    bootstrap_servers="",
+    key="",
+    secret="",
+    schema_registry_url="",
+    auto_offset_reset="",
+    sasl_mechanism="",
+    security_protocol="SASL_SSL",
+    skip_secrets=False,
+) -> Tuple[str, str, str, str, str, str, str, str, List[str]]:
+    client: TinyB = ctx.ensure_object(dict)["client"]
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    click.echo(FeedbackManager.gray(message="\n» Creating Kafka connection..."))
     project: Project = ctx.ensure_object(dict)["project"]
     name = get_kafka_connection_name(project.folder, name)
-    await generate_kafka_connection_with_secrets(name=name, folder=project.folder)
-    click.echo(FeedbackManager.success(message="✓ Done!"))
+    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 = not skip_secrets and click.confirm(
+        FeedbackManager.highlight(message="  ? Do you want to store the bootstrap server in a Tinybird Secret? [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:
+            await client.create_secret(name=tb_secret_bootstrap_servers, value=bootstrap_servers)
+        except Exception as e:
+            raise CLIConnectionException(FeedbackManager.error(message=str(e)))
+
+    if not key:
+        key = click.prompt(FeedbackManager.highlight(message="? Kafka key"))
+
+    validate_kafka_key(key)
+
+    secret_required = not skip_secrets and click.confirm(
+        FeedbackManager.highlight(message="  ? Do you want to store the Kafka key in a Tinybird Secret? [Y/n]"),
+        default=True,
+        show_default=False,
+    )
+
+    if secret_required:
+        tb_secret_key = str(click.prompt(FeedbackManager.highlight(message="    ? Secret name")))
+        try:
+            await client.create_secret(name=tb_secret_key, value=key)
+        except Exception as e:
+            raise CLIConnectionException(FeedbackManager.error(message=str(e)))
+
+    if not secret:
+        secret = click.prompt(FeedbackManager.highlight(message="? Kafka secret"), hide_input=True)
+
+    validate_kafka_secret(secret)
+
+    secret_required = not skip_secrets and click.confirm(
+        FeedbackManager.highlight(message="  ? Do you want to store the Kafka secret in a Tinybird Secret? [Y/n]"),
+        default=True,
+        show_default=False,
+    )
+
+    if secret_required:
+        tb_secret_secret = str(click.prompt(FeedbackManager.highlight(message="    ? Secret name")))
+        try:
+            await client.create_secret(name=tb_secret_secret, value=secret)
+        except Exception as e:
+            raise CLIConnectionException(FeedbackManager.error(message=str(e)))
+
+    if not sasl_mechanism:
+        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,
+            show_default=False,
+        )
+        if obj["env"] == "local"
+        and not skip_secrets
+        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:
+            await prod_client.create_secret(name=tb_secret_bootstrap_servers, value=bootstrap_servers)
+        if tb_secret_key:
+            await prod_client.create_secret(name=tb_secret_key, value=key)
+        if tb_secret_secret:
+            await prod_client.create_secret(name=tb_secret_secret, value=secret)
+        click.echo(FeedbackManager.success(message="✓ Secrets created!"))
+
+    schema_registry_url and validate_kafka_schema_registry_url(schema_registry_url)
+    auto_offset_reset and 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"))
+    await generate_kafka_connection_with_secrets(
+        name=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,
+        sasl_mechanism=sasl_mechanism,
+        folder=project.folder,
+    )
+    click.echo(FeedbackManager.info(message=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,
+    )
+
+
+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

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird/tb/modules/create.py

@@ -540,12 +540,32 @@ TYPE ENDPOINT
     return result
 
 
-async def generate_kafka_connection_with_secrets(name: str, folder: str) -> Path:
-    content = """TYPE kafka
-KAFKA_BOOTSTRAP_SERVERS {{ tb_secret("KAFKA_SERVERS", "localhost:9092") }}
-KAFKA_SECURITY_PROTOCOL SASL_SSL
-KAFKA_SASL_MECHANISM PLAIN
-KAFKA_KEY {{ tb_secret("KAFKA_USERNAME", "") }}
-KAFKA_SECRET {{ tb_secret("KAFKA_PASSWORD", "") }}
+async def generate_kafka_connection_with_secrets(
+    name: str,
+    bootstrap_servers: str,
+    key: str,
+    secret: str,
+    tb_secret_bootstrap_servers: Optional[str],
+    tb_secret_key: Optional[str],
+    tb_secret_secret: Optional[str],
+    security_protocol: str,
+    sasl_mechanism: str,
+    folder: str,
+) -> Path:
+    kafka_bootstrap_servers = (
+        inject_tb_secret(tb_secret_bootstrap_servers) if tb_secret_bootstrap_servers else bootstrap_servers
+    )
+    kafka_key = inject_tb_secret(tb_secret_key) if tb_secret_key else key
+    kafka_secret = inject_tb_secret(tb_secret_secret) if tb_secret_secret else secret
+    content = f"""TYPE kafka
+KAFKA_BOOTSTRAP_SERVERS {kafka_bootstrap_servers}
+KAFKA_SECURITY_PROTOCOL {security_protocol or "SASL_SSL"}
+KAFKA_SASL_MECHANISM {sasl_mechanism or "PLAIN"}
+KAFKA_KEY {kafka_key}
+KAFKA_SECRET {kafka_secret}
 """
     return generate_connection_file(name, content, folder, skip_feedback=True)
+
+
+def inject_tb_secret(secret_name: str) -> str:
+    return f"""{{{{ tb_secret("{secret_name}") }}}}"""

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird/tb/modules/datasource.py

@@ -10,7 +10,7 @@ import re
 import uuid
 from datetime import datetime
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 from urllib.parse import urlparse
 
 import click
@@ -33,11 +33,11 @@ from tinybird.tb.modules.common import (
     push_data,
 )
 from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.connection import connection_create_kafka
 from tinybird.tb.modules.create import (
     create_resources_from_prompt,
     generate_aws_iamrole_connection_file_with_secret,
     generate_gcs_connection_file_with_secrets,
-    generate_kafka_connection_with_secrets,
 )
 from tinybird.tb.modules.datafile.common import get_name_version
 from tinybird.tb.modules.datafile.fixture import persist_fixture
@@ -696,8 +696,9 @@ async def datasource_create(
     )
     datasource_type: Optional[str] = None
     connection_file: Optional[str] = None
-    ds_content = ""
-
+    ds_content = """SCHEMA >
+    `data` String `json:$`
+"""
     if file:
         datasource_type = "Local file"
     elif url:
@@ -780,6 +781,8 @@ async def datasource_create(
             return connection_files
 
         connection_files = get_connection_files()
+        connection_name = ""
+        topics: List[str] = []
         if len(connection_files) == 0:
             click.echo(FeedbackManager.error(message=f"x No {datasource_type} connections found."))
             if click.confirm(
@@ -787,15 +790,26 @@ async def datasource_create(
                 show_default=False,
                 default=True,
             ):
-                click.echo(FeedbackManager.gray(message="\n» Creating .connection file..."))
-                default_connection_name = f"{datasource_type.lower()}_{generate_short_id()}"
-                connection_name = click.prompt(
-                    FeedbackManager.highlight(message=f"? Connection name [{default_connection_name}]"),
-                    show_default=False,
-                    default=default_connection_name,
-                )
+                if datasource_type != "Kafka":
+                    click.echo(FeedbackManager.gray(message="\n» Creating .connection file..."))
+                    default_connection_name = f"{datasource_type.lower()}_{generate_short_id()}"
+                    connection_name = click.prompt(
+                        FeedbackManager.highlight(message=f"? Connection name [{default_connection_name}]"),
+                        show_default=False,
+                        default=default_connection_name,
+                    )
                 if datasource_type == "Kafka":
-                    await generate_kafka_connection_with_secrets(connection_name, folder=project.folder)
+                    (
+                        connection_name,
+                        bootstrap_servers,
+                        key,
+                        secret,
+                        schema_registry_url,
+                        auto_offset_reset,
+                        sasl_mechanism,
+                        security_protocol,
+                        topics,
+                    ) = await connection_create_kafka(ctx)
                 elif datasource_type == "S3":
                     await generate_aws_iamrole_connection_file_with_secret(
                         connection_name,
@@ -812,8 +826,9 @@ async def datasource_create(
                         svc_account_creds="GCS_SERVICE_ACCOUNT_CREDENTIALS_JSON",
                         folder=project.folder,
                     )
-                click.echo(FeedbackManager.info(message=f"/connections/{connection_name}.connection"))
-                click.echo(FeedbackManager.success(message="✓ .connection created!"))
+                if datasource_type != "Kafka":
+                    click.echo(FeedbackManager.info(message=f"/connections/{connection_name}.connection"))
+                    click.echo(FeedbackManager.success(message="✓ .connection created!"))
                 connection_files = get_connection_files()
             else:
                 click.echo(
@@ -826,11 +841,6 @@ async def datasource_create(
             connection_path = Path(connection_file)
             connection = connection_path.stem
 
-    ds_content = """SCHEMA >
-`timestamp` DateTime `json:$.timestamp`,
-`session_id` String `json:$.session_id`
-"""
-
     if datasource_type == "Local file":
         click.echo(FeedbackManager.gray(message="\n» Creating .datasource file..."))
         if not file:
@@ -879,16 +889,34 @@ async def datasource_create(
     if datasource_type == "Kafka":
         connections = await client.connections("kafka")
         connection_id = next((c["id"] for c in connections if c["name"] == connection), connection)
-        try:
-            topics = await client.kafka_list_topics(connection_id) if connection_id else []
-        except Exception:
-            topics = []
-        topic = topics[0] if len(topics) > 0 else "topic_0"
+
+        if not topics:
+            try:
+                topics = await client.kafka_list_topics(connection_id) if connection_id else []
+            except Exception:
+                topics = []
+
+        if len(topics) > 1:
+            click.echo(FeedbackManager.highlight(message="? Multiple topics found. Please select one:"))
+            topic_index = -1
+            while topic_index == -1:
+                for index, topic in enumerate(topics):
+                    click.echo(f"  [{index + 1}] {topic}")
+                topic_index = click.prompt("\nSelect option", default=1)
+                try:
+                    topic = topics[int(topic_index) - 1]
+                except Exception:
+                    topic_index = -1
+        else:
+            topic = topics[0] if len(topics) > 0 else "topic_0"
+
         group_id = generate_kafka_group_id(topic)
+
         ds_content += f"""
 KAFKA_CONNECTION_NAME {connection}
-KAFKA_TOPIC {{{{ tb_secret("KAFKA_TOPIC", "{topic}") }}}}
-KAFKA_GROUP_ID {{{{ tb_secret("KAFKA_GROUP_ID", "{group_id}") }}}}
+KAFKA_TOPIC {topic}
+KAFKA_GROUP_ID {group_id}
+KAFKA_STORE_RAW_VALUE 'False'
 """
 
     if datasource_type == "S3":

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tinybird
-Version: 0.0.1.dev196
+Version: 0.0.1.dev198
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/forward/commands
 Author: Tinybird
@@ -13,6 +13,7 @@ Requires-Dist: boto3
 Requires-Dist: click<8.2,>=8.1.6
 Requires-Dist: clickhouse-toolset==0.34.dev0
 Requires-Dist: colorama==0.4.6
+Requires-Dist: confluent-kafka==2.8.0
 Requires-Dist: cryptography~=41.0.0
 Requires-Dist: croniter==1.3.15
 Requires-Dist: docker==7.1.0

{tinybird-0.0.1.dev196 → tinybird-0.0.1.dev198}/tinybird.egg-info/requires.txt

@@ -4,6 +4,7 @@ boto3
 click<8.2,>=8.1.6
 clickhouse-toolset==0.34.dev0
 colorama==0.4.6
+confluent-kafka==2.8.0
 cryptography~=41.0.0
 croniter==1.3.15
 docker==7.1.0