tinybird 0.0.1.dev28__py3-none-any.whl → 0.0.1.dev30__py3-none-any.whl

tinybird/tb/modules/common.py

@@ -175,7 +175,7 @@ def generate_datafile(
     if not f.exists() or force:
         with open(f"{f}", "w") as ds_file:
             ds_file.write(datafile)
-        click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder)))
+        click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder or ".")))
 
         if data and (base / "fixtures").exists():
             # Generating a fixture for Parquet files is not so trivial, since Parquet format
@@ -1177,31 +1177,6 @@ async def print_current_workspace(config: CLIConfig) -> None:
     echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
 
 
-async def print_current_branch(config: CLIConfig) -> None:
-    _ = await try_update_config_with_remote(config, only_if_needed=True)
-
-    response = await config.get_client().user_workspaces_and_branches()
-
-    columns = ["name", "id", "workspace"]
-    table = []
-
-    for workspace in response["workspaces"]:
-        if config["id"] == workspace["id"]:
-            click.echo(FeedbackManager.info_current_branch())
-            if workspace.get("is_branch"):
-                name = workspace["name"]
-                main_workspace = await get_current_main_workspace(config)
-                assert isinstance(main_workspace, dict)
-                main_name = main_workspace["name"]
-            else:
-                name = MAIN_BRANCH
-                main_name = workspace["name"]
-            table.append([name, workspace["id"], main_name])
-            break
-
-    echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
-
-
 class ConnectionReplacements:
     _PARAMS_REPLACEMENTS: Dict[str, Dict[str, str]] = {
         "s3": {

tinybird/tb/modules/config.py

@@ -323,14 +323,6 @@ class CLIConfig:
         path: str = os.path.join(working_dir, ".tinyb")
         return CLIConfig(path, parent=CLIConfig.get_global_config())
 
-    @staticmethod
-    def get_llm_config(working_dir: Optional[str] = None) -> Dict[str, Any]:
-        return (
-            CLIConfig.get_project_config(working_dir)
-            .get("llms", {})
-            .get("openai", {"model": "gpt-4o-mini", "api_key": None})
-        )
-
     @staticmethod
     def reset() -> None:
         CLIConfig._global = None

tinybird/tb/modules/create.py

@@ -1,4 +1,5 @@
 import os
+import re
 from os import getcwd
 from pathlib import Path
 from typing import Optional
@@ -260,7 +261,25 @@ def init_git(folder: str):
 
 
 def generate_pipe_file(name: str, content: str, folder: str):
-    base = Path(folder) / "endpoints"
+    def is_copy(content: str) -> bool:
+        return re.search(r"TYPE copy", content, re.IGNORECASE) is not None
+
+    def is_materialization(content: str) -> bool:
+        return re.search(r"TYPE materialized", content, re.IGNORECASE) is not None
+
+    def is_sink(content: str) -> bool:
+        return re.search(r"TYPE sink", content, re.IGNORECASE) is not None
+
+    if is_copy(content):
+        pathname = "copies"
+    elif is_materialization(content):
+        pathname = "materializations"
+    elif is_sink(content):
+        pathname = "sinks"
+    else:
+        pathname = "endpoints"
+
+    base = Path(folder) / pathname
     if not base.exists():
         base = Path()
     f = base / (f"{name}.pipe")

tinybird/tb/modules/datafile/build.py

@@ -768,13 +768,13 @@ async def process(
             raise click.ClickException(FeedbackManager.error_forkdownstream_pipes_with_engine(pipe=resource_name))
 
         to_run[resource_name] = r
-        file_deps = r.get("deps", [])
+        file_deps: List[str] = r.get("deps", [])
         deps += file_deps
         # calculate and look for deps
         dep_list = []
         for x in file_deps:
             if x not in INTERNAL_TABLES or is_internal:
-                f, ds = find_file_by_name(dir_path, x, verbose, vendor_paths=vendor_paths, resource=r)
+                f, ds = find_file_by_name(dir_path or ".", x, verbose, vendor_paths=vendor_paths, resource=r)
                 if f:
                     dep_list.append(f.rsplit(".", 1)[0])
                 if ds:

tinybird/tb/modules/datafile/build_pipe.py

@@ -266,7 +266,19 @@ async def new_pipe(
 
     if data.get("type") == "endpoint":
         token = tb_client.token
-        click.echo(f"""** => Test endpoint with:\n** $ curl {host}/v0/pipes/{p["name"]}.json?token={token}""")
+        try:
+            example_params = {
+                "format": "json",
+                "pipe": p["name"],
+                "q": "",
+                "token": token,
+            }
+            endpoint_url = await tb_client._req(f"/examples/query.http?{urlencode(example_params)}")
+            if endpoint_url:
+                endpoint_url = endpoint_url.replace("http://localhost:8001", host)
+                click.echo(f"""** => Test endpoint with:\n** $ curl {endpoint_url}""")
+        except Exception:
+            pass
 
 
 async def get_token_from_main_branch(branch_tb_client: TinyB) -> Optional[str]:

tinybird/tb/modules/datasource.py

@@ -542,7 +542,7 @@ async def datasource_share(ctx: Context, datasource_name: str, workspace_name_or
     """Share a datasource"""
 
     config = CLIConfig.get_project_config()
-    client = config.get_client()
+    client: TinyB = ctx.ensure_object(dict)["client"]
     host = config.get_host() or CLIConfig.DEFAULTS["host"]
     ui_host = get_display_host(host)
 

tinybird/tb/modules/deploy.py

@@ -3,7 +3,7 @@ import json
 import logging
 import time
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
 import click
 import requests
@@ -24,8 +24,8 @@ def project_files(project_path: Path) -> List[str]:
     return project_files
 
 
-def promote_deployment(host: str, headers: dict) -> None:
-    TINYBIRD_API_URL = host + "/v1/deployments"
+def promote_deployment(host: Optional[str], headers: dict) -> None:
+    TINYBIRD_API_URL = f"{host}/v1/deployments"
     r = requests.get(TINYBIRD_API_URL, headers=headers)
     result = r.json()
     logging.debug(json.dumps(result, indent=2))
@@ -50,14 +50,14 @@ def promote_deployment(host: str, headers: dict) -> None:
     else:
         click.echo(FeedbackManager.success(message="Promoting deployment"))
 
-        TINYBIRD_API_URL = host + f"/v1/deployments/{candidate_deployment.get('id')}/set-live"
+        TINYBIRD_API_URL = f"{host}/v1/deployments/{candidate_deployment.get('id')}/set-live"
         r = requests.post(TINYBIRD_API_URL, headers=headers)
         result = r.json()
         logging.debug(json.dumps(result, indent=2))
 
     click.echo(FeedbackManager.success(message="Removing old deployment"))
 
-    TINYBIRD_API_URL = host + f"/v1/deployments/{last_deployment.get('id')}"
+    TINYBIRD_API_URL = f"{host}/v1/deployments/{last_deployment.get('id')}"
     r = requests.delete(TINYBIRD_API_URL, headers=headers)
     result = r.json()
     logging.debug(json.dumps(result, indent=2))
@@ -65,8 +65,8 @@ def promote_deployment(host: str, headers: dict) -> None:
     click.echo(FeedbackManager.success(message="Deployment promoted successfully"))
 
 
-def rollback_deployment(host: str, headers: dict) -> None:
-    TINYBIRD_API_URL = host + "/v1/deployments"
+def rollback_deployment(host: Optional[str], headers: dict) -> None:
+    TINYBIRD_API_URL = f"{host}/v1/deployments"
     r = requests.get(TINYBIRD_API_URL, headers=headers)
     result = r.json()
     logging.debug(json.dumps(result, indent=2))
@@ -91,14 +91,14 @@ def rollback_deployment(host: str, headers: dict) -> None:
     else:
         click.echo(FeedbackManager.success(message="Promoting previous deployment"))
 
-        TINYBIRD_API_URL = host + f"/v1/deployments/{previous_deployment.get('id')}/set-live"
+        TINYBIRD_API_URL = f"{host}/v1/deployments/{previous_deployment.get('id')}/set-live"
         r = requests.post(TINYBIRD_API_URL, headers=headers)
         result = r.json()
         logging.debug(json.dumps(result, indent=2))
 
     click.echo(FeedbackManager.success(message="Removing current deployment"))
 
-    TINYBIRD_API_URL = host + f"/v1/deployments/{current_deployment.get('id')}"
+    TINYBIRD_API_URL = f"{host}/v1/deployments/{current_deployment.get('id')}"
     r = requests.delete(TINYBIRD_API_URL, headers=headers)
     result = r.json()
     logging.debug(json.dumps(result, indent=2))
@@ -106,7 +106,15 @@ def rollback_deployment(host: str, headers: dict) -> None:
     click.echo(FeedbackManager.success(message="Deployment rolled back successfully"))
 
 
-@cli.command()
+@cli.group(name="deploy")
+def deploy_group() -> None:
+    """
+    Deploy commands.
+    """
+    pass
+
+
+@deploy_group.command(name="create")
 @click.argument("project_path", type=click.Path(exists=True), default=Path.cwd())
 @click.option(
     "--wait/--no-wait",
@@ -120,7 +128,7 @@ def rollback_deployment(host: str, headers: dict) -> None:
     default=False,
     help="Auto-promote the deployment. Only works if --wait is enabled. Disabled by default.",
 )
-def deploy(project_path: Path, wait: bool, auto: bool) -> None:
+def create(project_path: Path, wait: bool, auto: bool) -> None:
     """
     Validate and deploy the project server side.
     """
@@ -133,7 +141,7 @@ def deploy(project_path: Path, wait: bool, auto: bool) -> None:
     }
 
     config = CLIConfig.get_project_config(str(project_path))
-    TINYBIRD_API_URL = (config.get_host() or "") + "/v1/deploy"
+    TINYBIRD_API_URL = f"{config.get_host()}/v1/deploy"
     TINYBIRD_API_KEY = config.get_token()
 
     files = [
@@ -178,7 +186,7 @@ def deploy(project_path: Path, wait: bool, auto: bool) -> None:
     if deployment and wait:
         while deployment.get("status") != "data_ready":
             time.sleep(5)
-            TINYBIRD_API_URL = (config.get_host() or "") + f"/v1/deployments/{deployment.get('id')}"
+            TINYBIRD_API_URL = f"{config.get_host()}/v1/deployments/{deployment.get('id')}"
             r = requests.get(TINYBIRD_API_URL, headers=HEADERS)
             result = r.json()
             deployment = result.get("deployment")
@@ -189,28 +197,20 @@ def deploy(project_path: Path, wait: bool, auto: bool) -> None:
         click.echo(FeedbackManager.success(message="Deployment is ready"))
 
         if auto:
-            promote_deployment((config.get_host() or ""), HEADERS)
-
-
-@cli.group(name="releases")
-def releases_group() -> None:
-    """
-    Release commands.
-    """
-    pass
+            promote_deployment(config.get_host(), HEADERS)
 
 
-@releases_group.command(name="list")
+@deploy_group.command(name="list")
 @click.argument("project_path", type=click.Path(exists=True), default=Path.cwd())
-def release_list(project_path: Path) -> None:
+def deploy_list(project_path: Path) -> None:
     """
-    List all the releases you have in the project.
+    List all the deployments you have in the project.
     """
     config = CLIConfig.get_project_config(str(project_path))
 
     TINYBIRD_API_KEY = config.get_token()
     HEADERS = {"Authorization": f"Bearer {TINYBIRD_API_KEY}"}
-    TINYBIRD_API_URL = (config.get_host() or "") + "/v1/deployments"
+    TINYBIRD_API_URL = f"{config.get_host()}/v1/deployments"
 
     r = requests.get(TINYBIRD_API_URL, headers=HEADERS)
     result = r.json()
@@ -226,9 +226,9 @@ def release_list(project_path: Path) -> None:
     echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
 
 
-@releases_group.command(name="promote")
+@deploy_group.command(name="promote")
 @click.argument("project_path", type=click.Path(exists=True), default=Path.cwd())
-def release_promote(project_path: Path) -> None:
+def deploy_promote(project_path: Path) -> None:
     """
     Promote last deploy to ready and remove old one.
     """
@@ -237,18 +237,18 @@ def release_promote(project_path: Path) -> None:
     TINYBIRD_API_KEY = config.get_token()
     HEADERS = {"Authorization": f"Bearer {TINYBIRD_API_KEY}"}
 
-    promote_deployment((config.get_host() or ""), HEADERS)
+    promote_deployment(config.get_host(), HEADERS)
 
 
-@releases_group.command(name="rollback")
+@deploy_group.command(name="rollback")
 @click.argument("project_path", type=click.Path(exists=True), default=Path.cwd())
-def release_rollback(project_path: Path) -> None:
+def deploy_rollback(project_path: Path) -> None:
     """
-    Rollback to the previous release.
+    Rollback to the previous deployment.
     """
     config = CLIConfig.get_project_config(str(project_path))
 
     TINYBIRD_API_KEY = config.get_token()
     HEADERS = {"Authorization": f"Bearer {TINYBIRD_API_KEY}"}
 
-    rollback_deployment((config.get_host() or ""), HEADERS)
+    rollback_deployment(config.get_host(), HEADERS)

tinybird/tb/modules/llm.py

@@ -1,8 +1,10 @@
 import json
+import os
 import urllib.parse
 from copy import deepcopy
-from typing import List
+from typing import Any, List
 
+from anthropic import AnthropicVertex
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
@@ -37,7 +39,7 @@ class LLM:
         self.user_client = deepcopy(client)
         self.user_client.token = user_token
 
-    async def ask(self, prompt: str, system_prompt: str = "", model: str = "o1-mini") -> str:
+    async def ask(self, prompt: str, system_prompt: str = "") -> str:
         """
         Calls the model with the given prompt and returns the response.
 
@@ -47,7 +49,7 @@ class LLM:
         Returns:
             str: The response from the language model.
         """
-        messages = []
+        messages: List[Any] = []
 
         if system_prompt:
             messages.append({"role": "user", "content": system_prompt})
@@ -55,8 +57,21 @@ class LLM:
         if prompt:
             messages.append({"role": "user", "content": prompt})
 
+        if gcloud_access_token := os.getenv("GCLOUD_ACCESS_TOKEN"):
+            client = AnthropicVertex(
+                region="europe-west1",
+                project_id="gen-lang-client-0705305160",
+                access_token=gcloud_access_token,
+            )
+            message = client.messages.create(
+                max_tokens=8000,
+                messages=messages,
+                model="claude-3-5-sonnet-v2@20241022",
+            )
+            return message.content[0].text or ""  # type: ignore
+
         data = {
-            "model": model,
+            "model": "o1-mini",
             "messages": messages,
         }
         response = await self.user_client._req(

tinybird/tb/modules/local.py

@@ -82,7 +82,7 @@ def get_docker_client():
         client.ping()
         return client
     except Exception:
-        raise CLIException("Docker is not running or installed. Make sure Docker is installed and running.")
+        raise CLIException("No container runtime is running. Make sure a Docker-compatible runtime is installed and running.")
 
 
 def stop_tinybird_local(docker_client):

tinybird/tb/modules/login.py

@@ -135,7 +135,12 @@ async def login(host: str, workspace: str):
     server_thread.start()
 
     # Open the browser to the auth page
-    client_id = "T6excMo8IKguvUw4vFNYfqlt9pe6msCU"
+    if "wadus" in host:
+        client_id = "Rpl7Uy9aSjqoPCSvHgGl3zNQuZcSOXBe"
+        base_auth_url = "https://auth.wadus1.tinybird.co"
+    else:
+        client_id = "T6excMo8IKguvUw4vFNYfqlt9pe6msCU"
+        base_auth_url = "https://auth.tinybird.co"
     callback_url = f"http://localhost:{AUTH_SERVER_PORT}"
     params = {
         "client_id": client_id,
@@ -143,7 +148,7 @@ async def login(host: str, workspace: str):
         "response_type": "token",
         "scope": "openid profile email",
     }
-    auth_url = f"https://auth.tinybird.co/authorize?{urlencode(params)}"
+    auth_url = f"{base_auth_url}/authorize?{urlencode(params)}"
     webbrowser.open(auth_url)
 
     # Wait for the authentication to complete or timeout

tinybird/tb/modules/shell.py

@@ -274,7 +274,7 @@ class Shell:
         elif arg.startswith("mock"):
             self.handle_mock(arg)
         else:
-            subprocess.run(f"tb --local {arg}", shell=True, text=True)
+            subprocess.run(f"tb {arg}", shell=True, text=True)
 
     def default(self, argline):
         click.echo("")
@@ -286,7 +286,7 @@ class Shell:
         elif len(arg.split()) == 1 and arg in self.endpoints + self.pipes + self.datasources + self.shared_datasources:
             self.run_sql(f"select * from {arg}")
         else:
-            subprocess.run(f"tb --local {arg}", shell=True, text=True)
+            subprocess.run(f"tb {arg}", shell=True, text=True)
 
     def run_sql(self, query, rows_limit=20):
         try:

tinybird/tb/modules/test.py

@@ -6,11 +6,13 @@
 import difflib
 import glob
 import os
+import urllib.parse
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import click
 import yaml
+from requests import Response
 
 from tinybird.prompts import test_create_prompt
 from tinybird.tb.modules.cli import cli
@@ -66,9 +68,9 @@ def test(ctx: click.Context) -> None:
     type=click.Path(exists=True, file_okay=False),
     help="Folder where datafiles will be placed",
 )
-@click.option("--prompt", type=str, default=None, help="Prompt to be used to create the test")
+@click.option("--prompt", type=str, default="", help="Prompt to be used to create the test")
 @coro
-async def test_create(name_or_filename: str, prompt: Optional[str], folder: str) -> None:
+async def test_create(name_or_filename: str, prompt: str, folder: str) -> None:
     """
     Create a test for an existing pipe
     """
@@ -80,15 +82,13 @@ async def test_create(name_or_filename: str, prompt: Optional[str], folder: str)
                 raise CLIException(FeedbackManager.error(message=f"Pipe {name_or_filename} not found"))
         else:
             pipe_folders = ("endpoints", "copies", "materializations", "sinks", "pipes")
-            pipe_path = next(
-                (
+            try:
+                pipe_path = next(
                     root_path / folder / f"{name_or_filename}.pipe"
                     for folder in pipe_folders
                     if (root_path / folder / f"{name_or_filename}.pipe").exists()
-                ),
-                None,
-            )
-            if not pipe_path:
+                )
+            except Exception:
                 raise CLIException(FeedbackManager.error(message=f"Pipe {name_or_filename} not found"))
 
         pipe_name = pipe_path.stem
@@ -97,8 +97,8 @@ async def test_create(name_or_filename: str, prompt: Optional[str], folder: str)
         pipe_content = pipe_path.read_text()
 
         client = await get_tinybird_local_client(os.path.abspath(folder))
-        pipe_nodes = await client._req(f"/v0/pipes/{pipe_name}")
-        parameters = set([param["name"] for node in pipe_nodes["nodes"] for param in node["params"]])
+        pipe = await client._req(f"/v0/pipes/{pipe_name}")
+        parameters = set([param["name"] for node in pipe["nodes"] for param in node["params"]])
 
         system_prompt = test_create_prompt.format(
             name=pipe_name,
@@ -107,24 +107,26 @@ async def test_create(name_or_filename: str, prompt: Optional[str], folder: str)
         )
         config = CLIConfig.get_project_config(folder)
         user_token = config.get_user_token()
+        if not user_token:
+            raise CLIException(FeedbackManager.error(message="No user token found"))
         llm = LLM(user_token=user_token, client=config.get_client())
 
-        response = await llm.ask(prompt, system_prompt=system_prompt)
-        response = extract_xml(response, "response")
-        tests_content = parse_xml(response, "test")
+        response_llm = await llm.ask(prompt, system_prompt=system_prompt)
+        response_xml = extract_xml(response_llm, "response")
+        tests_content = parse_xml(response_xml, "test")
 
         tests: List[Dict[str, Any]] = []
         for test_content in tests_content:
-            test = {}
+            test: Dict[str, Any] = {}
             test["name"] = extract_xml(test_content, "name")
             test["description"] = extract_xml(test_content, "description")
-            parameters = extract_xml(test_content, "parameters")
-            test["parameters"] = parameters.split("?")[1] if "?" in parameters else parameters
+            parameters_api = extract_xml(test_content, "parameters")
+            test["parameters"] = parameters_api.split("?")[1] if "?" in parameters_api else parameters_api
             test["expected_result"] = ""
 
             response = None
             try:
-                response = await client._req_raw(f"/v0/pipes/{pipe_name}.ndjson?{parameters}")
+                response = await get_pipe_data(client, pipe_name=pipe_name, test_params=test["parameters"])
             except Exception:
                 pass
 
@@ -181,7 +183,7 @@ async def test_update(pipe: str, folder: str) -> None:
             test_params = test["parameters"].split("?")[1] if "?" in test["parameters"] else test["parameters"]
             response = None
             try:
-                response = await client._req_raw(f"/v0/pipes/{pipe_name}.ndjson?{test_params}")
+                response = await get_pipe_data(client, pipe_name=pipe_name, test_params=test_params)
             except Exception:
                 continue
 
@@ -221,20 +223,21 @@ async def run_tests(name: Tuple[str, ...], folder: str) -> None:
     client = await get_tinybird_local_client(os.path.abspath(folder))
     paths = [Path(n) for n in name]
     endpoints = [f"./tests/{p.stem}.yaml" for p in paths]
-    test_files: Iterable[str] = endpoints if len(endpoints) > 0 else glob.glob("./tests/**/*.y*ml", recursive=True)
+    test_files: List[str] = endpoints if len(endpoints) > 0 else glob.glob("./tests/**/*.y*ml", recursive=True)
 
     async def run_test(test_file):
         test_file_path = Path(test_file)
         click.echo(FeedbackManager.info(message=f"\n* {test_file_path.stem}{test_file_path.suffix}"))
         test_file_content = yaml.safe_load(test_file_path.read_text())
+
         for test in test_file_content:
             try:
                 test_params = test["parameters"].split("?")[1] if "?" in test["parameters"] else test["parameters"]
                 response = None
                 try:
-                    response = await client._req_raw(f"/v0/pipes/{test_file_path.stem}.ndjson?{test_params}")
+                    response = await get_pipe_data(client, pipe_name=test_file_path.stem, test_params=test_params)
                 except Exception:
-                    raise Exception("Expected to not fail but got an error")
+                    continue
 
                 expected_result = response.text
                 if response.status_code >= 400:
@@ -271,3 +274,19 @@ async def run_tests(name: Tuple[str, ...], folder: str) -> None:
         exit(1)
     else:
         click.echo(FeedbackManager.success(message=f"\n✓ {test_count}/{test_count} passed"))
+
+
+async def get_pipe_data(client, pipe_name: str, test_params: str) -> Response:
+    pipe = await client._req(f"/v0/pipes/{pipe_name}")
+    output_node = next(
+        (node for node in pipe["nodes"] if node["node_type"] != "default" and node["node_type"] != "standard"),
+        {"name": "not_found"},
+    )
+    if output_node["node_type"] == "endpoint":
+        return await client._req_raw(f"/v0/pipes/{pipe_name}.ndjson?{test_params}")
+
+    params = {
+        "q": output_node["sql"],
+        "pipeline": pipe_name,
+    }
+    return await client._req_raw(f"""/v0/sql?{urllib.parse.urlencode(params)}&{test_params}""")