tinybird 0.0.1.dev305__tar.gz → 0.0.1.dev307__tar.gz

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/PKG-INFO

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

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/context.py

@@ -12,7 +12,6 @@ table_id: ContextVar[str] = ContextVar("table_id")
 hfi_frequency: ContextVar[float] = ContextVar("hfi_frequency")
 hfi_frequency_gatherer: ContextVar[float] = ContextVar("hfi_frequency_gatherer")
 use_gatherer: ContextVar[bool] = ContextVar("use_gatherer")
-allow_gatherer_fallback: ContextVar[bool] = ContextVar("allow_gatherer_fallback")
 gatherer_allow_s3_backup_on_user_errors: ContextVar[bool] = ContextVar("gatherer_allow_s3_backup_on_user_errors")
 disable_template_security_validation: ContextVar[bool] = ContextVar("disable_template_security_validation")
 origin: ContextVar[str] = ContextVar("origin")

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/prompts.py

@@ -776,6 +776,7 @@ datasource_instructions = """
     - Use always json paths to define the schema. Example: `user_id` String `json:$.user_id`,
     - Array columns are supported with a special syntax. Example: `items` Array(String) `json:$.items[:]`
     - If the datasource is using an S3 or GCS connection, they need to set IMPORT_CONNECTION_NAME, IMPORT_BUCKET_URI and IMPORT_SCHEDULE (GCS @on-demand only, S3 supports @auto too)
+    - If the datasource is using a Kafka connection, they need to set KAFKA_CONNECTION_NAME as the name of the .connection file, KAFKA_TOPIC topic_name and KAFKA_GROUP_ID as the group id for the datasource
     - Unless the user asks for them, do not include ENGINE_PARTITION_KEY and ENGINE_PRIMARY_KEY.
     - DateTime64 type without precision is not supported. Use DateTime64(3) instead.
 </datasource_file_instructions>

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/service_datasources.py

@@ -455,6 +455,79 @@ def get_tinybird_service_datasources() -> List[Dict[str, Any]]:
                 {"name": "feature", "type": "String"},
             ],
         },
+        {
+            "name": "tinybird.llm_usage",
+            "description": "LLM usage metrics from Tinybird AI features including token consumption, costs, and model usage for each request in the workspace.",
+            "dateColumn": "start_time",
+            "engine": {
+                "engine": "MergeTree",
+                "sorting_key": "workspace_id, start_time, user_email, request_id",
+                "partition_key": "toYYYYMM(start_time)",
+            },
+            "columns": [
+                {"name": "start_time", "type": "DateTime"},
+                {"name": "end_time", "type": "DateTime"},
+                {"name": "organization_id", "type": "String"},
+                {"name": "organization_name", "type": "String"},
+                {"name": "workspace_id", "type": "String"},
+                {"name": "workspace_name", "type": "String"},
+                {"name": "user_email", "type": "String"},
+                {"name": "request_id", "type": "String"},
+                {"name": "prompt_tokens", "type": "UInt32"},
+                {"name": "completion_tokens", "type": "UInt32"},
+                {"name": "total_tokens", "type": "UInt32"},
+                {"name": "duration", "type": "Float32"},
+                {"name": "cost", "type": "Float32"},
+                {"name": "origin", "type": "String"},
+                {"name": "feature", "type": "String"},
+            ],
+        },
+        {
+            "name": "tinybird.query_metrics",
+            "description": "Query stats metrics from your workspace.",
+            "dateColumn": "event_time",
+            "engine": {
+                "engine": "ReplacingMergeTree",
+                "sorting_key": "event_time, organization_id, query_id",
+                "partition_key": "toStartOfDay(event_time)",
+            },
+            "columns": [
+                {"name": "event_time", "type": "DateTime"},
+                {"name": "organization_id", "type": "String"},
+                {"name": "workspace_id", "type": "String"},
+                {"name": "query", "type": "String"},
+                {"name": "query_id", "type": "String"},
+                {"name": "query_type", "type": "String"},
+                {"name": "query_start_time", "type": "DateTime"},
+                {"name": "query_duration_ms", "type": "Int32"},
+                {"name": "pipe_id", "type": "String"},
+                {"name": "job_id", "type": "String"},
+                {"name": "job_kind", "type": "String"},
+                {"name": "read_rows", "type": "Int32"},
+                {"name": "read_bytes", "type": "Int32"},
+                {"name": "written_rows", "type": "Int32"},
+                {"name": "written_bytes", "type": "Int32"},
+                {"name": "memory_usage", "type": "Int32"},
+                {"name": "vcpu_time", "type": "Float32"},
+                {"name": "exception_code", "type": "Int32"},
+                {"name": "exception", "type": "String"},
+            ],
+        },
+        {
+            "name": "tinybird.vcpu_time",
+            "description": "vCPU time metrics from your workspace.",
+            "dateColumn": "minute_slot",
+            "engine": {
+                "engine": "AggregatingMergeTree",
+                "sorting_key": "organization_id, minute_slot",
+                "partition_key": "toStartOfDay(minute_slot)",
+            },
+            "columns": [
+                {"name": "minute_slot", "type": "DateTime"},
+                {"name": "organization_id", "type": "String"},
+                {"name": "vcpu_time", "type": "Float64"},
+            ],
+        },
     ]
 
 
@@ -935,6 +1008,53 @@ def get_organization_service_datasources() -> List[Dict[str, Any]]:
                 {"name": "feature", "type": "String"},
             ],
         },
+        {
+            "name": "organization.query_metrics",
+            "description": "Query stats metrics from your workspace.",
+            "dateColumn": "event_time",
+            "engine": {
+                "engine": "ReplacingMergeTree",
+                "sorting_key": "event_time, organization_id, query_id",
+                "partition_key": "toStartOfDay(event_time)",
+            },
+            "columns": [
+                {"name": "event_time", "type": "DateTime"},
+                {"name": "organization_id", "type": "String"},
+                {"name": "workspace_id", "type": "String"},
+                {"name": "query", "type": "String"},
+                {"name": "query_id", "type": "String"},
+                {"name": "query_type", "type": "String"},
+                {"name": "query_start_time", "type": "DateTime"},
+                {"name": "query_duration_ms", "type": "Int32"},
+                {"name": "pipe_id", "type": "String"},
+                {"name": "job_id", "type": "String"},
+                {"name": "job_kind", "type": "String"},
+                {"name": "read_rows", "type": "Int32"},
+                {"name": "read_bytes", "type": "Int32"},
+                {"name": "written_rows", "type": "Int32"},
+                {"name": "written_bytes", "type": "Int32"},
+                {"name": "memory_usage", "type": "Int32"},
+                {"name": "vcpu_time", "type": "Float32"},
+                {"name": "exception_code", "type": "Int32"},
+                {"name": "exception", "type": "String"},
+            ],
+        },
+        {
+            "name": "organization.vcpu_time",
+            "description": "vCPU time metrics from your workspace.",
+            "dateColumn": "minute_slot",
+            "engine": {
+                "engine": "AggregatingMergeTree",
+                "sorting_key": "organization_id, minute_slot",
+                "partition_key": "toStartOfDay(minute_slot)",
+            },
+            "columns": [
+                {"name": "minute_slot", "type": "DateTime"},
+                {"name": "organization_id", "type": "String"},
+                {"name": "workspace_id", "type": "String"},
+                {"name": "vcpu_time", "type": "Float64"},
+            ],
+        },
     ]
 
 

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/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.dev305'
-__revision__ = 'cc453b3'
+__version__ = '0.0.1.dev307'
+__revision__ = '060369b'

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/tb/cli.py

@@ -9,6 +9,7 @@ import tinybird.tb.modules.datasource
 import tinybird.tb.modules.deployment
 import tinybird.tb.modules.deprecations
 import tinybird.tb.modules.endpoint
+import tinybird.tb.modules.environment
 import tinybird.tb.modules.info
 import tinybird.tb.modules.infra
 import tinybird.tb.modules.job

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/tb/client.py

@@ -711,7 +711,7 @@ class TinyB:
         return self._req(f"/{version}/user/workspaces/?with_environments=true&only_environments=true")
 
     def branches(self):
-        return self._req("/v0/environments")
+        return self._req("/v1/environments")
 
     def releases(self, workspace_id):
         return self._req(f"/v0/workspaces/{workspace_id}/releases")
@@ -740,7 +740,7 @@ class TinyB:
         }
         if ignore_datasources:
             params["ignore_datasources"] = ",".join(ignore_datasources)
-        return self._req(f"/v0/environments?{urlencode(params)}", method="POST", data=b"")
+        return self._req(f"/v1/environments?{urlencode(params)}", method="POST", data=b"")
 
     def branch_workspace_data(
         self,

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/tb/modules/build.py

@@ -35,7 +35,9 @@ 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":
+
+    # TODO: Explain that you can use custom environments too once they are open for everyone
+    if obj["env"] == "cloud" and not obj["environment"]:
         raise click.ClickException(FeedbackManager.error_build_only_supported_in_local())
 
     if project.has_deeper_level():
@@ -63,11 +65,18 @@ def build(ctx: click.Context, watch: bool) -> None:
 @click.option("--ui", is_flag=True, default=False, 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)
+
+    # TODO: Explain that you can use custom environments too once they are open for everyone
+    if obj["env"] == "cloud" and not obj["environment"]:
+        raise click.ClickException(FeedbackManager.error_build_only_supported_in_local())
+
     if data_origin == "cloud":
         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(

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/tb/modules/cli.py

@@ -76,6 +76,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("--environment", help="Run against a custom environment.")
 @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 +104,7 @@ def cli(
     version_warning: bool,
     show_tokens: bool,
     cloud: bool,
+    environment: Optional[str],
     staging: bool,
     output: str,
     max_depth: int,
@@ -110,7 +112,7 @@ def cli(
     prompt: Optional[str] = None,
 ) -> None:
     """
-    Use Tinybird Code to interact with your project.
+    Run just `tb` to use Tinybird Code to interact with your project.
     """
 
     # We need to unpatch for our tests not to break
@@ -200,17 +202,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(environment),
+        staging,
+        project=project,
+        show_warnings=version_warning,
+        environment=environment,
+    )
 
     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, environment)
+    ctx.ensure_object(dict)["environment"] = environment
     ctx.ensure_object(dict)["output"] = output
 
     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"]):
+            raise CLIException(
+                FeedbackManager.error(
+                    message="Tinybird Code does not support --cloud or --local flags. It will choose the correct environment based on your prompts."
+                )
+            )
+
         run_agent(config, project, dangerously_skip_permissions, prompt=prompt)
 
 
@@ -477,7 +495,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,
+    environment: Optional[str] = None,
 ):
     commands_without_ctx_client = [
         "auth",
@@ -499,8 +523,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", "environment"]
+    commands_always_local = ["dev", "create"]
     command_always_test = ["test"]
 
     if (
@@ -523,7 +547,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, environment=environment)
     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:
@@ -531,8 +555,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, environment: Optional[str]) -> str:
+    if cloud or bool(environment):
         return "cloud"
     return "local"
 

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/tb/modules/common.py

@@ -378,9 +378,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, environment: 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 environment:
+        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") == environment), None)
+    if not workspace:
+        raise CLIException(FeedbackManager.error_exception(error=f"Environment {environment} 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:
@@ -1165,10 +1189,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(
             (
@@ -2261,3 +2301,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)

{tinybird-0.0.1.dev305 → tinybird-0.0.1.dev307}/tinybird/tb/modules/deployment_common.py

@@ -59,9 +59,14 @@ def api_post(
     params: Optional[dict] = None,
 ) -> dict:
     r = requests.post(url, headers=headers, files=files, params=params)
-    if r.status_code < 300:
-        logging.debug(json.dumps(r.json(), indent=2))
-        return r.json()
+    try:
+        if r.status_code < 300:
+            logging.debug(json.dumps(r.json(), indent=2))
+            return r.json()
+    except json.JSONDecodeError:
+        message = "Error parsing response from API"
+        click.echo(FeedbackManager.error(message=message))
+        sys_exit("deployment_error", message)
 
     # Try to parse and print the error from the response
     try:

tinybird-0.0.1.dev307/tinybird/tb/modules/environment.py

@@ -0,0 +1,152 @@
+# This is a command file for our CLI. Please keep it clean.
+#
+# - If it makes sense and only when strictly necessary, you can create utility functions in this file.
+# - But please, **do not** interleave utility functions and command definitions.
+
+from typing import List, Optional, Tuple
+
+import click
+
+from tinybird.tb.modules.cli import cli
+from tinybird.tb.modules.common import (
+    MAIN_BRANCH,
+    create_workspace_branch,
+    echo_safe_humanfriendly_tables_format_smart_table,
+    get_current_main_workspace,
+    get_current_workspace_branches,
+    get_workspace_member_email,
+    switch_to_workspace_by_user_workspace_data,
+    try_update_config_with_remote,
+)
+from tinybird.tb.modules.config import CLIConfig
+from tinybird.tb.modules.exceptions import CLIBranchException, CLIException
+from tinybird.tb.modules.feedback_manager import FeedbackManager
+
+
+@cli.group(hidden=True)
+def environment() -> None:
+    """Environment commands. Custom environments is an experimental feature in beta."""
+    pass
+
+
+@environment.command(name="ls")
+@click.option("--sort/--no-sort", default=False, help="Sort the table rows by name")
+def branch_ls(sort: bool) -> None:
+    """List all the environments available in the current workspace"""
+
+    config = CLIConfig.get_project_config()
+    _ = try_update_config_with_remote(config, only_if_needed=True)
+
+    client = config.get_client()
+
+    current_main_workspace = get_current_main_workspace(config)
+    assert isinstance(current_main_workspace, dict)
+
+    if current_main_workspace["id"] != config["id"]:
+        client = config.get_client(token=current_main_workspace["token"])
+
+    response = client.branches()
+
+    columns = ["name", "id", "created_at", "owner", "current"]
+
+    table: List[Tuple[str, str, str, str, bool]] = []
+
+    for branch in response["environments"]:
+        branch_owner_email = get_workspace_member_email(branch, branch["owner"])
+
+        table.append(
+            (branch["name"], branch["id"], branch["created_at"], branch_owner_email, config["id"] == branch["id"])
+        )
+
+    current_branch = [row for row in table if row[4]]
+    other_branches = [row for row in table if not row[4]]
+
+    if sort:
+        other_branches.sort(key=lambda x: x[0])
+
+    sorted_table = current_branch + other_branches
+
+    click.echo(FeedbackManager.info(message="\n** Environments:"))
+    echo_safe_humanfriendly_tables_format_smart_table(sorted_table, column_names=columns)
+
+
+@environment.command(name="create", short_help="Create a new environment in the current Workspace")
+@click.argument("environment_name", required=False)
+@click.option(
+    "--last-partition",
+    is_flag=True,
+    default=False,
+    help="Attach the last modified partition from the current workspace to the new environment",
+)
+@click.option(
+    "-i",
+    "--ignore-datasource",
+    "ignore_datasources",
+    type=str,
+    multiple=True,
+    help="Ignore specified data source partitions",
+)
+@click.option(
+    "--wait/--no-wait",
+    is_flag=True,
+    default=True,
+    help="Wait for data branch jobs to finish, showing a progress bar. Disabled by default.",
+)
+def create_environment(
+    environment_name: Optional[str], last_partition: bool, ignore_datasources: List[str], wait: bool
+) -> None:
+    create_workspace_branch(environment_name, last_partition, False, list(ignore_datasources), wait)
+
+
+@environment.command(name="rm", short_help="Removes an environment from the workspace. It can't be recovered.")
+@click.argument("branch_name_or_id")
+@click.option("--yes", is_flag=True, default=False, help="Do not ask for confirmation")
+def delete_environment(branch_name_or_id: str, yes: bool) -> None:
+    """Remove an environment"""
+
+    config = CLIConfig.get_project_config()
+    _ = try_update_config_with_remote(config)
+
+    client = config.get_client()
+
+    if branch_name_or_id == MAIN_BRANCH:
+        raise CLIException(FeedbackManager.error_not_allowed_in_main_branch())
+
+    try:
+        workspace_branches = get_current_workspace_branches(config)
+        workspace_to_delete = next(
+            (
+                workspace
+                for workspace in workspace_branches
+                if workspace["name"] == branch_name_or_id or workspace["id"] == branch_name_or_id
+            ),
+            None,
+        )
+    except Exception as e:
+        raise CLIBranchException(FeedbackManager.error_exception(error=str(e)))
+
+    if not workspace_to_delete:
+        raise CLIBranchException(FeedbackManager.error_branch(branch=branch_name_or_id))
+
+    if yes or click.confirm(FeedbackManager.warning_confirm_delete_branch(branch=workspace_to_delete["name"])):
+        need_to_switch_to_main = workspace_to_delete.get("main") and config["id"] == workspace_to_delete["id"]
+        # get origin workspace if deleting current branch
+        if need_to_switch_to_main:
+            try:
+                workspaces = (client.user_workspaces()).get("workspaces", [])
+                workspace_main = next(
+                    (workspace for workspace in workspaces if workspace["id"] == workspace_to_delete["main"]), None
+                )
+            except Exception:
+                workspace_main = None
+        try:
+            client.delete_branch(workspace_to_delete["id"])
+            click.echo(FeedbackManager.success_branch_deleted(branch_name=workspace_to_delete["name"]))
+        except Exception as e:
+            raise CLIBranchException(FeedbackManager.error_exception(error=str(e)))
+        else:
+            if need_to_switch_to_main:
+                if workspace_main:
+                    switch_to_workspace_by_user_workspace_data(config, workspace_main)
+                else:
+                    raise CLIException(FeedbackManager.error_switching_to_main())