tinybird 0.0.1.dev5__py3-none-any.whl → 0.0.1.dev7__py3-none-any.whl

tinybird/tb/modules/datafile/pull.py

@@ -0,0 +1,157 @@
+import re
+from asyncio import Semaphore, gather
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import aiofiles
+import click
+
+from tinybird.client import AuthNoTokenException, TinyB
+from tinybird.feedback_manager import FeedbackManager
+from tinybird.tb.modules.datafile.common import get_name_version
+from tinybird.tb.modules.datafile.format_datasource import format_datasource
+from tinybird.tb.modules.datafile.format_pipe import format_pipe
+
+
+async def folder_pull(
+    client: TinyB,
+    folder: str,
+    auto: bool,
+    match: Optional[str],
+    force: bool,
+    verbose: bool = True,
+    progress_bar: bool = False,
+    fmt: bool = False,
+):
+    pattern = re.compile(match) if match else None
+
+    def _get_latest_versions(resources: List[str]):
+        versions: Dict[str, Any] = {}
+
+        for x in resources:
+            t = get_name_version(x)
+            t["original_name"] = x
+            if t["version"] is None:
+                t["version"] = -1
+            name = t["name"]
+
+            if name not in versions or name == x or versions[name]["version"] < t["version"]:
+                versions[name] = t
+        return versions
+
+    def get_file_folder(extension: str):
+        if not auto:
+            return None
+        if extension == "datasource":
+            return "datasources"
+        if extension == "pipe":
+            return "pipes"
+        if extension == "token":
+            return "tokens"
+        return None
+
+    async def write_files(
+        versions: Dict[str, Any],
+        resources: List[str],
+        extension: str,
+        get_resource_function: str,
+        progress_bar: bool = False,
+        fmt: bool = False,
+    ):
+        async def write_resource(k: Dict[str, Any]):
+            name = f"{k['name']}.{extension}"
+            try:
+                if pattern and not pattern.search(name):
+                    if verbose:
+                        click.echo(FeedbackManager.info_skipping_resource(resource=name))
+                    return
+
+                resource = await getattr(client, get_resource_function)(k["original_name"])
+                resource_to_write = resource
+
+                if fmt:
+                    if extension == "datasource":
+                        resource_to_write = await format_datasource(name, content=resource)
+                    elif extension == "pipe":
+                        resource_to_write = await format_pipe(name, content=resource)
+
+                dest_folder = folder
+                if "." in k["name"] and extension != "token":
+                    dest_folder = Path(folder) / "vendor" / k["name"].split(".", 1)[0]
+                    name = f"{k['name'].split('.', 1)[1]}.{extension}"
+
+                file_folder = get_file_folder(extension)
+                f = Path(dest_folder) / file_folder if file_folder is not None else Path(dest_folder)
+
+                if not f.exists():
+                    f.mkdir(parents=True)
+
+                f = f / name
+                resource_names = [x.split(".")[-1] for x in resources]
+
+                if verbose:
+                    click.echo(FeedbackManager.info_writing_resource(resource=f))
+                if not f.exists() or force:
+                    async with aiofiles.open(f, "w") as fd:
+                        # versions are a client only thing so
+                        # datafiles from the server do not contains information about versions
+                        if k["version"] >= 0:
+                            if fmt:
+                                resource_to_write = "\n" + resource_to_write  # fmt strips the first line
+
+                            resource_to_write = f"VERSION {k['version']}\n" + resource_to_write
+                        if resource_to_write:
+                            matches = re.findall(r"([^\s\.]*__v\d+)", resource_to_write)
+                            for match in set(matches):
+                                m = match.split("__v")[0]
+                                if m in resources or m in resource_names:
+                                    resource_to_write = resource_to_write.replace(match, m)
+                            await fd.write(resource_to_write)
+                else:
+                    if verbose:
+                        click.echo(FeedbackManager.info_skip_already_exists())
+            except Exception as e:
+                raise click.ClickException(FeedbackManager.error_exception(error=e))
+
+        values = versions.values()
+
+        if progress_bar:
+            with click.progressbar(values, label=f"Pulling {extension}s") as values:  # type: ignore
+                for k in values:
+                    await write_resource(k)
+        else:
+            tasks = [write_resource(k) for k in values]
+            await _gather_with_concurrency(5, *tasks)
+
+    try:
+        datasources = await client.datasources()
+        remote_datasources = sorted([x["name"] for x in datasources])
+        datasources_versions = _get_latest_versions(remote_datasources)
+
+        pipes = await client.pipes()
+        remote_pipes = sorted([x["name"] for x in pipes])
+        pipes_versions = _get_latest_versions(remote_pipes)
+
+        resources = list(datasources_versions.keys()) + list(pipes_versions.keys())
+
+        await write_files(
+            datasources_versions, resources, "datasource", "datasource_file", progress_bar=progress_bar, fmt=fmt
+        )
+        await write_files(pipes_versions, resources, "pipe", "pipe_file", progress_bar=progress_bar, fmt=fmt)
+
+        return
+
+    except AuthNoTokenException:
+        raise
+    except Exception as e:
+        raise click.ClickException(FeedbackManager.error_pull(error=str(e)))
+
+
+async def _gather_with_concurrency(n, *tasks):
+    semaphore = Semaphore(n)
+
+    async def sem_task(task):
+        async with semaphore:
+            return await task
+
+    return await gather(*(sem_task(task) for task in tasks))

tinybird/{tb_cli_modules → tb/modules}/datasource.py

@@ -15,16 +15,15 @@ from click import Context
 
 from tinybird.client import AuthNoTokenException, CanNotBeDeletedException, DoesNotExistException, TinyB
 from tinybird.config import get_display_host
-from tinybird.tb_cli_modules.config import CLIConfig
+from tinybird.tb.modules.config import CLIConfig
 
 if TYPE_CHECKING:
     from tinybird.connectors import Connector
 
-from tinybird.datafile import get_name_version, wait_job
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.branch import warn_if_in_live
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import (
+from tinybird.tb.modules.branch import warn_if_in_live
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import (
     _analyze,
     _generate_datafile,
     ask_for_user_token,
@@ -40,8 +39,10 @@ from tinybird.tb_cli_modules.common import (
     validate_kafka_auto_offset_reset,
     validate_kafka_group,
     validate_kafka_topic,
+    wait_job,
 )
-from tinybird.tb_cli_modules.exceptions import CLIDatasourceException
+from tinybird.tb.modules.datafile.common import get_name_version
+from tinybird.tb.modules.exceptions import CLIDatasourceException
 
 
 @cli.group()

tinybird/tb/modules/exceptions.py

@@ -0,0 +1,91 @@
+from typing import Any, Dict, Optional
+
+import click
+
+from tinybird.tb.modules.telemetry import add_telemetry_event
+
+
+class CLIException(click.exceptions.ClickException):
+    """Default exception for all exceptions raised in the CLI.
+
+    Allows to specift a custom telemetry event to be sent before
+    raising the exception (if not specified, it sends a generic
+    `error` beacon with the error message)
+    """
+
+    def __init__(self, message: str, telemetry_event: Optional[str] = None, **kw_telemetry_event_data: Any) -> None:
+        telemetry_event = telemetry_event or "error"
+        data: Dict[str, Any] = {"error": message}
+        data.update(kw_telemetry_event_data)
+        add_telemetry_event(telemetry_event, **data)
+        super().__init__(message)
+
+
+class CLIAuthException(CLIException):
+    """Exceptions generated by the auth commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "auth_error", **kw_telemetry_event_data)
+
+
+class CLIReleaseException(CLIException):
+    """Exceptions generated by the release commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "release_error", **kw_telemetry_event_data)
+
+
+class CLIBranchException(CLIException):
+    """Exceptions generated by the branch commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "branch_error", **kw_telemetry_event_data)
+
+
+class CLIGitReleaseException(CLIException):
+    """Exceptions generated by the git release related commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "git_release_error", **kw_telemetry_event_data)
+
+
+class CLIConnectionException(CLIException):
+    """Exceptions generated by the connection commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "connection_error", **kw_telemetry_event_data)
+
+
+class CLIDatasourceException(CLIException):
+    """Exceptions generated by the datasource commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "datasource_error", **kw_telemetry_event_data)
+
+
+class CLIPipeException(CLIException):
+    """Exceptions generated by the pipe commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "pipe_error", **kw_telemetry_event_data)
+
+
+class CLIWorkspaceMembersException(CLIException):
+    """Exceptions generated by the workspace members commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "workspace_members_error", **kw_telemetry_event_data)
+
+
+class CLIWorkspaceException(CLIException):
+    """Exceptions generated by the workspace commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "workspace_error", **kw_telemetry_event_data)
+
+
+class CLITokenException(CLIException):
+    """Exceptions generated by the token commands"""
+
+    def __init__(self, message: str, **kw_telemetry_event_data: Any) -> None:
+        super().__init__(message, "token_error", **kw_telemetry_event_data)

tinybird/{tb_cli_modules → tb/modules}/fmt.py

@@ -7,10 +7,13 @@ import aiofiles
 import click
 from click import Context
 
-from tinybird.datafile import color_diff, format_datasource, format_pipe, is_file_a_datasource, peek
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import coro
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import coro
+from tinybird.tb.modules.datafile.common import is_file_a_datasource
+from tinybird.tb.modules.datafile.diff import color_diff, peek
+from tinybird.tb.modules.datafile.format_datasource import format_datasource
+from tinybird.tb.modules.datafile.format_pipe import format_pipe
 
 
 @cli.command()

tinybird/{tb_cli_modules → tb/modules}/job.py

@@ -8,9 +8,9 @@ from click import Context
 
 from tinybird.client import DoesNotExistException, TinyB
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import coro, echo_safe_humanfriendly_tables_format_smart_table
-from tinybird.tb_cli_modules.exceptions import CLIException
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import coro, echo_safe_humanfriendly_tables_format_smart_table
+from tinybird.tb.modules.exceptions import CLIException
 
 
 @cli.group()

tinybird/{tb_cli_modules → tb/modules}/llm.py

@@ -6,7 +6,7 @@ from openai import OpenAI
 from pydantic import BaseModel
 
 from tinybird.client import TinyB
-from tinybird.tb_cli_modules.prompts import create_project_prompt, sample_data_sql_prompt
+from tinybird.tb.modules.prompts import create_project_prompt, sample_data_sql_prompt
 
 
 class DataFile(BaseModel):

tinybird/{tb_cli_modules → tb/modules}/local.py

@@ -6,12 +6,12 @@ import requests
 
 import docker
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import (
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import (
     coro,
 )
-from tinybird.tb_cli_modules.config import CLIConfig
-from tinybird.tb_cli_modules.exceptions import CLIException
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.exceptions import CLIException
 
 # TODO: Use the official Tinybird image once it's available 'tinybirdco/tinybird-local:latest'
 TB_IMAGE_NAME = "registry.gitlab.com/tinybird/analytics/tinybird-local-jammy-3.11:latest"
@@ -114,7 +114,11 @@ def remove_tinybird_local(docker_client):
 def get_tinybird_local_client():
     """Get a Tinybird client connected to the local environment."""
     config = CLIConfig.get_project_config()
-    tokens = requests.get(f"{TB_LOCAL_HOST}/tokens").json()
+    try:
+        tokens = requests.get(f"{TB_LOCAL_HOST}/tokens").json()
+    except Exception:
+        raise CLIException("Tinybird local environment is not running. Please start it with `tb local start`.")
+
     token = tokens["workspace_admin_token"]
     config.set_token(token)
     config.set_host(TB_LOCAL_HOST)

tinybird/{tb_cli_modules → tb/modules}/mock.py

@@ -4,11 +4,11 @@ from pathlib import Path
 import click
 
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import CLIException, coro
-from tinybird.tb_cli_modules.config import CLIConfig
-from tinybird.tb_cli_modules.llm import LLM
-from tinybird.tb_cli_modules.local import get_tinybird_local_client
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import CLIException, coro
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.llm import LLM
+from tinybird.tb.modules.local import get_tinybird_local_client
 
 
 @cli.command()

tinybird/{tb_cli_modules → tb/modules}/pipe.py

@@ -16,12 +16,18 @@ from click import Context
 import tinybird.context as context
 from tinybird.client import AuthNoTokenException, DoesNotExistException, TinyB
 from tinybird.config import DEFAULT_API_HOST, FeatureFlags
-from tinybird.datafile import PipeNodeTypes, PipeTypes, folder_push, get_name_version, process_file, wait_job
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.branch import warn_if_in_live
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import coro, create_tb_client, echo_safe_humanfriendly_tables_format_smart_table
-from tinybird.tb_cli_modules.exceptions import CLIPipeException
+from tinybird.tb.modules.branch import warn_if_in_live
+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.build import folder_push, process_file
+from tinybird.tb.modules.datafile.common import PipeNodeTypes, PipeTypes, get_name_version
+from tinybird.tb.modules.exceptions import CLIPipeException
 
 
 @cli.group()

tinybird/{tb_cli_modules → tb/modules}/prompts.py

@@ -44,7 +44,7 @@ SQL >
     - Create multiple pipes to show different use cases over the same datasource.
     - The SQL query must be a valid ClickHouse SQL query that mixes ClickHouse syntax and Tinybird templating syntax.
     - If you need to use dynamic parameters you must start the whole sql query ALWAYS with "%" symbol on top. e.g: SQL >\n    %\n SELECT * FROM <table> WHERE <condition> LIMIT 10
-    - The Parameter functions like this one {{String(<my_param_name>,<default_value>)}} can be one of the following: String, DateTime, Date, Float32, Float64, Int, Integer, UInt8, UInt16, UInt32, UInt64, UInt128, UInt256, Int8, Int16, Int32, Int64, Int128, Int256
+    - The Parameter functions like this one {{String(my_param_name,default_value)}} can be one of the following: String, DateTime, Date, Float32, Float64, Int, Integer, UInt8, UInt16, UInt32, UInt64, UInt128, UInt256, Int8, Int16, Int32, Int64, Int128, Int256
     - Parameter names must be different from column names. Pass always the param name and a default value to the function.
     - Nodes can't have the same exact name as the Pipe they belong to.
 </instructions>

tinybird/tb/modules/regions.py

@@ -0,0 +1,9 @@
+from typing import Optional, TypedDict
+
+
+class Region(TypedDict):
+    name: str
+    provider: str
+    api_host: str
+    host: str
+    default_password: Optional[str]

tinybird/{tb_cli_modules → tb/modules}/tag.py

@@ -4,8 +4,8 @@ import click
 from click import Context
 
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import coro, echo_safe_humanfriendly_tables_format_smart_table
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import coro, echo_safe_humanfriendly_tables_format_smart_table
 
 
 @cli.group()