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

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)
 
 
@@ -908,7 +948,7 @@ def get_format_from_filename_or_url(filename_or_url: str) -> str:
     'csv'
     """
     filename_or_url = filename_or_url.lower()
-    if filename_or_url.endswith("json") or filename_or_url.endswith("ndjson"):
+    if filename_or_url.endswith(("json", "ndjson")):
         return "ndjson"
     if filename_or_url.endswith("parquet"):
         return "parquet"
@@ -916,7 +956,7 @@ def get_format_from_filename_or_url(filename_or_url: str) -> str:
         return "csv"
     try:
         parsed = urlparse(filename_or_url)
-        if parsed.path.endswith("json") or parsed.path.endswith("ndjson"):
+        if parsed.path.endswith(("json", "ndjson")):
             return "ndjson"
         if parsed.path.endswith("parquet"):
             return "parquet"
@@ -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,
@@ -2197,18 +2132,17 @@ def ask_for_organization(
         user_token = get_user_token(config, user_token)
         organization = create_organization_and_add_workspaces(config, organization_name, user_token)
         organization_id = organization.get("id")
+    elif len(organizations) == 1:
+        organization_name = organizations[0]["name"]
+        organization_id = organizations[0]["id"]
     else:
-        if len(organizations) == 1:
-            organization_name = organizations[0]["name"]
-            organization_id = organizations[0]["id"]
+        sorted_organizations = sort_organizations_by_user(organizations, user_email=user_email)
+        current_organization = ask_for_organization_interactively(sorted_organizations)
+        if current_organization:
+            organization_id = current_organization.get("id")
+            organization_name = current_organization.get("name")
         else:
-            sorted_organizations = sort_organizations_by_user(organizations, user_email=user_email)
-            current_organization = ask_for_organization_interactively(sorted_organizations)
-            if current_organization:
-                organization_id = current_organization.get("id")
-                organization_name = current_organization.get("name")
-            else:
-                return None, None
+            return None, None
     return organization_id, organization_name
 
 
@@ -2262,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)