tinybird 0.0.1.dev42__py3-none-any.whl → 0.0.1.dev44__py3-none-any.whl

tinybird/tb/modules/endpoint.py

@@ -5,7 +5,7 @@
 
 import json
 import re
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 from urllib.parse import urlencode
 
 import click
@@ -21,7 +21,7 @@ from tinybird.tb.modules.exceptions import CLIPipeException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
-@cli.group(hidden=True)
+@cli.group()
 @click.pass_context
 def endpoint(ctx):
     """Endpoint commands"""
@@ -161,7 +161,7 @@ async def endpoint_url(ctx: Context, pipe: str):
     click.echo(build_endpoint_url(client, pipe, token))
 
 
-def build_endpoint_url(tb_client: TinyB, pipe_name: str, token: Optional[str]) -> str:
+def build_endpoint_url(tb_client: TinyB, pipe_name: str, token: Optional[str]) -> Optional[str]:
     try:
         token = token or tb_client.token
         example_params = {
@@ -185,3 +185,90 @@ def get_endpoint_token(tokens: List[Dict[str, Any]], pipe_name: str) -> Optional
                 break
 
     return token
+
+
+@endpoint.command(name="stats")
+@click.argument("pipes", nargs=-1)
+@click.option(
+    "--format",
+    "format_",
+    type=click.Choice(["json"], case_sensitive=False),
+    default=None,
+    help="Force a type of the output. To parse the output, keep in mind to use `tb --no-version-warning endpoint stats` option.",
+)
+@click.pass_context
+@coro
+async def endpoint_stats(ctx: click.Context, pipes: Tuple[str, ...], format_: str):
+    """
+    Print endpoint stats for the last 7 days
+    """
+    client: TinyB = ctx.ensure_object(dict)["client"]
+    all_pipes = await client.pipes()
+    pipes_to_get_stats = []
+    pipes_ids: Dict = {}
+
+    if pipes:
+        # We filter by the pipes we want to look for
+        all_pipes = [pipe for pipe in all_pipes if pipe["name"] in pipes]
+
+    for pipe in all_pipes:
+        name_version = get_name_version(pipe["name"])
+        if name_version["name"] in pipe["name"]:
+            pipes_to_get_stats.append(f"'{pipe['id']}'")
+            pipes_ids[pipe["id"]] = name_version
+
+    if not pipes_to_get_stats:
+        if format_ == "json":
+            click.echo(json.dumps({"pipes": []}, indent=2))
+        else:
+            click.echo(FeedbackManager.info_no_pipes_stats())
+        return
+
+    sql = f"""
+        SELECT
+            pipe_id id,
+            sumIf(view_count, date > now() - interval 7 day) requests,
+            sumIf(error_count, date > now() - interval 7 day) errors,
+            avgMergeIf(avg_duration_state, date > now() - interval 7 day) latency
+        FROM tinybird.pipe_stats
+        WHERE pipe_id in ({",".join(pipes_to_get_stats)})
+        GROUP BY pipe_id
+        ORDER BY requests DESC
+        FORMAT JSON
+    """
+
+    res = await client.query(sql)
+
+    if res and "error" in res:
+        raise CLIPipeException(FeedbackManager.error_exception(error=str(res["error"])))
+
+    columns = ["name", "request count", "error count", "avg latency"]
+    table_human_readable: List[Tuple] = []
+    table_machine_readable: List[Dict] = []
+    if res and "data" in res:
+        for x in res["data"]:
+            tk = pipes_ids[x["id"]]
+            table_human_readable.append(
+                (
+                    tk["name"],
+                    x["requests"],
+                    x["errors"],
+                    x["latency"],
+                )
+            )
+            table_machine_readable.append(
+                {
+                    "name": tk["name"],
+                    "requests": x["requests"],
+                    "errors": x["errors"],
+                    "latency": x["latency"],
+                }
+            )
+
+        table_human_readable.sort(key=lambda x: (x[1], x[0]))
+        table_machine_readable.sort(key=lambda x: x["name"])
+
+        if format_ == "json":
+            click.echo(json.dumps({"pipes": table_machine_readable}, indent=2))
+        else:
+            echo_safe_humanfriendly_tables_format_smart_table(table_human_readable, column_names=columns)

tinybird/tb/modules/llm_utils.py

@@ -36,7 +36,7 @@ def generate(llm_call, task: str, feedback: str = "") -> tuple[str, str]:
     thoughts = extract_xml(response, "thoughts")
     result = extract_xml(response, "response")
 
-    return thoughts, result, response
+    return thoughts, result
 
 
 def evaluate(llm_call, content: str, task: str) -> tuple[str, str]:
@@ -96,7 +96,7 @@ Output your evaluation concisely in the following format:
 """
 
 generator_prompt = """
-Your goal is to complete the task based on <task> tag. If there are feedback 
+Your goal is to complete the task based on <task> tag. If there are feedback
 from your previous generations, you should reflect on them to solve the task.
 All xml tags MUST be closed.
 

tinybird/tb/modules/materialization.py

@@ -0,0 +1,146 @@
+import json
+import re
+
+import click
+
+from tinybird.client import TinyB
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import (
+    coro,
+    create_tb_client,
+    echo_safe_humanfriendly_tables_format_smart_table,
+    wait_job,
+)
+from tinybird.tb.modules.datafile.common import PipeTypes, get_name_version
+from tinybird.tb.modules.exceptions import CLIPipeException
+from tinybird.tb.modules.feedback_manager import FeedbackManager
+
+
+@cli.group()
+@click.pass_context
+def materialization(ctx):
+    """Materialization commands"""
+
+
+@materialization.command(name="ls")
+@click.option("--match", default=None, help="Retrieve any resourcing matching the pattern. eg --match _test")
+@click.option(
+    "--format",
+    "format_",
+    type=click.Choice(["json"], case_sensitive=False),
+    default=None,
+    help="Force a type of the output",
+)
+@click.pass_context
+@coro
+async def materialization_ls(ctx: click.Context, match: str, format_: str):
+    """List materializations"""
+
+    client: TinyB = ctx.ensure_object(dict)["client"]
+    pipes = await client.pipes(dependencies=True, node_attrs="name,materialized", attrs="name,updated_at,endpoint,type")
+    materializations = [p for p in pipes if p.get("type") == PipeTypes.MATERIALIZED]
+    materializations = sorted(materializations, key=lambda p: p["updated_at"])
+    datasources = await client.datasources()
+    columns = ["name", "updated at", "nodes", "target datasource"]
+    table_human_readable = []
+    table_machine_readable = []
+    pattern = re.compile(match) if match else None
+    for t in materializations:
+        tk = get_name_version(t["name"])
+        if pattern and not pattern.search(tk["name"]):
+            continue
+        target_datasource_id = next((n["materialized"] for n in t["nodes"] if n.get("materialized")), None)
+        target_datasource = next((d for d in datasources if d["id"] == target_datasource_id), None)
+        target_datasource_name = target_datasource.get("name", "") if target_datasource else ""
+        table_human_readable.append((tk["name"], t["updated_at"][:-7], len(t["nodes"]), target_datasource_name))
+        table_machine_readable.append(
+            {
+                "name": tk["name"],
+                "updated at": t["updated_at"][:-7],
+                "nodes": len(t["nodes"]),
+                "target datasource": target_datasource_name,
+            }
+        )
+
+    if not format_:
+        click.echo(FeedbackManager.info_pipes())
+        echo_safe_humanfriendly_tables_format_smart_table(table_human_readable, column_names=columns)
+        click.echo("\n")
+    elif format_ == "json":
+        click.echo(json.dumps({"pipes": table_machine_readable}, indent=2))
+    else:
+        raise CLIPipeException(FeedbackManager.error_pipe_ls_type())
+
+
+@materialization.command(name="populate")
+@click.argument("pipe_name")
+@click.option("--node", type=str, help="Name of the materialized node.", default=None, required=False)
+@click.option(
+    "--sql-condition",
+    type=str,
+    default=None,
+    help="Populate with a SQL condition to be applied to the trigger Data Source of the Materialized View. For instance, `--sql-condition='date == toYYYYMM(now())'` it'll populate taking all the rows from the trigger Data Source which `date` is the current month. Use it together with --populate. --sql-condition is not taken into account if the --subset param is present. Including in the ``sql_condition`` any column present in the Data Source ``engine_sorting_key`` will make the populate job process less data.",
+)
+@click.option(
+    "--truncate", is_flag=True, default=False, help="Truncates the materialized Data Source before populating it."
+)
+@click.option(
+    "--unlink-on-populate-error",
+    is_flag=True,
+    default=False,
+    help="If the populate job fails the Materialized View is unlinked and new data won't be ingested in the Materialized View. First time a populate job fails, the Materialized View is always unlinked.",
+)
+@click.option(
+    "--wait",
+    is_flag=True,
+    default=False,
+    help="Waits for populate jobs to finish, showing a progress bar. Disabled by default.",
+)
+@click.pass_context
+@coro
+async def pipe_populate(
+    ctx: click.Context,
+    pipe_name: str,
+    node: str,
+    sql_condition: str,
+    truncate: bool,
+    unlink_on_populate_error: bool,
+    wait: bool,
+):
+    """Populate the result of a Materialized Node into the target Materialized View"""
+    cl = create_tb_client(ctx)
+
+    pipe = await cl.pipe(pipe_name)
+
+    if pipe["type"] != PipeTypes.MATERIALIZED:
+        raise CLIPipeException(FeedbackManager.error_pipe_not_materialized(pipe=pipe_name))
+
+    if not node:
+        materialized_ids = [pipe_node["id"] for pipe_node in pipe["nodes"] if pipe_node.get("materialized") is not None]
+
+        if not materialized_ids:
+            raise CLIPipeException(FeedbackManager.error_populate_no_materialized_in_pipe(pipe=pipe_name))
+
+        elif len(materialized_ids) > 1:
+            raise CLIPipeException(FeedbackManager.error_populate_several_materialized_in_pipe(pipe=pipe_name))
+
+        node = materialized_ids[0]
+
+    response = await cl.populate_node(
+        pipe_name,
+        node,
+        populate_condition=sql_condition,
+        truncate=truncate,
+        unlink_on_populate_error=unlink_on_populate_error,
+    )
+    if "job" not in response:
+        raise CLIPipeException(response)
+
+    job_id = response["job"]["id"]
+    job_url = response["job"]["job_url"]
+    if sql_condition:
+        click.echo(FeedbackManager.info_populate_condition_job_url(url=job_url, populate_condition=sql_condition))
+    else:
+        click.echo(FeedbackManager.info_populate_job_url(url=job_url))
+    if wait:
+        await wait_job(cl, job_id, job_url, "Populating")

tinybird/tb/modules/mock.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from pathlib import Path
+from typing import Optional
 
 import click
 
@@ -13,6 +14,7 @@ from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.llm import LLM
 from tinybird.tb.modules.llm_utils import extract_xml
 from tinybird.tb.modules.local_common import get_tinybird_local_client
+from tinybird.tb.modules.project import Project
 
 
 @cli.command()
@@ -21,24 +23,27 @@ from tinybird.tb.modules.local_common import get_tinybird_local_client
 @click.option(
     "--prompt",
     type=str,
-    default="Use the datasource schema to generate sample data",
+    default="",
     help="Extra context to use for data generation",
 )
-@click.option("--folder", type=str, default=os.getcwd(), help="Folder where datafiles will be placed")
+@click.option("--skip", is_flag=True, default=False, help="Skip following up on the generated data")
+@click.pass_context
 @coro
-async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
-    """Load sample data into a Data Source.
+async def mock(ctx: click.Context, datasource: str, rows: int, prompt: str, skip: bool) -> None:
+    """Generate sample data for a datasource.
 
     Args:
         datasource: Path to the datasource file to load sample data into
         rows: Number of events to send
         prompt: Extra context to use for data generation
-        folder: Folder where datafiles will be placed
+        skip: Skip following up on the generated data
     """
 
     try:
+        project: Project = ctx.ensure_object(dict)["project"]
         datasource_path = Path(datasource)
         datasource_name = datasource
+        folder = project.folder
         click.echo(FeedbackManager.highlight(message=f"\n» Creating fixture for {datasource_name}..."))
         if datasource_path.suffix == ".datasource":
             datasource_name = datasource_path.stem
@@ -46,6 +51,9 @@ async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
             datasource_path = Path("datasources", f"{datasource}.datasource")
         datasource_path = Path(folder) / datasource_path
 
+        if not datasource_path.exists():
+            raise CLIException(f"Datasource '{datasource_path.stem}' not found")
+
         prompt_path = Path(folder) / "fixtures" / f"{datasource_name}.prompt"
         if not prompt or prompt == "Use the datasource schema to generate sample data":
             # load the prompt from the fixture.prompt file if it exists
@@ -68,17 +76,49 @@ async def mock(datasource: str, rows: int, prompt: str, folder: str) -> None:
             click.echo(FeedbackManager.error(message="This action requires authentication. Run 'tb login' first."))
             return
         llm = LLM(user_token=user_token, host=user_client.host)
-        tb_client = await get_tinybird_local_client(os.path.abspath(folder))
+        tb_client = await get_tinybird_local_client(folder)
         prompt = f"<datasource_schema>{datasource_content}</datasource_schema>\n<user_input>{prompt}</user_input>"
-        response = llm.ask(system_prompt=mock_prompt(rows), prompt=prompt)
-        sql = extract_xml(response, "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().as_posix(), datasource_name, datasource_content)
-        persist_fixture(fixture_name, data, folder)
-        click.echo(FeedbackManager.success(message=f"✓ /fixtures/{fixture_name}.ndjson created with {rows} rows"))
+        iterations = 0
+        history = ""
+        fixture_path: Optional[Path] = None
+        sql = ""
+        while iterations < 10:
+            feedback = ""
+            if iterations > 0:
+                feedback = click.prompt("\nFollow-up instructions or continue", default="continue")
+            if iterations > 0 and (not feedback or feedback in ("continue", "ok", "exit", "quit", "q")):
+                break
+            else:
+                if iterations > 0:
+                    if fixture_path:
+                        fixture_path.unlink()
+                    fixture_path = None
+                    click.echo(FeedbackManager.highlight(message=f"\n» Creating fixture for {datasource_name}..."))
+
+                response = llm.ask(system_prompt=mock_prompt(rows, feedback, history), prompt=prompt)
+                sql = extract_xml(response, "sql")
+                result = await tb_client.query(f"{sql} FORMAT JSON")
+                data = result.get("data", [])[:rows]
+                fixture_name = build_fixture_name(str(datasource_path), datasource_name, datasource_content)
+                fixture_path = persist_fixture(fixture_name, data, folder)
+                click.echo(FeedbackManager.info(message=f"✓ /fixtures/{fixture_name}.ndjson created"))
+
+                if os.environ.get("TB_DEBUG", "") != "":
+                    logging.debug(sql)
+
+                history = (
+                    history
+                    + f"""
+                <result_iteration_{iterations}>
+                    {response}
+                </result_iteration_{iterations}>
+                """
+                )
+                if skip:
+                    break
+                iterations += 1
+
+        click.echo(FeedbackManager.success(message=f"✓ Sample data for {datasource_name} created with {rows} rows"))
 
     except Exception as e:
-        raise CLIException(FeedbackManager.error_exception(error=e))
+        click.echo(FeedbackManager.error_exception(error=f"Error: {e}"))