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

tinybird/tb/modules/build.py

@@ -3,7 +3,7 @@ import time
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 from urllib.parse import urlencode
 
 import click
@@ -35,7 +35,10 @@ def build(ctx: click.Context, watch: bool) -> None:
     project: Project = ctx.ensure_object(dict)["project"]
     tb_client: TinyB = ctx.ensure_object(dict)["client"]
     config: Dict[str, Any] = ctx.ensure_object(dict)["config"]
-    if obj["env"] == "cloud":
+    is_branch = bool(ctx.ensure_object(dict)["branch"])
+
+    # TODO: Explain that you can use custom branches too once they are open for everyone
+    if obj["env"] == "cloud" and not is_branch:
         raise click.ClickException(FeedbackManager.error_build_only_supported_in_local())
 
     if project.has_deeper_level():
@@ -48,47 +51,85 @@ def build(ctx: click.Context, watch: bool) -> None:
         )
 
     click.echo(FeedbackManager.highlight_building_project())
-    process(project=project, tb_client=tb_client, watch=False, config=config)
+    process(project=project, tb_client=tb_client, watch=False, config=config, is_branch=is_branch)
     if watch:
         run_watch(
             project=project,
             tb_client=tb_client,
             config=config,
-            process=partial(process, project=project, tb_client=tb_client, watch=True, config=config),
+            process=partial(
+                process,
+                project=project,
+                tb_client=tb_client,
+                watch=True,
+                config=config,
+                is_branch=is_branch,
+            ),
         )
 
 
 @cli.command("dev", help="Build the project server side and watch for changes.")
 @click.option("--data-origin", type=str, default="", help="Data origin: local or cloud")
-@click.option("--ui", is_flag=True, default=False, help="Connect your local project to Tinybird UI")
+@click.option("--ui/--skip-ui", is_flag=True, default=True, help="Connect your local project to Tinybird UI")
 @click.pass_context
 def dev(ctx: click.Context, data_origin: str, ui: bool) -> None:
+    obj: Dict[str, Any] = ctx.ensure_object(dict)
+    branch: Optional[str] = ctx.ensure_object(dict)["branch"]
+    is_branch = bool(branch)
+
+    if obj["env"] == "cloud" and not is_branch:
+        raise click.ClickException(FeedbackManager.error_build_only_supported_in_local())
+
     if data_origin == "cloud":
+        click.echo(
+            FeedbackManager.warning(
+                message="--data-origin=cloud is deprecated and will be removed in a future version. Create an branch and use `tb --branch <branch_name> dev`"
+            )
+        )
         return dev_cloud(ctx)
+
     project: Project = ctx.ensure_object(dict)["project"]
     tb_client: TinyB = ctx.ensure_object(dict)["client"]
     config: Dict[str, Any] = ctx.ensure_object(dict)["config"]
+
     build_status = BuildStatus()
     if ui:
         server_thread = threading.Thread(
-            target=start_server, args=(project, tb_client, process, build_status), daemon=True
+            target=start_server, args=(project, tb_client, process, build_status, branch), daemon=True
         )
         server_thread.start()
         # Wait for the server to start
         time.sleep(0.5)
 
     click.echo(FeedbackManager.highlight_building_project())
-    process(project=project, tb_client=tb_client, watch=True, config=config, build_status=build_status)
+    process(
+        project=project,
+        tb_client=tb_client,
+        watch=True,
+        config=config,
+        build_status=build_status,
+        is_branch=is_branch,
+    )
     run_watch(
         project=project,
         tb_client=tb_client,
         config=config,
-        process=partial(process, project=project, tb_client=tb_client, build_status=build_status, config=config),
+        branch=branch,
+        process=partial(
+            process,
+            project=project,
+            tb_client=tb_client,
+            build_status=build_status,
+            config=config,
+            is_branch=is_branch,
+        ),
     )
 
 
-def run_watch(project: Project, tb_client: TinyB, process: Callable, config: dict[str, Any]) -> None:
-    shell = Shell(project=project, tb_client=tb_client, playground=False)
+def run_watch(
+    project: Project, tb_client: TinyB, process: Callable, config: dict[str, Any], branch: Optional[str] = None
+) -> None:
+    shell = Shell(project=project, tb_client=tb_client, branch=branch)
     click.echo(FeedbackManager.gray(message="\nWatching for changes..."))
     watcher_thread = threading.Thread(
         target=watch_project,

tinybird/tb/modules/build_common.py

@@ -33,15 +33,17 @@ def process(
     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:
+    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
-    build_shared_with_workspaces(project=project, tb_client=tb_client, config=config)
+    if not is_branch:
+        build_shared_with_workspaces(project=project, tb_client=tb_client, config=config)
 
     build_failed = False
     build_error: Optional[str] = None

tinybird/tb/modules/cli.py

@@ -42,6 +42,7 @@ 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
@@ -76,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"
@@ -103,6 +105,7 @@ def cli(
     version_warning: bool,
     show_tokens: bool,
     cloud: bool,
+    branch: Optional[str],
     staging: bool,
     output: str,
     max_depth: int,
@@ -200,21 +203,33 @@ 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"]):
+        if any(arg in sys.argv for arg in ["--cloud", "--local", "--branch"]):
             raise CLIException(
                 FeedbackManager.error(
-                    message="Tinybird Code does not support --cloud or --local flags. It will choose the correct environment based on your prompts."
+                    message="Tinybird Code does not support --cloud, --local or --branch flags. It will choose the correct environment based on your prompts."
                 )
             )
 
@@ -484,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",
@@ -498,6 +519,7 @@ def create_ctx_client(
         "tag",
         "push",
         "branch",
+        "environment",
         "diff",
         "fmt",
         "init",
@@ -506,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", "create"]
+    commands_always_cloud = ["infra", "branch", "environment"]
+    commands_always_local = ["create"]
     command_always_test = ["test"]
 
     if (
@@ -530,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:
@@ -538,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"
 

tinybird/tb/modules/common.py

@@ -30,7 +30,6 @@ import requests
 from click import Context
 from click._termui_impl import ProgressBar
 from humanfriendly.tables import format_pretty_table
-from packaging.version import Version
 from thefuzz import process
 
 from tinybird.tb.client import (
@@ -61,7 +60,6 @@ from tinybird.tb.modules.exceptions import (
     CLIAuthException,
     CLIConnectionException,
     CLIException,
-    CLIReleaseException,
     CLIWorkspaceException,
 )
 from tinybird.tb.modules.feedback_manager import FeedbackManager, warning_message
@@ -378,9 +376,33 @@ def getenv_bool(key: str, default: bool) -> bool:
     return v.lower() == "true" or v == "1"
 
 
-def _get_tb_client(token: str, host: str, staging: bool = False) -> TinyB:
+def _get_tb_client(token: str, host: str, staging: bool = False, branch: Optional[str] = None) -> TinyB:
     disable_ssl: bool = getenv_bool("TB_DISABLE_SSL_CHECKS", False)
-    return TinyB(token, host, version=VERSION, disable_ssl_checks=disable_ssl, send_telemetry=True, staging=staging)
+    cloud_client = TinyB(
+        token,
+        host,
+        version=VERSION,
+        disable_ssl_checks=disable_ssl,
+        send_telemetry=True,
+        staging=staging,
+    )
+
+    if not branch:
+        return cloud_client
+
+    workspaces = cloud_client.user_workspaces_and_branches(version="v1")
+    workspace = next((w for w in workspaces.get("workspaces", []) if w.get("name") == branch), None)
+    if not workspace:
+        raise CLIException(FeedbackManager.error_exception(error=f"Branch {branch} not found"))
+
+    return TinyB(
+        workspace.get("token", ""),
+        host,
+        version=VERSION,
+        disable_ssl_checks=disable_ssl,
+        send_telemetry=True,
+        staging=staging,
+    )
 
 
 def create_tb_client(ctx: Context) -> TinyB:
@@ -414,7 +436,16 @@ def analyze_file(filename: str, client: TinyB, format: str):
     meta, data = _analyze(filename, client, format)
     schema = meta["analysis"]["schema"]
     schema = schema.replace(", ", ",\n    ")
-    content = f"""DESCRIPTION >\n    Generated from {filename}\n\nSCHEMA >\n    {schema}"""
+    content = f"""DESCRIPTION >
+    Generated from {filename}
+
+SCHEMA >
+    {schema}
+
+ENGINE "MergeTree"
+# ENGINE_SORTING_KEY "user_id, timestamp"
+# ENGINE_TTL "timestamp + toIntervalDay(60)"
+# Learn more at https://www.tinybird.co/docs/forward/dev-reference/datafiles/datasource-files"""
     return content
 
 
@@ -429,7 +460,16 @@ def _generate_datafile(
     meta, data = _analyze(filename, client, format, connector=connector)
     schema = meta["analysis"]["schema"]
     schema = schema.replace(", ", ",\n    ")
-    datafile = f"""DESCRIPTION >\n    Generated from {filename}\n\nSCHEMA >\n    {schema}"""
+    datafile = f"""DESCRIPTION >
+    Generated from {filename}
+
+SCHEMA >
+    {schema}
+
+ENGINE "MergeTree"
+# ENGINE_SORTING_KEY "user_id, timestamp"
+# ENGINE_TTL "timestamp + toIntervalDay(60)"
+# Learn more at https://www.tinybird.co/docs/forward/dev-reference/datafiles/datasource-files"""
     return generate_datafile(datafile, filename, data, force, _format=format, folder=folder)
 
 
@@ -965,7 +1005,10 @@ def push_data(
     cb.prev_done = 0  # type: ignore[attr-defined]
 
     if not silent:
-        click.echo(FeedbackManager.highlight(message=f"\n» Appending data to {datasource_name}..."))
+        if mode == "replace":
+            click.echo(FeedbackManager.highlight(message=f"\n» Replacing data in {datasource_name}..."))
+        else:
+            click.echo(FeedbackManager.highlight(message=f"\n» Appending data to {datasource_name}..."))
 
     if isinstance(url, list):
         urls = url
@@ -1165,10 +1208,26 @@ def _get_setting_value(connection, setting, sensitive_settings):
     return connection.get(setting, "")
 
 
-def switch_workspace(config: CLIConfig, workspace_name_or_id: str) -> None:
+def get_current_workspace_branches(config: CLIConfig) -> List[Dict[str, Any]]:
+    current_main_workspace: Optional[Dict[str, Any]] = get_current_main_workspace(config)
+    if not current_main_workspace:
+        raise CLIException(FeedbackManager.error_unable_to_identify_main_workspace())
+
+    client = config.get_client()
+    user_branches: List[Dict[str, Any]] = (client.user_workspace_branches("v1")).get("workspaces", [])
+    all_branches: List[Dict[str, Any]] = (client.branches()).get("environments", [])
+    branches = user_branches + [branch for branch in all_branches if branch not in user_branches]
+
+    return [branch for branch in branches if branch.get("main") == current_main_workspace["id"]]
+
+
+def switch_workspace(config: CLIConfig, workspace_name_or_id: str, only_environments: bool = False) -> None:
     try:
-        response = config.get_client().user_workspaces(version="v1")
-        workspaces = response["workspaces"]
+        if only_environments:
+            workspaces = get_current_workspace_branches(config)
+        else:
+            response = config.get_client().user_workspaces(version="v1")
+            workspaces = response["workspaces"]
 
         workspace = next(
             (
@@ -1563,130 +1622,6 @@ def _get_current_main_workspace_common(
     return current
 
 
-def is_post_semver(new_version: Version, current_version: Version) -> bool:
-    """
-    Check if only the post part of the semantic version has changed.
-
-    Args:
-    new_version (Version): The new version to check.
-    current_version (Version): The current version to compare with.
-
-    Returns:
-    bool: True if only the post part of the version has changed, False otherwise.
-
-    Examples:
-    >>> is_post_semver(Version("0.0.0-2"), Version("0.0.0-1"))
-    True
-    >>> is_post_semver(Version("0.0.0-1"), Version("0.0.0-1"))
-    False
-    >>> is_post_semver(Version("0.0.1-1"), Version("0.0.0-1"))
-    False
-    >>> is_post_semver(Version("0.1.0-1"), Version("0.0.0-1"))
-    False
-    >>> is_post_semver(Version("1.0.0-1"), Version("0.0.0-1"))
-    False
-    >>> is_post_semver(Version("0.0.1-1"), Version("0.0.0"))
-    False
-    """
-    if (
-        new_version.major == current_version.major
-        and new_version.minor == current_version.minor
-        and new_version.micro == current_version.micro
-    ):
-        return new_version.post is not None and new_version.post != current_version.post
-
-    return False
-
-
-def is_major_semver(new_version: Version, current_version: Version) -> bool:
-    """
-    Check if only the major part of the semantic version has changed.
-
-    Args:
-    new_version (Version): The new version to check.
-    current_version (Version): The current version to compare with.
-
-    Returns:
-    bool: True if only the major part of the version has changed, False otherwise.
-
-    Examples:
-    >>> is_major_semver(Version("1.0.0"), Version("0.0.0"))
-    True
-    >>> is_major_semver(Version("0.0.0"), Version("0.0.0"))
-    False
-    >>> is_major_semver(Version("1.0.1"), Version("1.0.0"))
-    False
-    >>> is_major_semver(Version("2.0.0-1"), Version("1.0.1-2"))
-    True
-    """
-
-    return new_version.major != current_version.major
-
-
-def print_release_summary(config: CLIConfig, semver: Optional[str], info: bool = False, dry_run=False):
-    if not semver:
-        click.echo(FeedbackManager.info_release_no_rollback())
-        return
-    try:
-        client = config.get_client()
-        response = client.release_rm(config["id"], semver, confirmation=config["name"], dry_run=True)
-    except Exception as e:
-        raise CLIReleaseException(FeedbackManager.error_exception(error=str(e)))
-    else:
-        columns = ["name", "id", "notes"]
-        if not response:
-            click.echo(FeedbackManager.info_release_no_rollback())
-            return
-
-        if len(response["datasources"]) or len(response["pipes"]):
-            semver = response.get("semver", semver)
-            if info:
-                if dry_run:
-                    click.echo(FeedbackManager.info_release_rm_resources_dry_run(semver=semver))
-                else:
-                    click.echo(FeedbackManager.info_release_rm_resources())
-            else:
-                click.echo(FeedbackManager.info_release_rollback(semver=semver))
-
-        if len(response["datasources"]):
-            click.echo("\nDatasources:")
-            rows = [
-                [ds, response["datasources"][ds], response["notes"].get(response["datasources"][ds], "")]
-                for ds in response["datasources"]
-            ]
-            echo_safe_humanfriendly_tables_format_smart_table(rows, column_names=columns)
-
-        if len(response["pipes"]):
-            click.echo("\nPipes:")
-            rows = [
-                [pipe, response["pipes"][pipe], response["notes"].get(response["pipes"][pipe], "")]
-                for pipe in response["pipes"]
-            ]
-            echo_safe_humanfriendly_tables_format_smart_table(rows, column_names=columns)
-
-
-def get_oldest_rollback(config: CLIConfig, client: TinyB) -> Optional[str]:
-    oldest_rollback_response = client.release_oldest_rollback(config["id"])
-    return oldest_rollback_response.get("semver") if oldest_rollback_response else None
-
-
-def remove_release(
-    dry_run: bool, config: CLIConfig, semver: Optional[str], client: TinyB, force: bool, show_print=True
-):
-    if semver == OLDEST_ROLLBACK:
-        semver = get_oldest_rollback(config, client)
-    if show_print:
-        print_release_summary(config, semver, info=True, dry_run=True)
-    if not dry_run:
-        if semver:
-            response = client.release_rm(
-                config["id"], semver, confirmation=config["name"], dry_run=dry_run, force=force
-            )
-            click.echo(FeedbackManager.success_release_delete(semver=response.get("semver")))
-        else:
-            click.echo(FeedbackManager.info_no_release_deleted())
-
-
 def run_aws_iamrole_connection_flow(
     client: TinyB,
     service: str,
@@ -2261,3 +2196,95 @@ def update_cli() -> None:
         click.echo(FeedbackManager.success(message="✓ Tinybird CLI updated"))
     else:
         click.echo(FeedbackManager.info(message="✓ Tinybird CLI is already up-to-date"))
+
+
+def create_workspace_branch(
+    branch_name: Optional[str],
+    last_partition: bool,
+    all: bool,
+    ignore_datasources: Optional[List[str]],
+    wait: Optional[bool],
+) -> None:
+    """
+    Creates a workspace branch
+    """
+    config = CLIConfig.get_project_config()
+    _ = try_update_config_with_remote(config)
+
+    try:
+        workspace = get_current_workspace(config)
+        if not workspace:
+            raise CLIWorkspaceException(FeedbackManager.error_workspace())
+
+        if not branch_name:
+            click.echo(FeedbackManager.info_workspace_branch_create_greeting())
+            default_name = f"{workspace['name']}_{uuid.uuid4().hex[0:4]}"
+            branch_name = click.prompt("\Branch name", default=default_name, err=True, type=str)
+        assert isinstance(branch_name, str)
+
+        response = config.get_client().create_workspace_branch(
+            branch_name,
+            last_partition,
+            all,
+            ignore_datasources,
+        )
+        assert isinstance(response, dict)
+
+        is_job: bool = "job" in response
+        is_summary: bool = "partitions" in response
+
+        if not is_job and not is_summary:
+            raise CLIException(str(response))
+
+        if all and not is_job:
+            raise CLIException(str(response))
+
+        click.echo(
+            FeedbackManager.success_workspace_branch_created(workspace_name=workspace["name"], branch_name=branch_name)
+        )
+
+        job_id: Optional[str] = None
+
+        if is_job:
+            job_id = response["job"]["job_id"]
+            job_url = response["job"]["job_url"]
+            click.echo(FeedbackManager.info_data_branch_job_url(url=job_url))
+
+        if wait and is_job:
+            assert isinstance(job_id, str)
+
+            # Await the job to finish and get the result dict
+            job_response = wait_job(config.get_client(), job_id, job_url, "Environment creation")
+            if job_response is None:
+                raise CLIException(f"Empty job API response (job_id: {job_id}, job_url: {job_url})")
+            else:
+                response = job_response.get("result", {})
+                is_summary = "partitions" in response
+
+    except Exception as e:
+        raise CLIException(FeedbackManager.error_exception(error=str(e)))
+
+
+async def print_current_branch(config: CLIConfig) -> None:
+    _ = try_update_config_with_remote(config, only_if_needed=True)
+
+    response = config.get_client().user_workspaces_and_branches("v1")
+
+    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 = 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)