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

tinybird/tb/modules/build_common.py

@@ -1,18 +1,22 @@
 import json
 import logging
+import re
 import time
+from copy import deepcopy
 from pathlib import Path
-from typing import Optional
+from typing import Any, Optional
 from urllib.parse import urlencode, urljoin
 
 import click
 import requests
 
+from tinybird.datafile.parse_datasource import parse_datasource
 from tinybird.tb.client import TinyB
 from tinybird.tb.modules.common import push_data, sys_exit
 from tinybird.tb.modules.datafile.fixture import FixtureExtension, get_fixture_dir, persist_fixture
 from tinybird.tb.modules.dev_server import BuildStatus
 from tinybird.tb.modules.feedback_manager import FeedbackManager
+from tinybird.tb.modules.local_common import get_local_tokens
 from tinybird.tb.modules.project import Project
 from tinybird.tb.modules.shell import print_table_formatted
 
@@ -21,14 +25,26 @@ def process(
     project: Project,
     tb_client: TinyB,
     watch: bool,
+    config: dict[str, Any],
     file_changed: Optional[str] = None,
     diff: Optional[str] = None,
     silent: bool = False,
     build_status: Optional[BuildStatus] = None,
     exit_on_error: bool = True,
     load_fixtures: bool = True,
+    project_with_vendors: Optional[Project] = None,
+    is_branch: bool = False,
 ) -> Optional[str]:
     time_start = time.time()
+
+    # Build vendored workspaces before build
+    if not project_with_vendors and not is_branch:
+        build_vendored_workspaces(project=project, tb_client=tb_client, config=config)
+
+    # Ensure SHARED_WITH workspaces exist before build
+    if not is_branch:
+        build_shared_with_workspaces(project=project, tb_client=tb_client, config=config)
+
     build_failed = False
     build_error: Optional[str] = None
     build_result: Optional[bool] = None
@@ -37,7 +53,7 @@ def process(
             return build_status.error
         else:
             build_status.building = True
-    if file_changed and (file_changed.endswith(FixtureExtension.NDJSON) or file_changed.endswith(FixtureExtension.CSV)):
+    if file_changed and file_changed.endswith((FixtureExtension.NDJSON, FixtureExtension.CSV)):
         rebuild_fixture(project, tb_client, file_changed)
         if build_status:
             build_status.building = False
@@ -47,13 +63,19 @@ def process(
         if build_status:
             build_status.building = False
             build_status.error = None
-    elif file_changed and (file_changed.endswith(".env.local") or file_changed.endswith(".env")):
+    elif file_changed and file_changed.endswith((".env.local", ".env")):
         if build_status:
             build_status.building = False
             build_status.error = None
     else:
         try:
-            build_result = build_project(project, tb_client, silent, load_fixtures)
+            build_result = build_project(
+                project,
+                tb_client,
+                silent,
+                load_fixtures,
+                project_with_vendors=project_with_vendors,
+            )
             if build_status:
                 build_status.building = False
                 build_status.error = None
@@ -82,12 +104,12 @@ def process(
             build_status.error = build_error
             build_status.building = False
         return build_error
-    else:
-        if not silent:
-            if build_result == False:  # noqa: E712
-                click.echo(FeedbackManager.info(message="No changes. Build skipped."))
-            else:
-                click.echo(FeedbackManager.success(message=f"\n✓ {rebuild_str} completed in {elapsed_time:.1f}s"))
+
+    if not silent:
+        if build_result == False:  # noqa: E712
+            click.echo(FeedbackManager.info(message="No changes. Build skipped."))
+        else:
+            click.echo(FeedbackManager.success(message=f"\n✓ {rebuild_str} completed in {elapsed_time:.1f}s"))
 
     return None
 
@@ -182,7 +204,11 @@ def show_data(tb_client: TinyB, filename: str, diff: Optional[str] = None):
 
 
 def build_project(
-    project: Project, tb_client: TinyB, silent: bool = False, load_fixtures: bool = True
+    project: Project,
+    tb_client: TinyB,
+    silent: bool = False,
+    load_fixtures: bool = True,
+    project_with_vendors: Optional[Project] = None,
 ) -> Optional[bool]:
     MULTIPART_BOUNDARY_DATA_PROJECT = "data_project://"
     DATAFILE_TYPE_TO_CONTENT_TYPE = {
@@ -194,6 +220,7 @@ def build_project(
     logging.debug(TINYBIRD_API_URL)
     TINYBIRD_API_KEY = tb_client.token
     error: Optional[str] = None
+
     try:
         files = [
             ("context://", ("cli-version", "1.0.0", "text/plain")),
@@ -208,9 +235,15 @@ def build_project(
             relative_path = Path(file_path).relative_to(project_path).as_posix()
             with open(file_path, "rb") as fd:
                 content_type = DATAFILE_TYPE_TO_CONTENT_TYPE.get(Path(file_path).suffix, "application/unknown")
-                files.append(
-                    (MULTIPART_BOUNDARY_DATA_PROJECT, (relative_path, fd.read().decode("utf-8"), content_type))
-                )
+                content = fd.read().decode("utf-8")
+                if project_with_vendors:
+                    # Replace 'SHARED_WITH' and everything that comes after, including new lines, with 'SHARED_WITH Tinybird_Local_Test_'
+                    content = replace_shared_with(
+                        content,
+                        [project_with_vendors.workspace_name],
+                    )
+
+                files.append((MULTIPART_BOUNDARY_DATA_PROJECT, (relative_path, content, content_type)))
         HEADERS = {"Authorization": f"Bearer {TINYBIRD_API_KEY}"}
 
         r = requests.post(TINYBIRD_API_URL, files=files, headers=HEADERS)
@@ -249,17 +282,16 @@ def build_project(
             )
             if no_changes:
                 return False
-            else:
-                if not silent:
-                    echo_changes(project, new_datasources, ".datasource", "created")
-                    echo_changes(project, changed_datasources, ".datasource", "changed")
-                    echo_changes(project, deleted_datasources, ".datasource", "deleted")
-                    echo_changes(project, new_pipes, ".pipe", "created")
-                    echo_changes(project, changed_pipes, ".pipe", "changed")
-                    echo_changes(project, deleted_pipes, ".pipe", "deleted")
-                    echo_changes(project, new_connections, ".connection", "created")
-                    echo_changes(project, changed_connections, ".connection", "changed")
-                    echo_changes(project, deleted_connections, ".connection", "deleted")
+            if not silent:
+                echo_changes(project, new_datasources, ".datasource", "created")
+                echo_changes(project, changed_datasources, ".datasource", "changed")
+                echo_changes(project, deleted_datasources, ".datasource", "deleted")
+                echo_changes(project, new_pipes, ".pipe", "created")
+                echo_changes(project, changed_pipes, ".pipe", "changed")
+                echo_changes(project, deleted_pipes, ".pipe", "deleted")
+                echo_changes(project, new_connections, ".connection", "created")
+                echo_changes(project, changed_connections, ".connection", "changed")
+                echo_changes(project, deleted_connections, ".connection", "deleted")
             if load_fixtures:
                 try:
                     for filename in project_files:
@@ -320,3 +352,155 @@ def echo_changes(project: Project, changes: list[str], extension: str, status: s
         if path_str:
             path_str = path_str.replace(f"{project.folder}/", "")
             click.echo(FeedbackManager.info(message=f"✓ {path_str} {status}"))
+
+
+def find_workspace_or_create(user_client: TinyB, workspace_name: str) -> Optional[str]:
+    # Get a client scoped to the vendored workspace using the user token
+    ws_token = None
+    org_id = None
+    try:
+        # Fetch org id and workspaces with tokens
+        info = user_client.user_workspaces_with_organization(version="v1")
+        org_id = info.get("organization_id")
+        workspaces = info.get("workspaces", [])
+        found = next((w for w in workspaces if w.get("name") == workspace_name), None)
+        if found:
+            ws_token = found.get("token")
+        # If still not found, try the generic listing
+        if not ws_token:
+            workspaces_full = user_client.user_workspaces_and_branches(version="v1")
+            created_ws = next(
+                (w for w in workspaces_full.get("workspaces", []) if w.get("name") == workspace_name), None
+            )
+            if created_ws:
+                ws_token = created_ws.get("token")
+    except Exception:
+        ws_token = None
+
+    # If workspace doesn't exist, try to create it and fetch its token
+    if not ws_token:
+        try:
+            user_client.create_workspace(workspace_name, assign_to_organization_id=org_id, version="v1")
+            # Fetch token for newly created workspace
+            info_after = user_client.user_workspaces_and_branches(version="v1")
+            created = next((w for w in info_after.get("workspaces", []) if w.get("name") == workspace_name), None)
+            ws_token = created.get("token") if created else None
+        except Exception as e:
+            click.echo(
+                FeedbackManager.warning(
+                    message=(f"Skipping vendored workspace '{workspace_name}': unable to create or resolve token ({e})")
+                )
+            )
+
+    return ws_token
+
+
+def build_vendored_workspaces(project: Project, tb_client: TinyB, config: dict[str, Any]) -> None:
+    """Build each vendored workspace under project.vendor_path if present.
+
+    Directory structure expected: vendor/<workspace_name>/<data_project_inside>
+    Each top-level directory under vendor is treated as a separate workspace
+    whose project files will be built using that workspace's token.
+    """
+    try:
+        vendor_root = Path(project.vendor_path)
+
+        if not vendor_root.exists() or not vendor_root.is_dir():
+            return
+
+        tokens = get_local_tokens()
+        user_token = tokens["user_token"]
+        user_client = deepcopy(tb_client)
+        user_client.token = user_token
+
+        # Iterate over vendored workspace folders
+        for ws_dir in sorted([p for p in vendor_root.iterdir() if p.is_dir()]):
+            workspace_name = ws_dir.name
+            ws_token = find_workspace_or_create(user_client, workspace_name)
+
+            if not ws_token:
+                click.echo(
+                    FeedbackManager.warning(
+                        message=f"Skipping vendored workspace '{workspace_name}': could not resolve token after creation"
+                    )
+                )
+                continue
+
+            # Build using a client scoped to the vendor workspace token
+            vendor_client = deepcopy(tb_client)
+            vendor_client.token = ws_token
+            vendor_project = Project(folder=str(ws_dir), workspace_name=workspace_name, max_depth=project.max_depth)
+            workspace_info = tb_client.workspace_info(version="v1")
+            project.workspace_name = workspace_info.get("name", "")
+            # Do not exit on error to allow main project to continue
+            process(
+                project=vendor_project,
+                tb_client=vendor_client,
+                watch=False,
+                silent=True,
+                exit_on_error=True,
+                load_fixtures=True,
+                config=config,
+                project_with_vendors=project,
+            )
+    except Exception as e:
+        # Never break the main build due to vendored build errors
+        click.echo(FeedbackManager.error_exception(error=e))
+
+
+def build_shared_with_workspaces(project: Project, tb_client: TinyB, config: dict[str, Any]) -> None:
+    """Scan project for .datasource files and ensure SHARED_WITH workspaces exist."""
+
+    try:
+        # Gather SHARED_WITH workspace names from all .datasource files
+        datasource_files = project.get_datasource_files()
+        shared_ws_names = set()
+
+        for filename in datasource_files:
+            try:
+                doc = parse_datasource(filename).datafile
+                for ws_name in doc.shared_with or []:
+                    shared_ws_names.add(ws_name)
+            except Exception:
+                # Ignore parse errors here; they'll be handled during the main process()
+                continue
+
+        if not shared_ws_names:
+            return
+
+        # Need a user token to list/create workspaces
+        tokens = get_local_tokens()
+        user_token = tokens.get("user_token")
+        if not user_token:
+            click.echo(FeedbackManager.info_skipping_shared_with_entry())
+            return
+
+        user_client = deepcopy(tb_client)
+        user_client.token = user_token
+
+        # Ensure each SHARED_WITH workspace exists
+        for ws_name in sorted(shared_ws_names):
+            find_workspace_or_create(user_client, ws_name)
+    except Exception as e:
+        click.echo(FeedbackManager.error_exception(error=e))
+
+
+def replace_shared_with(text: str, new_workspaces: list[str]) -> str:
+    replacement = ", ".join(new_workspaces)
+
+    # 1) Formato multilinea:
+    # SHARED_WITH >
+    #     workspace1, workspace2
+    #
+    # Solo sustituimos la LÍNEA de workspaces (grupo 3), no usamos DOTALL.
+    pat_multiline = re.compile(r"(?m)^(SHARED_WITH\s*>\s*)\n([ \t]*)([^\n]*)$")
+    if pat_multiline.search(text):
+        return pat_multiline.sub(lambda m: f"{m.group(1)}\n{m.group(2)}{replacement}", text)
+
+    # 2) Formato inline:
+    # SHARED_WITH workspace1, workspace2
+    pat_inline = re.compile(r"(?m)^(SHARED_WITH\s+)([^\n]*)$")
+    if pat_inline.search(text):
+        return pat_inline.sub(lambda m: f"{m.group(1)}{replacement}", text)
+
+    return text

tinybird/tb/modules/cli.py

@@ -2,7 +2,6 @@
 #
 # - If it makes sense and only when strictly necessary, you can create utility functions in this file.
 # - But please, **do not** interleave utility functions and command definitions.
-
 import json
 import logging
 import os
@@ -11,9 +10,11 @@ import sys
 from os import environ, getcwd
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from urllib.parse import urlencode
 
 import click
 import humanfriendly
+import requests
 from click import Context
 
 from tinybird.tb import __cli__
@@ -21,7 +22,9 @@ from tinybird.tb.check_pypi import CheckPypi
 from tinybird.tb.client import (
     AuthException,
     AuthNoTokenException,
+    TinyB,
 )
+from tinybird.tb.config import get_clickhouse_host
 from tinybird.tb.modules.agent import run_agent
 from tinybird.tb.modules.common import (
     CatchAuthExceptions,
@@ -36,8 +39,10 @@ from tinybird.tb.modules.common import (
 from tinybird.tb.modules.config import CURRENT_VERSION, CLIConfig
 from tinybird.tb.modules.datafile.build import build_graph
 from tinybird.tb.modules.datafile.pull import folder_pull
+from tinybird.tb.modules.exceptions import CLIChException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.local_common import get_tinybird_local_client
+from tinybird.tb.modules.login_common import check_current_folder_in_sessions
 from tinybird.tb.modules.project import Project
 
 __old_click_echo = click.echo
@@ -72,6 +77,7 @@ VERSION = f"{__cli__.__version__} (rev {__cli__.__revision__})"
 )
 @click.option("--show-tokens", is_flag=True, default=False, help="Enable the output of tokens.")
 @click.option("--cloud/--local", is_flag=True, default=False, help="Run against cloud or local.")
+@click.option("--branch", help="Run against a branch.")
 @click.option("--staging", is_flag=True, default=False, help="Run against a staging deployment.")
 @click.option(
     "--output", type=click.Choice(["human", "json", "csv"], case_sensitive=False), default="human", help="Output format"
@@ -99,6 +105,7 @@ def cli(
     version_warning: bool,
     show_tokens: bool,
     cloud: bool,
+    branch: Optional[str],
     staging: bool,
     output: str,
     max_depth: int,
@@ -106,17 +113,16 @@ def cli(
     prompt: Optional[str] = None,
 ) -> None:
     """
-    Use Tinybird Code to interact with your project.
+    Run just `tb` to use Tinybird Code to interact with your project.
     """
 
     # We need to unpatch for our tests not to break
     if output != "human":
         __hide_click_output()
+    elif show_tokens or not cloud or ctx.invoked_subcommand == "build":
+        __unpatch_click_output()
     else:
-        if show_tokens or not cloud or ctx.invoked_subcommand == "build":
-            __unpatch_click_output()
-        else:
-            __patch_click_output()
+        __patch_click_output()
 
     if getenv_bool("TB_DISABLE_SSL_CHECKS", False):
         click.echo(FeedbackManager.warning_disabled_ssl_checks())
@@ -197,17 +203,36 @@ def cli(
         return
 
     ctx.ensure_object(dict)["project"] = project
-    client = create_ctx_client(ctx, config, cloud, staging, project=project, show_warnings=version_warning)
+    client = create_ctx_client(
+        ctx,
+        config,
+        cloud or bool(branch),
+        staging,
+        project=project,
+        show_warnings=version_warning,
+        branch=branch,
+    )
 
     if client:
         ctx.ensure_object(dict)["client"] = client
 
-    ctx.ensure_object(dict)["env"] = get_target_env(cloud)
+    ctx.ensure_object(dict)["env"] = get_target_env(cloud, branch)
+    ctx.ensure_object(dict)["branch"] = branch
     ctx.ensure_object(dict)["output"] = output
 
+    # Check if current folder is tracked from previous sessions
+    check_current_folder_in_sessions(ctx)
+
     is_prompt_mode = prompt is not None
 
     if is_agent_mode or is_prompt_mode:
+        if any(arg in sys.argv for arg in ["--cloud", "--local", "--branch"]):
+            raise CLIException(
+                FeedbackManager.error(
+                    message="Tinybird Code does not support --cloud, --local or --branch flags. It will choose the correct environment based on your prompts."
+                )
+            )
+
         run_agent(config, project, dangerously_skip_permissions, prompt=prompt)
 
 
@@ -298,7 +323,7 @@ def sql(
             )
 
             query = ""
-            for _, elem in dependencies_graph.to_run.items():
+            for elem in dependencies_graph.to_run.values():
                 for _node in elem["nodes"]:
                     if _node["params"]["name"].lower() == node.lower():
                         query = "".join(_node["sql"])
@@ -336,6 +361,88 @@ def sql(
         click.echo(FeedbackManager.info_no_rows())
 
 
+@cli.command(
+    name="ch",
+    context_settings=dict(
+        ignore_unknown_options=True,
+        allow_extra_args=True,
+    ),
+)
+@click.option(
+    "--query",
+    type=str,
+    default=None,
+    required=False,
+    help="The query to run against ClickHouse.",
+)
+@click.option(
+    "--user",
+    required=False,
+    help="User field is not used for authentication but helps identify the connection.",
+)
+@click.option(
+    "--password",
+    required=False,
+    help="Your Tinybird Auth Token. If not provided, the token will be your current workspace token.",
+)
+@click.option(
+    "-m",
+    "--multiline",
+    is_flag=True,
+    default=False,
+    help="Enable multiline mode - read the query from multiple lines until a semicolon.",
+)
+@click.pass_context
+def ch(ctx: Context, query: str, user: Optional[str], password: Optional[str], multiline: bool) -> None:
+    """Run a query against ClickHouse native HTTP interface."""
+    try:
+        query_arg = next((arg for arg in ctx.args if not arg.startswith("--param_")), None)
+        if query_arg and not query:
+            query = query_arg
+
+        if not query and not sys.stdin.isatty():  # Check if there's piped input
+            query = sys.stdin.read().strip()
+
+        if not query:
+            click.echo(FeedbackManager.warning(message="Nothing to do. No query provided"))
+            return
+
+        if multiline:
+            queries = [query.strip() for query in query.split(";") if query.strip()]
+        else:
+            queries = [query]
+
+        client: TinyB = ctx.ensure_object(dict)["client"]
+        config = ctx.ensure_object(dict)["config"]
+        password = password or client.token
+        user = user or config.get("name", None)
+        ch_host = get_clickhouse_host(client.host)
+        headers = {"X-ClickHouse-Key": password}
+        if user:
+            headers["X-ClickHouse-User"] = user
+
+        params = {}
+
+        for param in ctx.args:
+            if param.startswith("--param_"):
+                param_name = param.split("=")[0].replace("--", "")
+                param_value = param.split("=")[1]
+                params[param_name] = param_value
+
+        for query in queries:
+            query_params = {**params, "query": query}
+            url = f"{ch_host}?{urlencode(query_params)}"
+            res = requests.get(url=url, headers=headers)
+
+            if res.status_code != 200:
+                raise Exception(res.text)
+
+            click.echo(res.text)
+
+    except Exception as e:
+        raise CLIChException(FeedbackManager.error(message=str(e)))
+
+
 def __patch_click_output():
     import re
 
@@ -392,7 +499,13 @@ def __hide_click_output() -> None:
 
 
 def create_ctx_client(
-    ctx: Context, config: Dict[str, Any], cloud: bool, staging: bool, project: Project, show_warnings: bool = True
+    ctx: Context,
+    config: Dict[str, Any],
+    cloud: bool,
+    staging: bool,
+    project: Project,
+    show_warnings: bool = True,
+    branch: Optional[str] = None,
 ):
     commands_without_ctx_client = [
         "auth",
@@ -402,11 +515,11 @@ def create_ctx_client(
         "logout",
         "update",
         "upgrade",
-        "create",
         "info",
         "tag",
         "push",
         "branch",
+        "environment",
         "diff",
         "fmt",
         "init",
@@ -415,8 +528,8 @@ def create_ctx_client(
     if not command or command in commands_without_ctx_client:
         return None
 
-    commands_always_cloud = ["infra"]
-    commands_always_local = ["build", "dev"]
+    commands_always_cloud = ["infra", "branch", "environment"]
+    commands_always_local = ["create"]
     command_always_test = ["test"]
 
     if (
@@ -439,7 +552,7 @@ def create_ctx_client(
         if method and show_warnings:
             click.echo(FeedbackManager.gray(message=f"Authentication method: {method}"))
 
-        return _get_tb_client(config.get("token", ""), config["host"], staging=staging)
+        return _get_tb_client(config.get("token", ""), config["host"], staging=staging, branch=branch)
     local = command in commands_always_local
     test = command in command_always_test
     if show_warnings and not local and command not in commands_always_local and command:
@@ -447,8 +560,8 @@ def create_ctx_client(
     return get_tinybird_local_client(config, test=test, staging=staging)
 
 
-def get_target_env(cloud: bool) -> str:
-    if cloud:
+def get_target_env(cloud: bool, branch: Optional[str]) -> str:
+    if cloud or bool(branch):
         return "cloud"
     return "local"