tinybird 0.0.1.dev18__py3-none-any.whl → 0.0.1.dev20__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

@@ -37,7 +37,8 @@ class LLM:
         self.client = client
         user_token = CLIConfig.get_project_config().get_user_token()
         user_client = deepcopy(client)
-        user_client.token = user_token
+        if user_token:
+            user_client.token = user_token
         self.user_client = user_client
         self.openai = OpenAI(api_key=api_key) if api_key else None
 
@@ -75,10 +76,11 @@ class LLM:
             data=f'{{"schema": "{urllib.parse.quote(schema)}", "rows": {rows}, "context": "{prompt}"}}',
             headers={"Content-Type": "application/json"},
         )
-        return response.get("result", "")
+        result = response.get("result", "")
+        return result.replace("elementAt", "arrayElement")
 
     async def create_test_commands(
-        self, pipe_content: str, pipe_params: set[str], context: str = ""
+        self, pipe_content: str, pipe_params: set[str], context: Optional[str] = None
     ) -> TestExpectations:
         if not self.openai:
             raise ValueError("OpenAI API key is not set")
@@ -86,7 +88,7 @@ class LLM:
         completion = self.openai.beta.chat.completions.parse(
             model="gpt-4o",
             messages=[
-                {"role": "system", "content": create_test_calls_prompt.format(context=context)},
+                {"role": "system", "content": create_test_calls_prompt.format(context=context or "")},
                 {"role": "user", "content": f"Pipe content: {pipe_content}\nPipe params: {pipe_params}"},
             ],
             temperature=0.2,

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,7 +58,7 @@ def start_tinybird_local(
             platform="linux/amd64",
         )
 
-    click.echo(FeedbackManager.info(message="** Waiting for Tinybird local to be ready..."))
+    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
@@ -68,16 +69,19 @@ def start_tinybird_local(
             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.")
+                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.success(message="✓ All set!\n"))
+    # 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 +109,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 +123,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,8 +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):
@@ -74,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:
@@ -115,11 +116,11 @@ def start_server(auth_callback):
     "--workspace",
     help="Set the workspace to authenticate to. If not set, the default workspace will be used.",
 )
-def login(host: str, workspace: str):
+@coro
+async def login(host: str, workspace: str):
     """Authenticate via browser."""
     auth_event = threading.Event()
     auth_code = [None]  # Using a list to store the code, as it's mutable
-    config = CLIConfig.get_project_config()
     host = host or "https://api.tinybird.co"
 
     def auth_callback(code):
@@ -150,19 +151,25 @@ def login(host: str, workspace: str):
         params = {}
         if workspace:
             params["workspace_id"] = workspace
-        response = requests.get(
+        response = requests.get(  # noqa: ASYNC210
             f"{host}/v0/user/tokens?{urlencode(params)}",
             headers={"Authorization": f"Bearer {auth_code[0]}"},
         )
+
         data = response.json()
         cli_config = CLIConfig.get_project_config()
-        workspace_token = data["workspace_token"]
-        user_token = data["user_token"]
-        cli_config.set_token(workspace_token)
-        cli_config.set_token_for_host(workspace_token, host)
-        cli_config.set_user_token(user_token)
-        config.set_host(host)
+        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)
+
+        ws = await cli_config.get_client().workspace_info()
+        for k in ("id", "name", "user_email", "user_id", "scope"):
+            if k in ws:
+                cli_config[k] = ws[k]
+
         cli_config.persist_to_file()
+
         click.echo(FeedbackManager.success(message="✓ Authentication successful!"))
     else:
         click.echo(FeedbackManager.error(message="Authentication failed or timed out."))

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:
@@ -49,15 +51,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()
+        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(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")