tinybird 0.0.1.dev4__py3-none-any.whl → 0.0.1.dev6__py3-none-any.whl

tinybird/tb/modules/mock.py

@@ -0,0 +1,53 @@
+import json
+from pathlib import Path
+
+import click
+
+from tinybird.feedback_manager import FeedbackManager
+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()
+@click.argument("datasource", type=str)
+@click.option("--rows", type=int, default=10, help="Number of events to send")
+@click.option("--context", type=str, default="", help="Extra context to use for data generation")
+@coro
+async def mock(datasource: str, rows: int, context: str) -> None:
+    """Load sample data into a datasource.
+
+    Args:
+        ctx: Click context object
+        datasource_file: Path to the datasource file to load sample data into
+    """
+    import llm
+
+    try:
+        datasource_path = Path(datasource)
+        datasource_name = datasource
+        if datasource_path.suffix == ".datasource":
+            datasource_name = datasource_path.stem
+        else:
+            datasource_path = Path("datasources", f"{datasource}.datasource")
+
+        datasource_content = datasource_path.read_text()
+        llm_config = CLIConfig.get_llm_config()
+        llm = LLM(key=llm_config["api_key"])
+        tb_client = get_tinybird_local_client()
+        sql = await llm.generate_sql_sample_data(tb_client, datasource_content, rows, context)
+        result = await tb_client.query(f"{sql} FORMAT JSON")
+        data = result.get("data", [])
+        max_rows_per_request = 100
+        sent_rows = 0
+        for i in range(0, len(data), max_rows_per_request):
+            batch = data[i : i + max_rows_per_request]
+            ndjson_data = "\n".join([json.dumps(row) for row in batch])
+            await tb_client.datasource_events(datasource_name, ndjson_data)
+            sent_rows += len(batch)
+        click.echo(f"Sent {sent_rows} events to datasource '{datasource_name}'")
+
+    except Exception as e:
+        raise CLIException(FeedbackManager.error_exception(error=str(e)))

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

@@ -16,12 +16,12 @@ 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
+from tinybird.tb.modules.datafile import PipeNodeTypes, PipeTypes, folder_push, get_name_version, process_file, wait_job
+from tinybird.tb.modules.exceptions import CLIPipeException
 
 
 @cli.group()

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

@@ -130,4 +130,7 @@ FROM numbers({row_count})
 - 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.
 
+# Extra context:
+{context}
+
 """

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/modules/table.py

@@ -0,0 +1,185 @@
+# Standard library modules.
+import collections
+import re
+
+# Modules included in our package.
+from humanfriendly.compat import coerce_string
+from humanfriendly.tables import format_robust_table
+from humanfriendly.terminal import (
+    ansi_strip,
+    ansi_width,
+    ansi_wrap,
+    find_terminal_size,
+    terminal_supports_colors,
+)
+
+from tinybird.feedback_manager import bcolors
+
+NUMERIC_DATA_PATTERN = re.compile(r"^\d+(\.\d+)?$")
+
+
+def format_table(data, columns):
+    """
+    Render tabular data using the most appropriate representation.
+
+    :param data: An iterable (e.g. a :func:`tuple` or :class:`list`)
+                 containing the rows of the table, where each row is an
+                 iterable containing the columns of the table (strings).
+    :param column_names: An iterable of column names (strings).
+    :returns: The rendered table (a string).
+
+    If you want an easy way to render tabular data on a terminal in a human
+    friendly format then this function is for you! It works as follows:
+
+    - If the input data doesn't contain any line breaks the function
+      :func:`format_pretty_table()` is used to render a pretty table. If the
+      resulting table fits in the terminal without wrapping the rendered pretty
+      table is returned.
+
+    - If the input data does contain line breaks or if a pretty table would
+      wrap (given the width of the terminal) then the function
+      :func:`format_robust_table()` is used to render a more robust table that
+      can deal with data containing line breaks and long text.
+    """
+    # Normalize the input in case we fall back from a pretty table to a robust
+    # table (in which case we'll definitely iterate the input more than once).
+    data = [normalize_columns(r) for r in data]
+
+    column_names = []
+    column_types = []
+
+    for c in columns:
+        column_names.append(c["name"])
+        column_types.append(c["type"])
+
+    column_names = normalize_columns(column_names)
+    # Make sure the input data doesn't contain any line breaks (because pretty
+    # tables break horribly when a column's text contains a line break :-).
+    if not any(any("\n" in c for c in r) for r in data):
+        # Render a pretty table.
+        pretty_table = format_pretty_table(data, column_names, column_types)
+        # Check if the pretty table fits in the terminal.
+        table_width = max(map(ansi_width, pretty_table.splitlines()))
+        num_rows, num_columns = find_terminal_size()
+        if table_width <= num_columns:
+            # The pretty table fits in the terminal without wrapping!
+            return pretty_table
+    # Fall back to a robust table when a pretty table won't work.
+    return format_robust_table(data, column_names)
+
+
+def format_pretty_table(data, column_names=None, column_types=None, horizontal_bar="─", vertical_bar=" "):
+    """
+    Render a table using characters like dashes and vertical bars to emulate borders.
+
+    :param data: An iterable (e.g. a :func:`tuple` or :class:`list`)
+                 containing the rows of the table, where each row is an
+                 iterable containing the columns of the table (strings).
+    :param column_names: An iterable of column names (strings).
+    :param column_types: An iterable of column types (strings).
+    :param horizontal_bar: The character used to represent a horizontal bar (a
+                           string).
+    :param vertical_bar: The character used to represent a vertical bar (a
+                         string).
+    :returns: The rendered table (a string).
+
+    Here's an example:
+
+    >>> from humanfriendly.tables import format_pretty_table
+    >>> column_names = ['Version', 'Uploaded on', 'Downloads']
+    >>> humanfriendly_releases = [
+    ... ['1.23', '2015-05-25', '218'],
+    ... ['1.23.1', '2015-05-26', '1354'],
+    ... ['1.24', '2015-05-26', '223'],
+    ... ['1.25', '2015-05-26', '4319'],
+    ... ['1.25.1', '2015-06-02', '197'],
+    ... ]
+    >>> print(format_pretty_table(humanfriendly_releases, column_names))
+    -------------------------------------
+    | Version | Uploaded on | Downloads |
+    -------------------------------------
+    | 1.23    | 2015-05-25  |       218 |
+    | 1.23.1  | 2015-05-26  |      1354 |
+    | 1.24    | 2015-05-26  |       223 |
+    | 1.25    | 2015-05-26  |      4319 |
+    | 1.25.1  | 2015-06-02  |       197 |
+    -------------------------------------
+
+    Notes about the resulting table:
+
+    - If a column contains numeric data (integer and/or floating point
+      numbers) in all rows (ignoring column names of course) then the content
+      of that column is right-aligned, as can be seen in the example above. The
+      idea here is to make it easier to compare the numbers in different
+      columns to each other.
+
+    - The column names are highlighted in color so they stand out a bit more
+      (see also :data:`.HIGHLIGHT_COLOR`). The following screen shot shows what
+      that looks like (my terminals are always set to white text on a black
+      background):
+
+      .. image:: images/pretty-table.png
+    """
+    # Normalize the input because we'll have to iterate it more than once.
+    data = [normalize_columns(r, expandtabs=True) for r in data]
+    if column_names is not None:
+        column_names = normalize_columns(column_names)
+        if column_names:
+            if terminal_supports_colors():
+                column_names = [highlight_column_name(n) for n in column_names]
+            data.insert(0, column_names)
+            if column_types is not None:
+                column_types = normalize_columns(column_types)
+                column_types = [highlight_column_type(t) for t in column_types]
+                data.insert(1, column_types)
+    # Calculate the maximum width of each column.
+    widths = collections.defaultdict(int)
+    numeric_data = collections.defaultdict(list)
+    for row_index, row in enumerate(data):
+        for column_index, column in enumerate(row):
+            widths[column_index] = max(widths[column_index], ansi_width(column))
+            if not (column_names and row_index <= 1):  # Skip both header rows for numeric check
+                numeric_data[column_index].append(bool(NUMERIC_DATA_PATTERN.match(ansi_strip(column))))
+
+    horizontal_bar = highlight_horizontal_bar(horizontal_bar)
+    # Create a horizontal bar of dashes as a delimiter.
+    line_delimiter = horizontal_bar * (sum(widths.values()) + len(widths) * 3 + 1)
+    # Start the table with a vertical bar.
+    lines = []
+    # Format the rows and columns.
+    for row_index, row in enumerate(data):
+        line = [vertical_bar]
+        for column_index, column in enumerate(row):
+            padding = " " * (widths[column_index] - ansi_width(column))
+            if all(numeric_data[column_index]):
+                line.append(" " + padding + column + " ")
+            else:
+                line.append(" " + column + padding + " ")
+            line.append(vertical_bar)
+        lines.append("".join(line))
+        if column_names and column_types and row_index > 0 and row_index < len(data) - 1:
+            lines.append(line_delimiter)
+    # Join the lines, returning a single string.
+    return "\n".join(lines)
+
+
+def normalize_columns(row, expandtabs=False):
+    results = []
+    for value in row:
+        text = coerce_string(value)
+        if expandtabs:
+            text = text.expandtabs()
+        results.append(text)
+    return results
+
+
+def highlight_column_name(name):
+    return ansi_wrap(name)
+
+
+def highlight_column_type(type):
+    return f"{bcolors.CGREY}\033[3m{type}\033[23m{bcolors.ENDC}"
+
+
+def highlight_horizontal_bar(horizontal_bar):
+    return f"{bcolors.CGREY}\033[3m{horizontal_bar}\033[23m{bcolors.ENDC}"

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()

tinybird/tb/modules/telemetry.py

@@ -0,0 +1,310 @@
+import functools
+import json
+import os
+import platform
+import re
+import sys
+import threading
+import uuid
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Tuple
+from urllib.parse import urlencode
+
+import requests
+
+from tinybird.config import CURRENT_VERSION
+
+TELEMETRY_TIMEOUT: int = 1
+TELEMETRY_DATASOURCE: str = "tb_cli_telemetry"
+
+
+def get_ci_product_name() -> Optional[str]:
+    if _is_env_true("TB_DISABLE_CI_DETECTION"):
+        return None
+
+    CI_CHECKS: List[Tuple[str, Callable[[], bool]]] = [
+        ("Azure pipelines", lambda: _is_env_true("TF_BUILD")),
+        ("GitHub Actions", lambda: _is_env_true("GITHUB_ACTIONS")),
+        ("Appveyor", lambda: _is_env_true("APPVEYOR")),
+        ("Travis CI", lambda: _is_env_true("TRAVIS")),
+        ("Circle CI", lambda: _is_env_true("CIRCLECI")),
+        ("Amazon Web Services CodeBuild", lambda: _is_env_present(["CODEBUILD_BUILD_ID", "AWS_REGION"])),
+        ("Jenkins", lambda: _is_env_present(["BUILD_ID", "BUILD_URL"])),
+        ("Google Cloud Build", lambda: _is_env_present(["BUILD_ID", "PROJECT_ID"])),
+        ("TeamCity", lambda: _is_env_present(["TEAMCITY_VERSION"])),
+        ("JetBrains Space", lambda: _is_env_present(["JB_SPACE_API_URL"])),
+        ("Generic CI", lambda: _is_env_true("CI")),
+    ]
+
+    return next((check[0] for check in CI_CHECKS if check[1]()), None)
+
+
+def is_ci_environment() -> bool:
+    ci_product: Optional[str] = get_ci_product_name()
+    return ci_product is not None
+
+
+def silence_errors(f: Callable) -> Callable:
+    """Decorator to silence all errors in the decorated
+    function.
+    """
+
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs) -> Any:
+        try:
+            return f(*args, **kwargs)
+        except Exception:
+            return None
+
+    return wrapper
+
+
+def _is_env_true(env_var: str) -> bool:
+    """Checks if `env_var` is `true` or `1`."""
+    return os.getenv(env_var, "").lower() in ("true", "1")
+
+
+def _is_env_present(envs: List[str]) -> bool:
+    """Checks if all of the variables passed in `envs`
+    are defined (ie: not empty)
+    """
+    return all(os.getenv(env_var, None) is not None for env_var in envs)
+
+
+def _hide_tokens(text: str) -> str:
+    """Cuts any token in a way that they get unusable if leaked,
+    but we still can use them for debugging if needed.
+    """
+    return re.sub(r"p\.ey[A-Za-z0-9-_\.]+", lambda s: f"{s[0][:10]}...{s[0][-10:]}", text)
+
+
+class TelemetryHelper:
+    def __init__(self, tb_host: Optional[str] = None, max_enqueued_events: int = 5) -> None:
+        self.tb_host = tb_host or os.getenv("TB_CLI_TELEMETRY_HOST", "https://api.tinybird.co")
+        self.max_enqueued_events: int = max_enqueued_events
+
+        self.enabled: bool = True
+        self.events: List[Dict[str, Any]] = []
+        self.telemetry_token: Optional[str] = None
+
+        run_id = str(uuid.uuid4())
+
+        self._defaults: Dict[str, Any] = {
+            # Per-event values
+            "event": "<the event>",
+            "event_data": "<the event data>",
+            "timestamp": "<the timestamp>",
+            # Static values
+            "run_id": run_id,
+        }
+
+        self._threads: List[threading.Thread] = []
+        self.log(f"Telemetry initialized with run_id: {run_id}")
+
+    @silence_errors
+    def add_event(self, event: str, event_data: Dict[str, Any]) -> None:
+        if not self.enabled:
+            self.log("Helper is disabled")
+            return
+
+        if "x.y.z" in CURRENT_VERSION and not _is_env_true("TB_CLI_TELEMETRY_SEND_IN_LOCAL"):
+            self.log("Not sending events in local development mode")
+            return
+
+        # Let's save deep copies to not interfere with original objects
+        event_dict: Dict[str, Any] = deepcopy(self._defaults)
+        event_dict["event"] = event
+        event_dict["event_data"] = json.dumps(event_data)
+        event_dict["timestamp"] = datetime.utcnow().isoformat()
+
+        self.events.append(event_dict)
+        if len(self.events) >= self.max_enqueued_events:
+            self.flush()
+
+    @silence_errors
+    def flush(self, wait: bool = False) -> None:
+        if self.enabled and len(self.events) > 0:
+            # Take the ownership for the pending events.
+            #
+            # We need this because the proper flush() is done in
+            # a thread to avoid blocking  the user and we could send
+            # the same event twice if we maintain the same list after
+            # during the sending.
+
+            events: List[Dict[str, Any]] = self.events
+            self.events = []
+
+            self.log(f"Flusing {len(events)} events in a new thread...")
+            thread: threading.Thread = threading.Thread(target=self._flush, args=[events])
+            self._threads.append(thread)
+            thread.start()
+
+        if wait:
+            for t in self._threads:
+                t.join()
+                if t.is_alive():
+                    self.log(f"Couldn't wait for the end of the thread {t.name}")
+            self._threads.clear()
+
+    @silence_errors
+    def _flush(self, events: List[Dict[str, Any]]) -> None:
+        """Actual flush. This is where we use HFI to ingest events."""
+
+        timeout: int
+        try:
+            timeout = int(os.getenv("TB_CLI_TELEMETRY_TIMEOUT", TELEMETRY_TIMEOUT))
+            timeout = max(TELEMETRY_TIMEOUT, timeout)
+        except ValueError:
+            timeout = TELEMETRY_TIMEOUT
+
+        if not self.telemetry_token:
+            self.telemetry_token = os.getenv("TB_CLI_TELEMETRY_TOKEN")
+            if self.telemetry_token:
+                self.log("Got telemetry token from environment TB_CLI_TELEMETRY_TOKEN")
+
+        with requests.Session() as session:
+            if not self.telemetry_token:
+                url: str = f"{self.tb_host}/v0/regions"
+                self.log(f"Requesting token from {url}...")
+                try:
+                    r = session.get(url, timeout=timeout)
+                    regions: List[Dict[str, Any]] = json.loads(r.content.decode())["regions"]
+                    self.telemetry_token = next(
+                        (r.get("telemetry_token", None) for r in regions if r["api_host"] == self.tb_host), None
+                    )
+                    if self.telemetry_token:
+                        self.log(f"Got telemetry token from {url}")
+                except requests.exceptions.Timeout:
+                    self.log(f"Disabling due to timeout after {timeout} seconds")
+                    self.enabled = False
+                    return
+                except Exception as ex:
+                    self.log(str(ex))
+
+            if not self.telemetry_token:
+                self.log("Disabling due to lack of token")
+                self.enabled = False
+                return
+
+            self.log(f"token={self.telemetry_token}")
+
+            data: str = _hide_tokens("\n".join(json.dumps(e) for e in events))
+
+            # Note we don't use `wait` as this telemetry isn't a critical
+            # operation to support and we don't want to generate overhead
+            params: Dict[str, Any] = {"name": TELEMETRY_DATASOURCE, "token": self.telemetry_token}
+            url = f"{self.tb_host}/v0/events?{urlencode(params)}"
+
+            try:
+                self.log(f"Sending data to {url}...")
+                r = session.post(url, data=data, timeout=timeout)
+            except requests.exceptions.Timeout:
+                self.log(f"Disabling due to timeout after {timeout} seconds")
+                self.enabled = False
+                return
+
+            self.log(f"Received status {r.status_code}: {r.text}")
+
+            if r.status_code == 200 or r.status_code == 202:
+                self.log(f"Successfully sent {len(events)} events to {self.tb_host}")
+                self.events.clear()
+                return
+
+            if r.status_code in (403, 404):
+                self.log(f"Disabling due to {r.status_code} errors")
+                self.enabled = False
+                return
+
+            if r.status_code >= 500:
+                self.log(f"Disabling telemetry and discarding {len(events)} events")
+                self.enabled = False
+
+    @silence_errors
+    def log(self, msg: str) -> None:
+        """Internal logging function to help with development and debugging."""
+        if not _is_env_true("TB_CLI_TELEMETRY_DEBUG"):
+            return
+        print(f"> Telemetry: {msg}")  # noqa: T201
+
+
+_helper_instance: Optional[TelemetryHelper] = None
+
+
+@silence_errors
+def init_telemetry() -> None:
+    """Setups the telemetry helper with the config present in `config`.
+    If no config is provided, it tries to get it from the passed Click context.
+
+    We need to call this method any time we suspect the config changes any value.
+    """
+
+    telemetry = _get_helper()
+    if telemetry:
+        telemetry.log("Initialized")
+
+
+@silence_errors
+def add_telemetry_event(event: str, **kw_event_data: Any) -> None:
+    """Adds a new telemetry event."""
+
+    telemetry = _get_helper()
+    if not telemetry:
+        return
+
+    try:
+        telemetry.add_event(event, dict(**kw_event_data))
+    except Exception as ex:
+        telemetry.log(str(ex))
+
+
+@silence_errors
+def add_telemetry_sysinfo_event() -> None:
+    """Collects system info and sends a `system_info` event
+    with the data.
+    """
+
+    ci_product: Optional[str] = get_ci_product_name()
+
+    add_telemetry_event(
+        "system_info",
+        platform=platform.platform(),
+        system=platform.system(),
+        arch=platform.machine(),
+        processor=platform.processor(),
+        python_runtime=platform.python_implementation(),
+        python_version=platform.python_version(),
+        is_ci=ci_product is not None,
+        ci_product=ci_product,
+        cli_version=CURRENT_VERSION,
+        cli_args=sys.argv[1:] if len(sys.argv) > 1 else [],
+    )
+
+
+@silence_errors
+def flush_telemetry(wait: bool = False) -> None:
+    """Flushes all pending telemetry events."""
+
+    telemetry = _get_helper()
+    if not telemetry:
+        return
+
+    try:
+        telemetry.flush(wait=wait)
+    except Exception as ex:
+        telemetry.log(str(ex))
+
+
+@silence_errors
+def _get_helper() -> Optional[TelemetryHelper]:
+    """Returns the shared TelemetryHelper instance."""
+
+    if _is_env_true("TB_CLI_TELEMETRY_OPTOUT"):
+        return None
+
+    global _helper_instance
+
+    if not _helper_instance:
+        _helper_instance = TelemetryHelper()
+
+    return _helper_instance

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

@@ -10,11 +10,11 @@ import click
 
 from tinybird.client import AuthNoTokenException
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.cli import cli
-from tinybird.tb_cli_modules.common import coro, create_tb_client, gather_with_concurrency
-from tinybird.tb_cli_modules.config import CLIConfig
-from tinybird.tb_cli_modules.exceptions import CLIException
-from tinybird.tb_cli_modules.tinyunit.tinyunit import (
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import coro, create_tb_client, gather_with_concurrency
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.exceptions import CLIException
+from tinybird.tb.modules.tinyunit.tinyunit import (
     TestSummaryResults,
     generate_file,
     parse_file,

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

@@ -9,7 +9,7 @@ from typing_extensions import override
 
 from tinybird.client import TinyB
 from tinybird.feedback_manager import FeedbackManager
-from tinybird.tb_cli_modules.common import CLIException
+from tinybird.tb.modules.common import CLIException
 
 
 @dataclass

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

@@ -8,13 +8,13 @@ from humanfriendly import parse_timespan
 
 from tinybird.client import AuthNoTokenException, TinyB
 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 (
     DoesNotExistException,
     coro,
     echo_safe_humanfriendly_tables_format_smart_table,
 )
-from tinybird.tb_cli_modules.exceptions import CLITokenException
+from tinybird.tb.modules.exceptions import CLITokenException
 
 
 @cli.group()

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

@@ -10,10 +10,9 @@ from click import Context
 
 from tinybird.client import CanNotBeDeletedException, DoesNotExistException, TinyB
 from tinybird.config import get_display_host
-from tinybird.datafile import PipeTypes
 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 (
     _get_workspace_plan_name,
     ask_for_user_token,
     check_user_token,
@@ -26,8 +25,9 @@ from tinybird.tb_cli_modules.common import (
     print_current_workspace,
     switch_workspace,
 )
-from tinybird.tb_cli_modules.config import CLIConfig
-from tinybird.tb_cli_modules.exceptions import CLIWorkspaceException
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.datafile import PipeTypes
+from tinybird.tb.modules.exceptions import CLIWorkspaceException
 
 
 @cli.group()