tinybird 0.0.1.dev282__tar.gz → 0.0.1.dev284__tar.gz

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: tinybird
-Version: 0.0.1.dev282
+Version: 0.0.1.dev284
 Summary: Tinybird Command Line Tool
 Home-page: https://www.tinybird.co/docs/forward/commands
 Author: Tinybird

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/ch_utils/constants.py

@@ -257,4 +257,5 @@ VALID_QUERY_FORMATS = (
     "Native",
     "RowBinaryWithNamesAndTypes",
     "TabSeparated",
+    "JSONCompactEachRowWithNamesAndTypes",
 )

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/datafile/common.py

@@ -184,7 +184,7 @@ class CopyParameters(Parameters):
 
 class MaterializedParameters(Parameters):
     MANDATORY_ATTRIBUTES = PipeParameters.MANDATORY_ATTRIBUTES.union({"datasource"})
-    ACCEPTED_ATTRIBUTES = PipeParameters.ACCEPTED_ATTRIBUTES.union(MANDATORY_ATTRIBUTES)
+    ACCEPTED_ATTRIBUTES = PipeParameters.ACCEPTED_ATTRIBUTES.union(MANDATORY_ATTRIBUTES).union({"deployment_method"})
 
 
 class SinkParameters(Parameters):
@@ -294,6 +294,13 @@ class Datafile:
     def set_kind(self, kind: DatafileKind):
         self.kind = kind
 
+    def validate_standard_node(self, node: Dict[str, Any]):
+        for key in node.keys():
+            if key not in PipeParameters.valid_params():
+                raise DatafileValidationError(
+                    f"Standard node {repr(node['name'])} has an invalid attribute ({PipeParameters.canonical_name(key)})"
+                )
+
     def validate_copy_node(self, node: Dict[str, Any]):
         if missing := [param for param in CopyParameters.required_params() if param not in node]:
             raise DatafileValidationError(
@@ -412,13 +419,6 @@ class Datafile:
 
     def validate(self):
         if self.kind == DatafileKind.pipe:
-            # TODO(eclbg):
-            #  [x] node names are unique
-            #  [x] SQL in all nodes
-            #  [x] Materialized nodes have target datasource
-            #  [x] Only one materialized node
-            #  [x] Only one node of any specific type
-            #  (rbarbadillo): there's a HUGE amount of validations in api_pipes.py, we should somehow merge them
             non_standard_nodes_count = 0
             for node in self.nodes:
                 node_type = node.get("type", "").lower()
@@ -432,10 +432,13 @@ class Datafile:
                     self.validate_copy_node(node)
                 if node_type == PipeNodeTypes.DATA_SINK:
                     self.validate_sink_node(node)
+                if node_type in {PipeNodeTypes.STANDARD, ""}:
+                    self.validate_standard_node(node)
                 if node_type not in VALID_PIPE_NODE_TYPES:
                     raise DatafileValidationError(
                         f"Invalid node '{repr(node['name'])}' of type ({node_type}). Allowed node types: {VISIBLE_PIPE_NODE_TYPES}"
                     )
+
             for token in self.tokens:
                 if token["permission"].upper() != "READ":
                     raise DatafileValidationError(
@@ -1983,6 +1986,7 @@ def parse(
             "include": include,
             "sql": sql("sql"),
             "version": version,
+            "deployment_method": assign_var("deployment_method", allowed_values={"alter"}),
             "export_connection_name": assign_var("export_connection_name"),
             "export_schedule": assign_var("export_schedule"),
             "export_bucket_uri": assign_var("export_bucket_uri"),

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/__cli__.py

@@ -4,5 +4,5 @@ __description__ = 'Tinybird Command Line Tool'
 __url__ = 'https://www.tinybird.co/docs/forward/commands'
 __author__ = 'Tinybird'
 __author_email__ = 'support@tinybird.co'
-__version__ = '0.0.1.dev282'
-__revision__ = '77eb85f'
+__version__ = '0.0.1.dev284'
+__revision__ = '8beb9ee'

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/modules/agent/agent.py

@@ -28,6 +28,7 @@ from tinybird.tb.modules.agent.memory import (
     get_last_messages_from_last_user_prompt,
     save_messages,
 )
+from tinybird.tb.modules.agent.mock_agent import MockAgent
 from tinybird.tb.modules.agent.models import create_model
 from tinybird.tb.modules.agent.prompts import (
     agent_system_prompt,
@@ -45,7 +46,6 @@ from tinybird.tb.modules.agent.tools.deploy_check import deploy_check
 from tinybird.tb.modules.agent.tools.diff_resource import diff_resource
 from tinybird.tb.modules.agent.tools.get_endpoint_stats import get_endpoint_stats
 from tinybird.tb.modules.agent.tools.get_openapi_definition import get_openapi_definition
-from tinybird.tb.modules.agent.tools.mock import mock
 from tinybird.tb.modules.agent.tools.plan import plan
 from tinybird.tb.modules.agent.tools.secret import create_or_update_secrets
 from tinybird.tb.modules.agent.utils import AgentRunCancelled, TinybirdAgentContext, show_confirmation, show_input
@@ -111,7 +111,6 @@ class TinybirdAgent:
                 Tool(build, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(deploy, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(deploy_check, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
-                Tool(mock, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(analyze_file, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(analyze_url, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
                 Tool(append_file, docstring_format="google", require_parameter_descriptions=True, takes_ctx=True),
@@ -166,6 +165,16 @@ class TinybirdAgent:
             workspace_id=workspace_id,
             project=self.project,
         )
+        self.mock_agent = MockAgent(
+            dangerously_skip_permissions=self.dangerously_skip_permissions,
+            prompt_mode=prompt_mode,
+            thinking_animation=self.thinking_animation,
+            token=self.token,
+            user_token=self.user_token,
+            host=self.host,
+            workspace_id=workspace_id,
+            project=self.project,
+        )
 
         @self.agent.tool
         def manage_tests(ctx: RunContext[TinybirdAgentContext], task: str) -> str:
@@ -210,6 +219,25 @@ class TinybirdAgent:
             result = self.explore_agent.run(task, deps=ctx.deps, usage=ctx.usage)
             return result.output or "No result returned"
 
+        @self.agent.tool
+        def mock(
+            ctx: RunContext[TinybirdAgentContext], datasource_name: str, rows: int, data_format: str, task: str
+        ) -> str:
+            """Generate mock data for a datasource.
+
+            Args:
+                datasource_name (str): The datasource name to generate mock data for. Required.
+                rows (int): Number of rows to create. If not provided, the default is 10. Required.
+                data_format (str): Format of the mock data to create. Options: ndjson, csv. Required.
+                task (str): Extra details about how to generate the mock data (nested json if any, sample row to help with the generation, etc). Required.
+
+            Returns:
+                str: The result of the mock data generation.
+            """
+            user_input = f"Datasource name: {datasource_name}\nRows: {rows}\nData format: {data_format}\nTask: {task}"
+            result = self.mock_agent.run(user_input, deps=ctx.deps, usage=ctx.usage)
+            return result.output or "No result returned"
+
         @self.agent.instructions
         def get_local_host(ctx: RunContext[TinybirdAgentContext]) -> str:
             return f"""
@@ -355,20 +383,18 @@ class TinybirdAgent:
             workspace_id = config.get("id", "")
             workspace = client.workspace(workspace_id, with_organization=True, version="v1")
             limits_data = client.organization_limits(workspace["organization"]["id"])
-            ai_requests_limits = limits_data.get("limits", {}).get("ai_requests", {})
-            current_ai_requests = ai_requests_limits.get("quantity") or 0
-            max_ai_requests = ai_requests_limits.get("max") or 0
-            remaining_requests = max(max_ai_requests - current_ai_requests, 0)
-            current_ai_requests = min(max_ai_requests, current_ai_requests)
-            if not max_ai_requests:
+            ai_credits_limits = limits_data.get("limits", {}).get("ai_credits", {})
+            current_ai_credits = ai_credits_limits.get("quantity") or 0
+            ai_credits = ai_credits_limits.get("max") or 0
+            remaining_credits = max(ai_credits - current_ai_credits, 0)
+            current_ai_credits = min(ai_credits, current_ai_credits)
+            if not ai_credits:
                 return
-            warning_threshold = max_ai_requests * 0.8
-            message_color = (
-                FeedbackManager.warning if current_ai_requests >= warning_threshold else FeedbackManager.gray
-            )
+            warning_threshold = ai_credits * 0.8
+            message_color = FeedbackManager.warning if current_ai_credits >= warning_threshold else FeedbackManager.gray
             click.echo(
                 message_color(
-                    message=f"{remaining_requests} requests left ({current_ai_requests}/{max_ai_requests}). You can continue using Tinybird Code. Limits will be enforced soon."
+                    message=f"{remaining_credits} credits left ({current_ai_credits}/{ai_credits}). You can continue using Tinybird Code. Limits will be enforced soon."
                 )
             )
         except Exception:

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/modules/agent/explore_agent.py

@@ -50,6 +50,11 @@ You can do the following:
 - Executing SQL queries against Tinybird Cloud or Tinybird Local.
 - Requesting endpoints in Tinybird Cloud or Tinybird Local.
 - Visualizing data as a chart using execute_query tool with the `script` parameter.
+
+IMPORTANT: Use always the last environment used in previous queries or endpoint requests (cloud_or_local: str). If you don't have any information about the last environment, use None.
+IMPORTANT: If some resource is not found in a environment, you can use the `diff_resource` tool to check the status across environments.
+
+Once you finish the task, return a valid response for the task to complete.
 """,
                 tone_and_style_instructions,
                 explore_data_instructions,

tinybird-0.0.1.dev284/tinybird/tb/modules/agent/mock_agent.py

@@ -0,0 +1,206 @@
+from datetime import datetime
+
+from pydantic_ai import Agent, RunContext, Tool
+from pydantic_ai.messages import ModelMessage
+from pydantic_ai.usage import Usage
+
+from tinybird.tb.modules.agent.animations import ThinkingAnimation
+from tinybird.tb.modules.agent.models import create_model
+from tinybird.tb.modules.agent.prompts import resources_prompt
+from tinybird.tb.modules.agent.tools.mock import generate_mock_fixture
+from tinybird.tb.modules.agent.utils import TinybirdAgentContext
+from tinybird.tb.modules.project import Project
+
+mock_sql_instructions = """
+## When generating the SQL query to generate mock data
+
+Given the schema for a Tinybird datasource, return a valid clickhouse sql query to generate some random data that matches that schema.
+Response format MUST be just a valid clickhouse sql query.
+
+<example>
+    <example_datasource_schema>
+SCHEMA >
+    experience_gained Int16 `json:$.experience_gained`,
+    level Int16 `json:$.level`,
+    monster_kills Int16 `json:$.monster_kills`,
+    player_id String `json:$.player_id`,
+    pvp_kills Int16 `json:$.pvp_kills`,
+    quest_completions Int16 `json:$.quest_completions`,
+    timestamp DateTime `json:$.timestamp`
+    </example_datasource_schema>
+    <example_output>
+
+SELECT
+    rand() % 1000 AS experience_gained, -- Random experience gained between 0 and 999
+    1 + rand() % 100 AS level,          -- Random level between 1 and 100
+    rand() % 500 AS monster_kills,      -- Random monster kills between 0 and 499
+    concat('player_', toString(rand() % 10000)) AS player_id, -- Random player IDs like "player_1234"
+    rand() % 50 AS pvp_kills,           -- Random PvP kills between 0 and 49
+    rand() % 200 AS quest_completions,  -- Random quest completions between 0 and 199
+    now() - rand() % 86400 AS timestamp -- Random timestamp within the last day
+FROM numbers(10)
+    </example_output>
+</example>
+
+<instructions>
+- The query MUST return a random sample of data that matches the schema.
+- The query MUST return a valid clickhouse sql query.
+- The query MUST be valid for clickhouse and Tinybird.
+- FROM numbers([number_of_rows]) part is mandatory.
+- If json paths are present (e.g. `userAgent` String `json:$.userAgent`), rely on the json paths to generate the sample record.
+- If the schema has nested json paths (e.g. `json:$.location.country`), generate nested JSON objects accordingly.
+- Use recent dates to avoid generating dates that are too far in the past.
+- Do NOT include ```clickhouse or ```sql or any other wrapping text to the sql query.
+- Do NOT use any of these functions: elementAt
+- Do NOT add a semicolon at the end of the query
+- Do NOT add any FORMAT at the end of the query, because it will be added later by Tinybird.
+- Do not use any function that is not present in the list of general functions, character insensitive functions and aggregate functions.
+- If the function is not present in the list, the sql query will fail, so avoid at all costs to use any function that is not present in the list.
+</instructions>
+
+<more_examples>
+# Examples with different schemas, like an array field or a nested JSON field:
+
+## Example schema with an array field:
+
+### Schema:
+
+SCHEMA >
+    `order_id` UInt64 `json:$.order_id`,
+    `customer_id` UInt64 `json:$.customer_id`,
+    `order_date` DateTime `json:$.order_date`,
+    `total_amount` Float64 `json:$.total_amount`,
+    `items` Array(String) `json:$.items[:]` // This is an array field
+
+### Desired final output of the query:
+{
+  "order_id": 123456,
+  "customer_id": 7890,
+  "order_date": "2024-11-30T10:30:00.000Z",
+  "total_amount": 150.0,
+  "items": ["item1", "item2", "item3"]
+}
+
+### Example SQL output with an array field:
+
+SELECT
+  concat('ord_', toString(rand() % 10000)) AS order_id,
+  concat('cust_', toString(rand() % 10000)) AS customer_id,
+  now() - rand() % 86400 AS order_date,
+  rand() % 1000 AS total_amount,
+  arrayMap(x -> concat('item_', toString(x)), range(1, rand() % 5 + 1)) AS items
+FROM numbers(ROWS)
+
+## Example schema with nested JSON paths:
+
+### Schema:
+
+SCHEMA >
+    `timestamp` DateTime `json:$.timestamp`,
+    `location_country` String `json:$.location.country`,
+    `location_region` String `json:$.location.region`,
+    `location_city` String `json:$.location.city`,
+    `location_latitude` String `json:$.location.latitude`,
+    `location_longitude` String `json:$.location.longitude`
+
+### Important: Understanding JSON paths
+When you see json paths like `json:$.location.country`, it means the data should be structured as nested JSON:
+- `json:$.location.country` → the country field is nested inside the location object
+- `json:$.location.region` → the region field is nested inside the location object
+
+### Desired final output structure:
+{
+  "timestamp": "2024-11-30T10:30:00.000Z",
+  "location": {
+    "country": "United States",
+    "region": "California", 
+    "city": "San Francisco",
+    "latitude": "37.7749",
+    "longitude": "-122.4194"
+  }
+}
+
+### Example SQL output for nested JSON paths:
+
+SELECT
+    timestamp,
+    CAST(concat('{
+        "country": "', country, '",
+        "region": "', region, '",
+        "city": "', city, '",
+        "latitude": "', latitude, '",
+        "longitude": "', longitude, '"
+    }'), 'JSON') AS location
+FROM
+(
+    SELECT
+        now() - rand() % 86400 AS timestamp,
+        ['United States', 'Canada', 'United Kingdom', 'Germany', 'France'][(rand() % 5) + 1] AS country,
+        ['California', 'Texas', 'New York', 'Ontario', 'London'][(rand() % 5) + 1] AS region,
+        ['San Francisco', 'Los Angeles', 'New York', 'Toronto', 'London'][(rand() % 5) + 1] AS city,
+        toString(round(rand() * 180 - 90, 4)) AS latitude,
+        toString(round(rand() * 360 - 180, 4)) AS longitude
+    FROM numbers(10)
+)
+</more_examples>
+"""
+
+
+class MockAgent:
+    def __init__(
+        self,
+        dangerously_skip_permissions: bool,
+        prompt_mode: bool,
+        token: str,
+        user_token: str,
+        host: str,
+        workspace_id: str,
+        project: Project,
+        thinking_animation: ThinkingAnimation,
+    ):
+        self.dangerously_skip_permissions = dangerously_skip_permissions or prompt_mode
+        self.token = token
+        self.user_token = user_token
+        self.host = host
+        self.workspace_id = workspace_id
+        self.project = project
+        self.thinking_animation = thinking_animation
+        self.messages: list[ModelMessage] = []
+        self.agent = Agent(
+            deps_type=TinybirdAgentContext,
+            instructions=[
+                f"""
+You are part of Tinybird Code, an agentic CLI that can help users to work with Tinybird.                 
+You are a sub-agent of the main Tinybird Code agent. You are responsible for generating mock data for a datasource.
+You will be given a datasource name and you will use `generate_mock_fixture` tool to generate a sql query to execute to generate the mock data.
+When finish return the result of the mock data generation: the path of the fixture file, the number of rows generated, the datasource name and if the data was appended to the datasource.
+
+# Info
+Today is {datetime.now().strftime("%Y-%m-%d")}
+""",
+                mock_sql_instructions,
+            ],
+            tools=[
+                Tool(
+                    generate_mock_fixture,
+                    docstring_format="google",
+                    require_parameter_descriptions=True,
+                    takes_ctx=True,
+                ),
+            ],
+        )
+
+        @self.agent.instructions
+        def get_project_files(ctx: RunContext[TinybirdAgentContext]) -> str:
+            return resources_prompt(self.project)
+
+    def run(self, task: str, deps: TinybirdAgentContext, usage: Usage):
+        result = self.agent.run_sync(
+            task,
+            deps=deps,
+            usage=usage,
+            message_history=self.messages,
+            model=create_model(self.user_token, self.host, self.workspace_id, run_id=deps.run_id),
+        )
+        self.messages.extend(result.new_messages())
+        return result

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/modules/agent/tools/datafile.py

@@ -19,7 +19,7 @@ from tinybird.tb.modules.feedback_manager import FeedbackManager
 def create_datafile(
     ctx: RunContext[TinybirdAgentContext], name: str, type: str, description: str, content: str, pathname: str
 ) -> str:
-    """Given a resource representation, create or update a datafile in the project folder
+    """Given a resource representation, create or update a datafile (.datasource, .connection, .pipe) in the project folder
 
     Args:
         name (str): The name of the datafile. Required.

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/modules/agent/tools/deploy.py

@@ -35,7 +35,12 @@ def deploy(ctx: RunContext[TinybirdAgentContext], allow_destructive_operations:
             ctx.deps.thinking_animation.start()
             return f"User did not confirm deployment and gave the following feedback: {feedback}"
 
-        click.echo(FeedbackManager.highlight(message="» Deploying project..."))
+        allow_destructive_operations_flag = " --allow-destructive-operations" if allow_destructive_operations else ""
+        click.echo(
+            FeedbackManager.highlight(
+                message=f"» Running command: tb --cloud deploy{allow_destructive_operations_flag}"
+            )
+        )
         ctx.deps.deploy_project(allow_destructive_operations=allow_destructive_operations)
         click.echo(FeedbackManager.success(message="✓ Project deployed successfully"))
         ctx.deps.thinking_animation.start()

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/modules/agent/tools/deploy_check.py

@@ -1,7 +1,9 @@
+import click
 from pydantic_ai import RunContext
 
 from tinybird.tb.modules.agent.utils import TinybirdAgentContext, show_confirmation, show_input
 from tinybird.tb.modules.exceptions import CLIDeploymentException
+from tinybird.tb.modules.feedback_manager import FeedbackManager
 
 
 def deploy_check(ctx: RunContext[TinybirdAgentContext]) -> str:
@@ -18,6 +20,7 @@ def deploy_check(ctx: RunContext[TinybirdAgentContext]) -> str:
             ctx.deps.thinking_animation.start()
             return f"User did not confirm deployment check and gave the following feedback: {feedback}"
 
+        click.echo(FeedbackManager.highlight(message="» Running command: tb --cloud deploy --check"))
         ctx.deps.deploy_check_project()
         ctx.deps.thinking_animation.start()
         return "Project can be deployed"

{tinybird-0.0.1.dev282 → tinybird-0.0.1.dev284}/tinybird/tb/modules/agent/tools/execute_query.py

@@ -8,7 +8,7 @@ import click
 import humanfriendly
 from pydantic_ai import RunContext
 
-from tinybird.tb.modules.agent.utils import TinybirdAgentContext
+from tinybird.tb.modules.agent.utils import TinybirdAgentContext, show_env_options
 from tinybird.tb.modules.common import echo_safe_humanfriendly_tables_format_pretty_table
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 
@@ -35,23 +35,24 @@ def execute_query(
     ctx: RunContext[TinybirdAgentContext],
     query: str,
     task: str,
-    cloud: Optional[bool] = None,
+    cloud_or_local: Optional[str] = None,
     script: Optional[str] = None,
     export_format: Optional[str] = None,
+    explanation_why_not_know_about_last_environment: Optional[str] = None,
 ):
     """Execute a query and return the result as a table, chart or exported file.
 
     Args:
         query (str): The query to execute. Required.
         task (str): The purpose of the query. Required.
-        cloud (bool): Whether to execute the query on cloud or local. If None (user didn't specify), will ask user to clarify. Defaults to local (False) in dangerous skip permissions mode.
+        cloud_or_local (str): Whether to execute the query on cloud or local. Use the last environment used in previous queries or endpoint requests. If you don't have any information about the last environment, use None. Options: cloud, local.
         script (str): Python script using plotext to render the query results as a chart. The script will have access to 'data' (list of dicts), 'meta' (list of column info dicts), 'terminal_width' and 'terminal_height' variables. Always use plt.theme("clear") for transparent background and plt.plot_size(terminal_width, terminal_height) for proper sizing. For bar charts, use the simple versions: plt.simple_bar(), plt.simple_multiple_bar(), and plt.simple_stacked_bar(). Optional.
         export_format (str): The format to export the query results to. Options: csv, json, ndjson. Optional.
+        explanation_why_not_know_about_last_environment (str): Why you don't know about the last environment used in previous queries or endpoint requests. Required.
 
     Returns:
         str: The result of the query.
     """
-
     try:
         for forbidden_command in forbidden_commands:
             if forbidden_command in query.lower():
@@ -61,24 +62,21 @@ def execute_query(
             if query.lower().startswith(forbidden_command):
                 return f"Error executing query: {forbidden_command} is not allowed."
 
-        # Handle cloud parameter - ask user if uncertain and not in dangerous skip mode
-        if cloud is None:
+        # Handle cloud_or_local parameter - ask user if uncertain and not in dangerous skip mode
+        if cloud_or_local is None:
             if ctx.deps.dangerously_skip_permissions:
                 # Default to local when in dangerous skip mode
-                cloud = False
+                cloud_or_local = "local"
             else:
                 # Ask the user to choose execution mode
-                from tinybird.tb.modules.agent.utils import show_env_options
 
                 cloud = show_env_options(ctx)
                 if cloud is None:
                     return "Query execution cancelled by user."
+                cloud_or_local = "cloud" if cloud else "local"
 
-        cloud_or_local = "cloud" if cloud else "local"
         ctx.deps.thinking_animation.stop()
-
         click.echo(FeedbackManager.highlight(message=f"» Executing query to {cloud_or_local}:\n{query}\n"))
-
         is_templating = query.strip().startswith("%")
         query_format = "JSON"
         if export_format == "csv":
@@ -94,7 +92,7 @@ def execute_query(
         else:
             query = f"SELECT * FROM ({query}) FORMAT {query_format}"
 
-        execute_query = ctx.deps.execute_query_cloud if cloud else ctx.deps.execute_query_local
+        execute_query = ctx.deps.execute_query_cloud if cloud_or_local == "cloud" else ctx.deps.execute_query_local
         result = execute_query(query=query)
         if export_format:
             file_extension = f".{export_format}"

tinybird-0.0.1.dev284/tinybird/tb/modules/agent/tools/mock.py

@@ -0,0 +1,135 @@
+import click
+from pydantic_ai import RunContext
+
+from tinybird.tb.modules.agent.utils import (
+    AgentRunCancelled,
+    TinybirdAgentContext,
+    create_terminal_box,
+    show_confirmation,
+    show_input,
+)
+from tinybird.tb.modules.common import echo_safe_humanfriendly_tables_format_pretty_table, format_data_to_ndjson
+from tinybird.tb.modules.datafile.fixture import persist_fixture
+from tinybird.tb.modules.feedback_manager import FeedbackManager
+
+
+def generate_mock_fixture(
+    ctx: RunContext[TinybirdAgentContext], datasource_name: str, sql: str, data_format: str, rows: int, task: str
+) -> str:
+    """Given a datasource name and a sql query to execute, generate a fixture file with mock data and append it to the datasource.
+
+    Args:
+        datasource_name (str): Name of the datasource to create mock data for. Required.
+        sql (str): SQL query to execute to generate the mock data. Required.
+        data_format (str): Format of the mock data to create. Options: ndjson, csv. Required.
+        rows (int): Number of rows to create. If not provided, the default is 10. Required.
+        task (str): Extra details about how to generate the mock data (nested json if any, sample row to help with the generation, etc). You can use this to fix issues with the mock data generation. Required.
+
+    Returns:
+        str: Message indicating the success or failure of the mock data generation
+    """
+    try:
+        ctx.deps.thinking_animation.stop()
+
+        click.echo(FeedbackManager.highlight(message=f"» Generating mock data for datasource '{datasource_name}'..."))
+        try:
+            sql_format = "JSON" if data_format == "ndjson" else "CSV"
+            sql = f"SELECT * FROM ({sql}) LIMIT {rows} FORMAT {sql_format}"
+            result = ctx.deps.execute_query_local(query=sql)
+        except Exception as e:
+            click.echo(
+                FeedbackManager.error(message="✗ Failed to generate a valid SQL query for generating mock data.\n{e}")
+            )
+            ctx.deps.thinking_animation.start()
+            return f"Failed to generate a valid sql for generating mock data for datasource '{datasource_name}'. SQL: {sql}\nError: {e}"
+
+        preview_content = ""
+        if sql_format == "JSON":
+            data = result.get("data", [])[:rows]
+            preview_content = str(format_data_to_ndjson(data[:10]))
+
+            if len(data) != rows:
+                raise Exception(
+                    f"Failed to generate a valid sql for generating mock data for datasource '{datasource_name}'. Rows generated: {len(data)} != {rows}. SQL: {sql}\nError: {result}"
+                )
+
+            error_response = result.get("error", None)
+            if error_response:
+                raise Exception(
+                    f"Failed to generate a valid sql for generating mock data for datasource '{datasource_name}'. SQL: {sql}\nError: {error_response}"
+                )
+
+        else:
+            data = result
+            preview_content = str(data[:1000])
+
+        if isinstance(data, dict):
+            data = data.get("data", [])[:rows]
+            preview_content = str(format_data_to_ndjson(data[:10]))
+
+        content = create_terminal_box(preview_content, title=f"fixtures/{datasource_name}.{data_format}")
+        click.echo(content)
+        click.echo("Showing a preview of the file.\n")
+        confirmation = show_confirmation(
+            title=f"Create fixture file for datasource '{datasource_name}'?",
+            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+        )
+
+        if confirmation == "review":
+            feedback = show_input(ctx.deps.workspace_name)
+            ctx.deps.thinking_animation.start()
+            return f"User did not confirm mock data for datasource '{datasource_name}' in Tinybird Local and gave the following feedback: {feedback}"
+
+        fixture_path = persist_fixture(datasource_name, data, ctx.deps.folder, format=data_format)
+        fixture_path_name = f"fixtures/{fixture_path.name}"
+        click.echo(FeedbackManager.success(message=f"✓ {fixture_path_name} created"))
+        confirmation = show_confirmation(
+            title=f"Append {fixture_path_name} to datasource '{datasource_name}'?",
+            skip_confirmation=ctx.deps.dangerously_skip_permissions,
+        )
+        if confirmation == "review":
+            feedback = show_input(ctx.deps.workspace_name)
+            ctx.deps.thinking_animation.start()
+            return f"Mock data was generated in {fixture_path_name} but user did not confirm appending {fixture_path_name} to datasource '{datasource_name}' in Tinybird Local and gave the following feedback: {feedback}"
+
+        ctx.deps.append_data_local(datasource_name=datasource_name, path=str(fixture_path))
+        click.echo(FeedbackManager.success(message=f"✓ Data appended to datasource '{datasource_name}'"))
+        ctx.deps.thinking_animation.start()
+        return f"Mock data generated in {fixture_path_name} and appended to datasource '{datasource_name}' in Tinybird Local"
+    except AgentRunCancelled as e:
+        raise e
+    except Exception as e:
+        ctx.deps.thinking_animation.stop()
+        error_message = str(e)
+        click.echo(FeedbackManager.error(message=error_message))
+        try:
+            if "in quarantine" in error_message:
+                click.echo(
+                    FeedbackManager.highlight(message=f"» Looking for errors in {datasource_name}_quarantine...")
+                )
+                query = f"select * from {datasource_name}_quarantine order by insertion_date desc limit 5 FORMAT JSON"
+                quarantine_result = ctx.deps.execute_query_local(query=query)
+                quarantine_data = quarantine_result["data"] or []
+                quarantine_meta = quarantine_result["meta"] or []
+                column_names = [c["name"] for c in quarantine_meta]
+                echo_safe_humanfriendly_tables_format_pretty_table(
+                    data=[d.values() for d in quarantine_data], column_names=column_names
+                )
+                click.echo(
+                    FeedbackManager.info(
+                        message=f"These are the first 5 rows of the quarantine table for datasource '{datasource_name}':"
+                    )
+                )
+                error_message = (
+                    error_message
+                    + f"\nThese are the first 5 rows of the quarantine table for datasource '{datasource_name}':\n{quarantine_data}. Use again `mock` tool but add this issue to the context."
+                )
+
+        except Exception as quarantine_error:
+            error_message = error_message + f"\nError accessing to {datasource_name}_quarantine: {quarantine_error}"
+
+        if "must be created first with 'mode=create'" in error_message:
+            error_message = error_message + "\nBuild the project again."
+
+        ctx.deps.thinking_animation.start()
+        return f"Error generating mock data for datasource '{datasource_name}' in Tinybird Local: {error_message}"