tinybird 0.0.1.dev19__py3-none-any.whl → 0.0.1.dev21__py3-none-any.whl

tinybird/tb/modules/fmt.py

@@ -7,13 +7,13 @@ import aiofiles
 import click
 from click import Context
 
-from tinybird.feedback_manager import FeedbackManager
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import coro
 from tinybird.tb.modules.datafile.common import is_file_a_datasource
 from tinybird.tb.modules.datafile.diff import color_diff, peek
 from tinybird.tb.modules.datafile.format_datasource import format_datasource
 from tinybird.tb.modules.datafile.format_pipe import format_pipe
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
 @cli.command()

tinybird/tb/modules/job.py

@@ -7,10 +7,10 @@ import click
 from click import Context
 
 from tinybird.client import DoesNotExistException, TinyB
-from tinybird.feedback_manager import FeedbackManager
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import coro, echo_safe_humanfriendly_tables_format_smart_table
 from tinybird.tb.modules.exceptions import CLIException
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
 @cli.group()

tinybird/tb/modules/llm.py

@@ -8,8 +8,6 @@ from openai import OpenAI
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
-from tinybird.prompts import create_test_calls_prompt
-from tinybird.tb.modules.config import CLIConfig
 
 
 class DataFile(BaseModel):
@@ -33,12 +31,10 @@ class TestExpectations(BaseModel):
 
 
 class LLM:
-    def __init__(self, client: TinyB, api_key: Optional[str] = None):
-        self.client = client
-        user_token = CLIConfig.get_project_config().get_user_token()
-        user_client = deepcopy(client)
-        user_client.token = user_token
-        self.user_client = user_client
+    def __init__(self, user_token: str, client: TinyB, api_key: Optional[str] = None):
+        self.user_client = deepcopy(client)
+        self.user_client.token = user_token
+
         self.openai = OpenAI(api_key=api_key) if api_key else None
 
     async def _execute(self, action_fn: Callable[[], Awaitable[str]], checker_fn: Callable[[str], bool]):
@@ -75,22 +71,15 @@ class LLM:
             data=f'{{"schema": "{urllib.parse.quote(schema)}", "rows": {rows}, "context": "{prompt}"}}',
             headers={"Content-Type": "application/json"},
         )
-        return response.get("result", "")
-
-    async def create_test_commands(
-        self, pipe_content: str, pipe_params: set[str], context: str = ""
-    ) -> TestExpectations:
-        if not self.openai:
-            raise ValueError("OpenAI API key is not set")
-
-        completion = self.openai.beta.chat.completions.parse(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": create_test_calls_prompt.format(context=context)},
-                {"role": "user", "content": f"Pipe content: {pipe_content}\nPipe params: {pipe_params}"},
-            ],
-            temperature=0.2,
-            seed=42,
-            response_format=TestExpectations,
+        result = response.get("result", "")
+        return result.replace("elementAt", "arrayElement")
+
+    async def create_tests(self, pipe_content: str, pipe_params: set[str], prompt: str = "") -> TestExpectations:
+        response = await self.user_client._req(
+            "/v0/llm/create/tests",
+            method="POST",
+            data=json.dumps({"pipe_content": pipe_content, "pipe_params": list(pipe_params), "prompt": prompt}),
+            headers={"Content-Type": "application/json"},
         )
-        return completion.choices[0].message.parsed or TestExpectations(tests=[])
+        result = response.get("result", "")
+        return TestExpectations.model_validate(result)

tinybird/tb/modules/local.py

@@ -1,13 +1,14 @@
 import os
+import re
 import time
 
 import click
 
 import docker
-from tinybird.feedback_manager import FeedbackManager
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import coro
 from tinybird.tb.modules.exceptions import CLIException
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.local_common import TB_CONTAINER_NAME, TB_IMAGE_NAME, TB_LOCAL_PORT
 
 
@@ -28,9 +29,9 @@ def start_tinybird_local(
 
     if (
         pull_show_prompt
-        and click.prompt(FeedbackManager.info(message="** New version detected, download? [y/N]")).lower() == "y"
+        and click.prompt(FeedbackManager.warning(message="△ New version detected, download? [y/N]")).lower() == "y"
     ):
-        click.echo(FeedbackManager.info(message="** Downloading latest version of Tinybird local..."))
+        click.echo(FeedbackManager.info(message="* Downloading latest version of Tinybird Local..."))
         pull_required = True
 
     if pull_required:
@@ -57,27 +58,26 @@ def start_tinybird_local(
             platform="linux/amd64",
         )
 
-    click.echo(FeedbackManager.info(message="** Waiting for Tinybird local to be ready..."))
-    for attempt in range(10):
-        try:
-            run = container.exec_run("tb --no-version-warning sql 'SELECT 1 AS healthcheck' --format json").output
-            # dont parse the json as docker sometimes returns warning messages
-            # todo: rafa, make this rigth
-            if b'"healthcheck": 1' in run:
-                break
-            raise RuntimeError("Unexpected response from Tinybird")
-        except Exception:
-            if attempt == 9:  # Last attempt
-                raise CLIException("Tinybird local not ready yet. Please try again in a few seconds.")
-            time.sleep(5)  # Wait 5 seconds before retrying
+    click.echo(FeedbackManager.info(message="* Waiting for Tinybird Local to be ready..."))
+    while True:
+        health = container.attrs.get("State", {}).get("Health", {}).get("Status")
+        if health == "healthy":
+            break
+        if health == "unhealthy":
+            raise CLIException("Tinybird Local is unhealthy. Please try running `tb local restart` in a few seconds.")
 
-    click.echo(FeedbackManager.success(message="✓ All set!\n"))
+        time.sleep(5)
+
+    # Remove tinybird-local dangling images to avoid running out of disk space
+    images = docker_client.images.list(name=re.sub(r":.*$", "", TB_IMAGE_NAME), all=True, filters={"dangling": True})
+    for image in images:
+        image.remove(force=True)
 
 
 def get_docker_client():
     """Check if Docker is installed and running."""
     try:
-        client = docker.from_env()
+        client = docker.from_env()  # type: ignore
         client.ping()
         return client
     except Exception:
@@ -105,9 +105,9 @@ def remove_tinybird_local(docker_client):
 @cli.command()
 def upgrade():
     """Upgrade Tinybird CLI to the latest version"""
-    click.echo(FeedbackManager.info(message="Upgrading Tinybird CLI..."))
+    click.echo(FeedbackManager.highlight(message="» Upgrading Tinybird CLI..."))
     os.system(f"{os.getenv('HOME')}/.local/bin/uv tool upgrade tinybird")
-    click.echo(FeedbackManager.success(message="Tinybird CLI upgraded"))
+    click.echo(FeedbackManager.success(message="✓ Tinybird CLI upgraded"))
 
 
 @cli.group()
@@ -119,39 +119,40 @@ def local(ctx):
 @local.command()
 @coro
 async def stop() -> None:
-    """Stop Tinybird local"""
-    click.echo(FeedbackManager.info(message="Shutting down Tinybird local..."))
+    """Stop Tinybird Local"""
+    click.echo(FeedbackManager.highlight(message="» Shutting down Tinybird Local..."))
     docker_client = get_docker_client()
     stop_tinybird_local(docker_client)
-    click.echo(FeedbackManager.success(message="Tinybird local stopped"))
+    click.echo(FeedbackManager.success(message="✓ Tinybird Local stopped"))
 
 
 @local.command()
 @coro
 async def remove() -> None:
-    """Remove Tinybird local"""
-    click.echo(FeedbackManager.info(message="Removing Tinybird local..."))
+    """Remove Tinybird Local"""
+    click.echo(FeedbackManager.highlight(message="» Removing Tinybird Local..."))
     docker_client = get_docker_client()
     remove_tinybird_local(docker_client)
-    click.echo(FeedbackManager.success(message="Tinybird local removed"))
+    click.echo(FeedbackManager.success(message="✓ Tinybird Local removed"))
 
 
 @local.command()
 @coro
 async def start() -> None:
-    """Start Tinybird local"""
-    click.echo(FeedbackManager.info(message="Starting Tinybird local..."))
+    """Start Tinybird Local"""
+    click.echo(FeedbackManager.highlight(message="» Starting Tinybird Local..."))
     docker_client = get_docker_client()
     start_tinybird_local(docker_client)
-    click.echo(FeedbackManager.success(message="Tinybird local started"))
+    click.echo(FeedbackManager.success(message="✓ Tinybird Local is ready!"))
 
 
 @local.command()
 @coro
 async def restart() -> None:
-    """Restart Tinybird local"""
-    click.echo(FeedbackManager.info(message="Restarting Tinybird local..."))
+    """Restart Tinybird Local"""
+    click.echo(FeedbackManager.highlight(message="» Restarting Tinybird Local..."))
     docker_client = get_docker_client()
     remove_tinybird_local(docker_client)
+    click.echo(FeedbackManager.info(message="✓ Tinybird Local stopped"))
     start_tinybird_local(docker_client)
-    click.echo(FeedbackManager.success(message="Tinybird local restarted"))
+    click.echo(FeedbackManager.success(message="✓ Tinybird Local is ready!"))

tinybird/tb/modules/local_common.py

@@ -8,7 +8,6 @@ from tinybird.client import TinyB
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.exceptions import CLIException
 
-# TODO: Use the official Tinybird image once it's available 'tinybirdco/tinybird-local:latest'
 TB_IMAGE_NAME = "tinybirdco/tinybird-local:latest"
 TB_CONTAINER_NAME = "tinybird-local"
 TB_LOCAL_PORT = int(os.getenv("TB_LOCAL_PORT", 80))
@@ -40,6 +39,8 @@ async def get_tinybird_local_client(path: Optional[str] = None) -> TinyB:
             await user_client.create_workspace(ws_name, template=None)
             user_workspaces = await user_client.user_workspaces()
             ws = next((ws for ws in user_workspaces["workspaces"] if ws["name"] == ws_name), None)
+            if not ws:
+                raise CLIException(f"Workspace {ws_name} not found after creation")
 
         ws_token = ws["token"]
 

tinybird/tb/modules/login.py

@@ -9,9 +9,9 @@ from urllib.parse import urlencode
 import click
 import requests
 
-from tinybird.feedback_manager import FeedbackManager
 from tinybird.tb.modules.cli import CLIConfig, cli
 from tinybird.tb.modules.common import coro
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
 class AuthHandler(http.server.SimpleHTTPRequestHandler):
@@ -75,7 +75,7 @@ class AuthHandler(http.server.SimpleHTTPRequestHandler):
 
         if "token" in query_params:
             token = query_params["token"][0]
-            self.server.auth_callback(token)
+            self.server.auth_callback(token)  # type: ignore
             self.send_response(200)
             self.end_headers()
         else:
@@ -158,15 +158,15 @@ async def login(host: str, workspace: str):
 
         data = response.json()
         cli_config = CLIConfig.get_project_config()
-        cli_config.set_token(data["workspace_token"])
-        cli_config.set_token_for_host(data["workspace_token"], host)
-        cli_config.set_user_token(data["user_token"])
+        cli_config.set_token(data.get("workspace_token", ""))
+        cli_config.set_token_for_host(data.get("workspace_token", ""), host)
+        cli_config.set_user_token(data.get("user_token", ""))
         cli_config.set_host(host)
 
-        response = await cli_config.get_client().workspace_info()
+        ws = await cli_config.get_client().workspace_info()
         for k in ("id", "name", "user_email", "user_id", "scope"):
-            if k in response:
-                cli_config[k] = response[k]
+            if k in ws:
+                cli_config[k] = ws[k]
 
         cli_config.persist_to_file()
 

tinybird/tb/modules/mock.py

@@ -1,13 +1,14 @@
+import logging
 import os
 from pathlib import Path
 
 import click
 
-from tinybird.feedback_manager import FeedbackManager
 from tinybird.tb.modules.cli import cli
-from tinybird.tb.modules.common import CLIException, coro
+from tinybird.tb.modules.common import CLIException, check_user_token, coro
 from tinybird.tb.modules.config import CLIConfig
 from tinybird.tb.modules.datafile.fixture import build_fixture_name, persist_fixture
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.llm import LLM
 from tinybird.tb.modules.local_common import get_tinybird_local_client
 
@@ -17,8 +18,9 @@ from tinybird.tb.modules.local_common import get_tinybird_local_client
 @click.option("--rows", type=int, default=10, help="Number of events to send")
 @click.option("--prompt", type=str, default="", help="Extra context to use for data generation")
 @click.option("--folder", type=str, default=".", help="Folder where datafiles will be placed")
+@click.pass_context
 @coro
-async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
+async def mock(ctx: click.Context, datasource: str, rows: int, prompt: str, folder: str) -> None:
     """Load sample data into a Data Source.
 
     Args:
@@ -39,7 +41,6 @@ async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
         if not prompt:
             # load the prompt from the fixture.prompt file if it exists
             if prompt_path.exists():
-                click.echo(FeedbackManager.gray(message=f"Using prompt for {prompt_path}..."))
                 prompt = prompt_path.read_text()
         else:
             click.echo(FeedbackManager.gray(message=f"Overriding prompt for {datasource_name}..."))
@@ -49,15 +50,24 @@ async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
         datasource_content = datasource_path.read_text()
         config = CLIConfig.get_project_config()
         user_client = config.get_client()
-        user_client.token = config.get_user_token()
-        llm = LLM(client=user_client)
+        user_token = config.get_user_token()
+
+        try:
+            if not user_token:
+                raise CLIException("No user token found")
+            await check_user_token(ctx, token=user_token)
+        except Exception:
+            click.echo(FeedbackManager.error(message="This action requires authentication. Run 'tb login' first."))
+            return
+
+        llm = LLM(user_token=user_token, client=user_client)
         tb_client = await get_tinybird_local_client(os.path.abspath(folder))
         sql = await llm.generate_sql_sample_data(datasource_content, rows=rows, prompt=prompt)
-        if os.environ.get('TB_DEBUG', '') != '':
-            print(sql)
+        if os.environ.get("TB_DEBUG", "") != "":
+            logging.debug(sql)
         result = await tb_client.query(f"{sql} FORMAT JSON")
         data = result.get("data", [])[:rows]
-        fixture_name = build_fixture_name(datasource_path.absolute(), datasource_name, datasource_content)
+        fixture_name = build_fixture_name(datasource_path.absolute().as_posix(), datasource_name, datasource_content)
         persist_fixture(fixture_name, data)
         click.echo(FeedbackManager.success(message="✓ Done!"))
 

tinybird/tb/modules/pipe.py

@@ -13,10 +13,8 @@ import click
 import humanfriendly
 from click import Context
 
-import tinybird.context as context
 from tinybird.client import AuthNoTokenException, DoesNotExistException, TinyB
-from tinybird.config import DEFAULT_API_HOST, FeatureFlags
-from tinybird.feedback_manager import FeedbackManager
+from tinybird.config import DEFAULT_API_HOST
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import (
     coro,
@@ -24,12 +22,13 @@ from tinybird.tb.modules.common import (
     echo_safe_humanfriendly_tables_format_smart_table,
     wait_job,
 )
-from tinybird.tb.modules.datafile.build import folder_push, process_file
+from tinybird.tb.modules.datafile.build import process_file
 from tinybird.tb.modules.datafile.common import PipeNodeTypes, PipeTypes, get_name_version
 from tinybird.tb.modules.exceptions import CLIPipeException
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
-@cli.group()
+@cli.group(hidden=True)
 @click.pass_context
 def pipe(ctx):
     """Pipes commands"""
@@ -585,127 +584,6 @@ async def print_pipe(ctx: Context, pipe: str, query: str, format_: str):
         click.echo(res)
 
 
-@pipe.command(name="regression-test", short_help="Run regression tests using last requests")
-@click.option(
-    "--debug",
-    is_flag=True,
-    default=False,
-    help="Prints internal representation, can be combined with any command to get more information.",
-)
-@click.option("--only-response-times", is_flag=True, default=False, help="Checks only response times")
-@click.argument("filenames", type=click.Path(exists=True), nargs=-1, default=None)
-@click.option("--workspace_map", nargs=2, type=str, multiple=True)
-@click.option(
-    "--workspace",
-    nargs=2,
-    type=str,
-    multiple=True,
-    help="add a workspace path to the list of external workspaces, usage: --workspace name path/to/folder",
-)
-@click.option(
-    "--no-versions",
-    is_flag=True,
-    default=False,
-    help="when set, resource dependency versions are not used, it pushes the dependencies as-is",
-)
-@click.option(
-    "-l", "--limit", type=click.IntRange(0, 100), default=0, required=False, help="Number of requests to validate"
-)
-@click.option(
-    "--sample-by-params",
-    type=click.IntRange(1, 100),
-    default=1,
-    required=False,
-    help="When set, we will aggregate the pipe_stats_rt requests by extractURLParameterNames(assumeNotNull(url)) and for each combination we will take a sample of N requests",
-)
-@click.option(
-    "-m",
-    "--match",
-    multiple=True,
-    required=False,
-    help="Filter the checker requests by specific parameter. You can pass multiple parameters -m foo -m bar",
-)
-@click.option(
-    "-ff", "--failfast", is_flag=True, default=False, help="When set, the checker will exit as soon one test fails"
-)
-@click.option(
-    "--ignore-order", is_flag=True, default=False, help="When set, the checker will ignore the order of list properties"
-)
-@click.option(
-    "--validate-processed-bytes",
-    is_flag=True,
-    default=False,
-    help="When set, the checker will validate that the new version doesn't process more than 25% than the current version",
-)
-@click.option(
-    "--check-requests-from-main",
-    is_flag=True,
-    default=False,
-    help="When set, the checker will get Main Workspace requests",
-    hidden=True,
-)
-@click.option(
-    "--relative-change",
-    type=float,
-    default=0.01,
-    help="When set, the checker will validate the new version has less than this distance with the current version",
-)
-@click.pass_context
-@coro
-async def regression_test(
-    ctx: click.Context,
-    filenames: Optional[List[str]],
-    debug: bool,
-    only_response_times: bool,
-    workspace_map,
-    workspace: str,
-    no_versions: bool,
-    limit: int,
-    sample_by_params: int,
-    match: List[str],
-    failfast: bool,
-    ignore_order: bool,
-    validate_processed_bytes: bool,
-    check_requests_from_main: bool,
-    relative_change: float,
-):
-    """
-    Run regression tests on Tinybird
-    """
-
-    ignore_sql_errors = FeatureFlags.ignore_sql_errors()
-
-    context.disable_template_security_validation.set(True)
-    await folder_push(
-        create_tb_client(ctx),
-        filenames,
-        dry_run=False,
-        check=True,
-        push_deps=False,
-        debug=debug,
-        force=False,
-        populate=False,
-        upload_fixtures=False,
-        wait=False,
-        ignore_sql_errors=ignore_sql_errors,
-        skip_confirmation=False,
-        only_response_times=only_response_times,
-        workspace_map=dict(workspace_map),
-        workspace_lib_paths=workspace,
-        no_versions=no_versions,
-        run_tests=True,
-        tests_to_run=limit,
-        tests_relative_change=relative_change,
-        tests_sample_by_params=sample_by_params,
-        tests_filter_by=match,
-        tests_failfast=failfast,
-        tests_ignore_order=ignore_order,
-        tests_validate_processed_bytes=validate_processed_bytes,
-        tests_check_requests_from_branch=check_requests_from_main,
-    )
-    return
-
-
 @pipe_copy.command(name="run", short_help="Run an on-demand copy job")
 @click.argument("pipe_name_or_id")
 @click.option("--wait", is_flag=True, default=False, help="Wait for the copy job to finish")

tinybird/tb/modules/{build_shell.py → shell.py}

@@ -1,7 +1,6 @@
 import asyncio
 import concurrent.futures
 import os
-import random
 import subprocess
 import sys
 from typing import List
@@ -16,8 +15,8 @@ from prompt_toolkit.shortcuts import CompleteStyle
 from prompt_toolkit.styles import Style
 
 from tinybird.client import TinyB
-from tinybird.feedback_manager import FeedbackManager, bcolors
 from tinybird.tb.modules.exceptions import CLIException
+from tinybird.tb.modules.feedback_manager import FeedbackManager, bcolors
 from tinybird.tb.modules.table import format_table
 
 
@@ -176,7 +175,7 @@ def _(event):
         b.start_completion(select_first=False)
 
 
-class BuildShell:
+class Shell:
     def __init__(
         self,
         folder: str,
@@ -195,8 +194,7 @@ class BuildShell:
         self.endpoints = endpoints
         self.prompt_message = "\ntb > "
         self.commands = ["create", "mock", "test", "tb", "select"]
-
-        self.session = PromptSession(
+        self.session: PromptSession = PromptSession(
             completer=DynamicCompleter(self.datasources, self.shared_datasources, self.endpoints, self.pipes),
             complete_style=CompleteStyle.COLUMN,
             complete_while_typing=True,
@@ -210,7 +208,7 @@ class BuildShell:
         except Exception:
             return None
 
-    def run_shell(self):
+    def run(self):
         while True:
             try:
                 user_input = self.session.prompt(
@@ -328,9 +326,6 @@ class BuildShell:
 
 
 def print_table_formatted(res: dict, name: str):
-    rebuild_colors = [bcolors.FAIL, bcolors.OKBLUE, bcolors.WARNING, bcolors.OKGREEN, bcolors.HEADER]
-    rebuild_index = random.randint(0, len(rebuild_colors) - 1)
-    rebuild_color = rebuild_colors[rebuild_index % len(rebuild_colors)]
     data = []
     limit = 5
     for d in res["data"][:5]:
@@ -341,24 +336,17 @@ def print_table_formatted(res: dict, name: str):
     elapsed = stats.get("elapsed", 0)
     cols = len(meta)
     try:
-
-        def print_message(message: str, color=bcolors.CGREY):
-            return f"{color}{message}{bcolors.ENDC}"
-
         table = format_table(data, meta)
-        colored_char = print_message("│", rebuild_color)
-        table_with_marker = "\n".join(f"{colored_char} {line}" for line in table.split("\n"))
-        click.echo(f"\n{colored_char} {print_message('⚡', rebuild_color)} Running {name}")
-        click.echo(colored_char)
-        click.echo(table_with_marker)
-        click.echo(colored_char)
+        click.echo(FeedbackManager.highlight(message=f"\n» Running {name}\n"))
+        click.echo(table)
+        click.echo("")
         rows_read = humanfriendly.format_number(stats.get("rows_read", 0))
         bytes_read = humanfriendly.format_size(stats.get("bytes_read", 0))
         elapsed = humanfriendly.format_timespan(elapsed) if elapsed >= 1 else f"{elapsed * 1000:.2f}ms"
         stats_message = f"» {bytes_read} ({rows_read} rows x {cols} cols) in {elapsed}"
         rows_message = f"» Showing {limit} first rows" if row_count > limit else "» Showing all rows"
-        click.echo(f"{colored_char} {print_message(stats_message, bcolors.OKGREEN)}")
-        click.echo(f"{colored_char} {print_message(rows_message, bcolors.CGREY)}")
+        click.echo(FeedbackManager.success(message=stats_message))
+        click.echo(FeedbackManager.gray(message=rows_message))
     except ValueError as exc:
         if str(exc) == "max() arg is an empty sequence":
             click.echo("------------")