tinybird 0.0.1.dev102__py3-none-any.whl → 0.0.1.dev104__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of tinybird might be problematic. Click here for more details.

@@ -1,30 +1,30 @@
1
+ import asyncio
1
2
  import json
2
3
  import subprocess
4
+ import sys
3
5
  import time
4
6
  import uuid
5
7
  from pathlib import Path
6
- from typing import Optional
8
+ from typing import Any, Dict, List, Optional
7
9
 
8
10
  import click
9
- import pyperclip
10
11
  import requests
11
12
  from click import Context
12
13
 
13
14
  from tinybird.client import TinyB
14
- from tinybird.syncasync import async_to_sync
15
15
  from tinybird.tb.modules.cli import CLIException, cli
16
16
  from tinybird.tb.modules.common import coro, echo_safe_humanfriendly_tables_format_smart_table
17
17
  from tinybird.tb.modules.config import CLIConfig
18
18
  from tinybird.tb.modules.feedback_manager import FeedbackManager
19
19
 
20
- from .common import CONTEXT_SETTINGS
20
+ from .common import ask_for_organization, get_organizations_by_user
21
21
 
22
22
  K8S_YML = """
23
23
  ---
24
24
  apiVersion: v1
25
25
  kind: Namespace
26
26
  metadata:
27
- name: %(namespace)s
27
+ name: %(kubernetes_namespace)s
28
28
  labels:
29
29
  name: tinybird
30
30
  ---
@@ -32,7 +32,7 @@ apiVersion: v1
32
32
  kind: ServiceAccount
33
33
  metadata:
34
34
  name: tinybird
35
- namespace: %(namespace)s
35
+ namespace: %(kubernetes_namespace)s
36
36
  labels:
37
37
  name: tinybird
38
38
  automountServiceAccountToken: true
@@ -41,7 +41,7 @@ apiVersion: v1
41
41
  kind: Service
42
42
  metadata:
43
43
  name: tinybird
44
- namespace: %(namespace)s
44
+ namespace: %(kubernetes_namespace)s
45
45
  labels:
46
46
  name: tinybird
47
47
  spec:
@@ -57,15 +57,16 @@ spec:
57
57
  apiVersion: networking.k8s.io/v1
58
58
  kind: Ingress
59
59
  metadata:
60
- namespace: %(namespace)s
60
+ namespace: %(kubernetes_namespace)s
61
61
  name: tinybird
62
62
  annotations:
63
+ external-dns.alpha.kubernetes.io/aws-evaluate-target-health: "false"
63
64
  alb.ingress.kubernetes.io/listen-ports: '[{"HTTPS":443}]'
64
65
  alb.ingress.kubernetes.io/scheme: internet-facing
65
66
  alb.ingress.kubernetes.io/ssl-policy: ELBSecurityPolicy-TLS13-1-2-2021-06
66
67
  alb.ingress.kubernetes.io/ssl-redirect: '443'
67
68
  alb.ingress.kubernetes.io/target-type: 'ip'
68
- alb.ingress.kubernetes.io/load-balancer-name: %(namespace)s
69
+ alb.ingress.kubernetes.io/load-balancer-name: %(kubernetes_namespace)s
69
70
  alb.ingress.kubernetes.io/success-codes: '200,301,302'
70
71
  spec:
71
72
  ingressClassName: alb
@@ -87,7 +88,7 @@ apiVersion: apps/v1
87
88
  kind: StatefulSet
88
89
  metadata:
89
90
  name: tinybird
90
- namespace: %(namespace)s
91
+ namespace: %(kubernetes_namespace)s
91
92
  spec:
92
93
  serviceName: "tinybird"
93
94
  replicas: 1
@@ -120,6 +121,8 @@ spec:
120
121
  volumeMounts:
121
122
  - name: clickhouse-data
122
123
  mountPath: /var/lib/clickhouse
124
+ - name: redis-data
125
+ mountPath: /redis-data
123
126
  volumeClaimTemplates:
124
127
  - metadata:
125
128
  name: clickhouse-data
@@ -127,11 +130,17 @@ spec:
127
130
  accessModes: [ "ReadWriteOnce" ]
128
131
  resources:
129
132
  requests:
130
- storage: 100Gi
131
- storageClassName: %(storage_class)s
133
+ storage: 100Gi%(storage_class_line)s
134
+ - metadata:
135
+ name: redis-data
136
+ spec:
137
+ accessModes: [ "ReadWriteOnce" ]
138
+ resources:
139
+ requests:
140
+ storage: 50Gi%(storage_class_line)s
132
141
  """
133
142
 
134
- TERRAFORM_FIRST_TEMPLATE = """
143
+ TERRAFORM_TEMPLATE = """
135
144
  terraform {
136
145
  required_providers {
137
146
  aws = {
@@ -189,7 +198,7 @@ resource "aws_acm_certificate_validation" "cert" {
189
198
  TERRAFORM_SECOND_TEMPLATE = """
190
199
  # Create Route 53 record for the load balancer
191
200
  data "aws_alb" "tinybird" {
192
- name = "%(namespace)s"
201
+ name = "%(kubernetes_namespace)s"
193
202
  }
194
203
 
195
204
  resource "aws_route53_record" "tinybird" {
@@ -210,50 +219,158 @@ output "tinybird_dns" {
210
219
  """
211
220
 
212
221
 
213
- @cli.group(context_settings=CONTEXT_SETTINGS, hidden=True)
222
+ class Infrastructure:
223
+ """Class to manage Tinybird infrastructure operations."""
224
+
225
+ def __init__(self, client: TinyB):
226
+ self.client = client
227
+
228
+ async def get_organizations_info(self, config: CLIConfig):
229
+ self.orgs = await get_organizations_by_user(config)
230
+
231
+ async def create_infra(self, name: str, host: str, organization_id: str) -> Dict[str, Any]:
232
+ """Create a new infrastructure."""
233
+ infra = await self.client.infra_create(organization_id=organization_id, name=name, host=host)
234
+ return infra
235
+
236
+ async def list_infras(self) -> None:
237
+ """List all self-managed regions for the admin organization."""
238
+
239
+ infras = await self.get_infra_list()
240
+ columns = [
241
+ "name",
242
+ "host",
243
+ "organization",
244
+ ]
245
+ table_human_readable = []
246
+ table_machine_readable = []
247
+
248
+ for infra in infras:
249
+ name = infra["name"]
250
+ host = infra["host"]
251
+ organization = infra["organization"]
252
+
253
+ table_human_readable.append((name, host, organization))
254
+ table_machine_readable.append({"name": name, "host": host, "organization": organization})
255
+
256
+ click.echo(FeedbackManager.info(message="\n** Infras:"))
257
+ echo_safe_humanfriendly_tables_format_smart_table(table_human_readable, column_names=columns)
258
+ click.echo("\n")
259
+
260
+ async def remove_infra(self, name: str):
261
+ try:
262
+ click.echo(FeedbackManager.highlight(message=f"» Deleting infrastructure '{name}' from Tinybird..."))
263
+ infras = await self.get_infra_list()
264
+ infra = next((infra for infra in infras if infra["name"] == name), None)
265
+ if not infra:
266
+ raise CLIException(f"Infrastructure '{name}' not found")
267
+ await self.client.infra_delete(infra["id"], infra["organization_id"])
268
+ click.echo(FeedbackManager.success(message=f"\n✓ Infrastructure '{name}' deleted"))
269
+ except Exception as e:
270
+ click.echo(FeedbackManager.error(message=f"✗ Error: {e}"))
271
+
272
+ async def update(self, infra_name: str, name: str, host: str):
273
+ infras_list = await self.get_infra_list()
274
+ infra = next((infra for infra in infras_list if infra["name"] == infra_name), None)
275
+ if not infra:
276
+ return None
277
+ await self.client.infra_update(
278
+ infra_id=infra["id"],
279
+ organization_id=infra["organization_id"],
280
+ name=name,
281
+ host=host,
282
+ )
283
+ return infra
284
+
285
+ async def get_infra_list(self) -> List[Dict[str, str]]:
286
+ """Get a list of all infrastructures across organizations.
287
+
288
+ Returns:
289
+ List[Dict[str, Any]]: List of infrastructure objects, each containing:
290
+ - id: str
291
+ - name: str
292
+ - host: str
293
+ - organization: str
294
+ """
295
+ try:
296
+ all_infras = []
297
+ for org in self.orgs:
298
+ org_id = org.get("id") or ""
299
+ org_name = org.get("name")
300
+ try:
301
+ infras = await self.client.infra_list(organization_id=org_id)
302
+ for infra in infras:
303
+ infra["organization"] = org_name
304
+ all_infras.append(infra)
305
+ except Exception as e:
306
+ click.echo(
307
+ FeedbackManager.warning(message=f"Could not fetch infras for organization {org_name}: {str(e)}")
308
+ )
309
+ continue
310
+ return all_infras
311
+ except Exception as e:
312
+ click.echo(FeedbackManager.error(message=f"Failed to list infrastructures: {str(e)}"))
313
+ return []
314
+
315
+
316
+ @cli.group(hidden=True)
214
317
  @click.pass_context
215
318
  def infra(ctx: Context) -> None:
216
319
  """Infra commands."""
320
+ if "--help" in sys.argv or "-h" in sys.argv:
321
+ return
322
+ client: TinyB = ctx.ensure_object(dict)["client"]
323
+ config = CLIConfig.get_project_config()
324
+ infra = Infrastructure(client)
325
+ asyncio.run(infra.get_organizations_info(config))
326
+ ctx.ensure_object(dict)["infra"] = infra
217
327
 
218
328
 
219
329
  @infra.command(name="init")
220
- @click.option("--name", type=str, help="Name for identifying the self-managed infrastructure in Tinybird")
221
- @click.option("--provider", type=str, help="Infrastructure provider. Possible values are: aws, gcp, azure)")
222
- @click.option("--region", type=str, help="AWS region, when using aws as the provider")
330
+ @click.option("--name", type=str, help="Name for identifying the self-managed region in Tinybird")
331
+ @click.option("--cloud-provider", type=str, help="Infrastructure provider. Possible values are: aws, gcp, azure)")
332
+ @click.option("--cloud-region", type=str, help="AWS region, when using aws as the provider")
223
333
  @click.option("--dns-zone-name", type=str, help="DNS zone name")
224
- @click.option("--namespace", type=str, help="Kubernetes namespace for the deployment")
225
- @click.option("--dns-record", type=str, help="DNS record name to create, without domain. For example, 'tinybird')")
226
- @click.option("--storage-class", type=str, help="Storage class for the k8s StatefulSet")
334
+ @click.option("--dns-record", type=str, help="DNS record name to create, without domain. For example, 'tinybird'")
335
+ @click.option("--kubernetes-namespace", type=str, help="Kubernetes namespace for the deployment")
336
+ @click.option("--kubernetes-storage-class", type=str, help="Storage class for the k8s StatefulSet")
227
337
  @click.option(
228
338
  "--auto-apply", is_flag=True, help="Automatically apply Terraform and kubectl configuration without prompting"
229
339
  )
230
340
  @click.option("--skip-apply", is_flag=True, help="Skip Terraform and kubectl configuration and application")
341
+ @click.option("--organization-id", type=str, help="Organization ID for the self-managed region")
231
342
  @click.pass_context
232
- def infra_init(
343
+ @coro
344
+ async def infra_init(
233
345
  ctx: Context,
234
346
  name: str,
235
- provider: str,
236
- region: Optional[str] = None,
347
+ cloud_provider: str,
348
+ cloud_region: Optional[str] = None,
237
349
  dns_zone_name: Optional[str] = None,
238
- namespace: Optional[str] = None,
350
+ kubernetes_namespace: Optional[str] = None,
239
351
  dns_record: Optional[str] = None,
240
- storage_class: Optional[str] = None,
352
+ kubernetes_storage_class: Optional[str] = None,
241
353
  auto_apply: bool = False,
242
354
  skip_apply: bool = False,
355
+ organization_id: Optional[str] = None,
243
356
  ) -> None:
244
357
  """Init infra"""
245
358
  # Check if provider is specified
246
- if not provider:
247
- click.echo("Error: --provider option is required. Specify a provider. Possible values are: aws, gcp, azure.")
359
+ if not cloud_provider:
360
+ click.echo(
361
+ FeedbackManager.error(
362
+ message="✗ Error: --cloud-provider option is required. Specify a cloud provider. Possible values are: aws, gcp, azure."
363
+ )
364
+ )
248
365
  return
249
366
 
250
367
  # AWS-specific Terraform template creation
251
- if provider.lower() != "aws":
252
- click.echo("Provider not supported yet.")
368
+ if cloud_provider.lower() != "aws":
369
+ click.echo(FeedbackManager.error(message="✗ Error: Provider not supported yet."))
253
370
  return
254
371
 
255
372
  # Create infra directory if it doesn't exist
256
- infra_dir = Path(f"infra/{provider}")
373
+ infra_dir = Path(f"infra/{cloud_provider}")
257
374
  infra_dir.mkdir(exist_ok=True)
258
375
  yaml_path = infra_dir / "k8s.yaml"
259
376
  tf_path = infra_dir / "main.tf"
@@ -263,145 +380,193 @@ def infra_init(
263
380
  config = {}
264
381
  if config_path.exists():
265
382
  try:
266
- with open(config_path, "r") as f:
383
+ with open(config_path, "r") as f: # noqa: ASYNC230
267
384
  config = json.load(f)
268
- click.echo("Loaded existing configuration from config.json")
385
+ click.echo(FeedbackManager.info(message="** Loaded existing configuration from config.json"))
269
386
  except json.JSONDecodeError:
270
- click.echo("Warning: Could not parse existing config.json. Creating a new file...")
387
+ click.echo(
388
+ FeedbackManager.warning(
389
+ message="** Warning: Could not parse existing config.json. Creating a new file..."
390
+ )
391
+ )
271
392
 
272
393
  # Generate a random ID for default values
273
394
  random_id = str(uuid.uuid4())[:8]
274
395
 
275
396
  # Get or prompt for configuration values
276
- name = name or click.prompt("Enter the name for your self-managed region", type=str)
277
- region = region or config.get("region") or click.prompt("Enter the AWS region", default="us-east-1", type=str)
278
- dns_zone_name = dns_zone_name or config.get("dns_zone_name") or click.prompt("Enter the DNS zone name", type=str)
279
- namespace = (
280
- namespace
281
- or config.get("namespace")
282
- or click.prompt("Enter the Kubernetes namespace", default=f"tinybird-{random_id}", type=str)
397
+ name = name or config.get("name") or click.prompt("** Enter the name for your self-managed region", type=str)
398
+ cloud_region = (
399
+ cloud_region
400
+ or config.get("cloud_region")
401
+ or click.prompt("** Enter the AWS region", default="us-east-1", type=str)
283
402
  )
403
+ dns_zone_name = dns_zone_name or config.get("dns_zone_name") or click.prompt("** Enter the DNS zone name", type=str)
284
404
  dns_record = (
285
405
  dns_record
286
406
  or config.get("dns_record")
287
- or click.prompt("Enter the DNS record name, without domain", default=f"tinybird-{random_id}", type=str)
407
+ or click.prompt("** Enter the DNS record name, without domain", default=f"tinybird-{random_id}", type=str)
288
408
  )
289
- storage_class = config.get("storage_class") or click.prompt(
290
- "Enter the Kubernetes storage class", default="gp3-encrypted", type=str
409
+ kubernetes_namespace = (
410
+ kubernetes_namespace
411
+ or config.get("kubernetes_namespace")
412
+ or click.prompt("** Enter the Kubernetes namespace", default=f"tinybird-{random_id}", type=str)
291
413
  )
292
414
 
415
+ # Special handling for kubernetes_storage_class to handle None values
416
+ if kubernetes_storage_class is not None:
417
+ # Use the provided value
418
+ pass
419
+ elif "kubernetes_storage_class" in config:
420
+ # Use the value from config, even if it's None
421
+ kubernetes_storage_class = config["kubernetes_storage_class"]
422
+ else:
423
+ # Prompt the user
424
+ kubernetes_storage_class = click.prompt(
425
+ "** Enter the Kubernetes storage class (leave empty for None)", default="", show_default=False, type=str
426
+ )
427
+ # Convert empty string to None
428
+ if kubernetes_storage_class == "":
429
+ kubernetes_storage_class = None
430
+
431
+ kubernetes_context = config.get("kubernetes_context")
432
+
293
433
  # Save configuration
294
434
  config = {
295
- "provider": provider,
296
- "region": region,
435
+ "name": name,
436
+ "cloud_provider": cloud_provider,
437
+ "cloud_region": cloud_region,
297
438
  "dns_zone_name": dns_zone_name,
298
- "namespace": namespace,
299
439
  "dns_record": dns_record,
300
- "storage_class": storage_class,
440
+ "kubernetes_namespace": kubernetes_namespace,
441
+ "kubernetes_storage_class": kubernetes_storage_class,
442
+ "kubernetes_context": kubernetes_context,
301
443
  }
302
444
 
303
- with open(config_path, "w") as f:
445
+ with open(config_path, "w") as f: # noqa: ASYNC230
304
446
  json.dump(config, f, indent=2)
305
447
 
306
- click.echo(f"Configuration saved to {config_path}")
448
+ click.echo(FeedbackManager.info(message=f"** Configuration saved to {config_path}"))
307
449
 
308
- client: TinyB = ctx.obj["client"]
450
+ infra: Infrastructure = ctx.ensure_object(dict)["infra"]
309
451
  cli_config = CLIConfig.get_project_config()
310
- user_client = cli_config.get_client(token=cli_config.get_user_token() or "")
311
- user_workspaces = async_to_sync(user_client.user_workspaces_with_organization)()
312
- admin_org_id = user_workspaces.get("organization_id")
313
- admin_org_name = user_workspaces.get("organization_name", "")
314
- infras = async_to_sync(client.infra_list)(organization_id=admin_org_id)
315
- infra = next((infra for infra in infras if infra["name"] == name), None)
316
- if not infra:
452
+
453
+ infras = await infra.get_infra_list()
454
+ infra_name = next((infra for infra in infras if infra["name"] == name), None)
455
+ if not infra_name:
456
+ # Handle organization if not provided
457
+ organization_id, organization_name = await ask_for_organization(infra.orgs, organization_id)
458
+ if not organization_id:
459
+ return
317
460
  click.echo(FeedbackManager.highlight(message=f"\n» Creating infrastructure '{name}' in Tinybird..."))
318
461
  host = f"https://{dns_record}.{dns_zone_name}"
319
- infra = async_to_sync(client.infra_create)(organization_id=admin_org_id, name=name, host=host)
320
-
321
- infra_token = infra["token"]
462
+ infra_obj = await infra.create_infra(name, host, organization_id)
463
+ else:
464
+ click.echo(FeedbackManager.highlight(message=f"» Infrastructure '{name}' already exists."))
465
+ if infra_name["host"] != f"https://{dns_record}.{dns_zone_name}":
466
+ click.echo(
467
+ FeedbackManager.highlight(
468
+ message="» Infrastructure host is different in the current config than the one provisioned at Tinybird"
469
+ )
470
+ )
471
+ if click.confirm("Would you like to update the host in the infra provisioned at Tinybird?"):
472
+ await infra.update(infra_name=name, name=name, host=f"https://{dns_record}.{dns_zone_name}")
473
+ organization_name = infra_name["organization"]
474
+ infra_obj = infra_name
322
475
 
323
476
  # Write the Terraform template
324
- terraform_content = TERRAFORM_FIRST_TEMPLATE % {
325
- "aws_region": region,
477
+ terraform_content = TERRAFORM_TEMPLATE % {
478
+ "aws_region": cloud_region,
326
479
  "dns_zone_name": dns_zone_name,
327
480
  "dns_record": dns_record,
328
481
  }
329
482
 
330
- with open(tf_path, "w") as f:
483
+ with open(tf_path, "w") as f: # noqa: ASYNC230
331
484
  f.write(terraform_content.lstrip())
332
485
 
333
- click.echo(f"Created Terraform configuration in {tf_path}")
486
+ click.echo(FeedbackManager.info(message=f"** Created Terraform configuration in {tf_path}"))
487
+
488
+ # Prepare the storage class line based on whether kubernetes_storage_class is None
489
+ storage_class_line = f"\n storageClassName: {kubernetes_storage_class}" if kubernetes_storage_class else ""
334
490
 
335
491
  new_content = K8S_YML % {
336
- "namespace": namespace,
337
- "storage_class": storage_class,
492
+ "kubernetes_namespace": kubernetes_namespace,
493
+ "storage_class_line": storage_class_line,
338
494
  "full_dns_name": f"{dns_record}.{dns_zone_name}",
339
- "infra_token": infra_token,
495
+ "infra_token": infra_obj["token"],
340
496
  "infra_workspace": cli_config.get("name", ""),
341
- "infra_organization": admin_org_name,
497
+ "infra_organization": organization_name,
342
498
  "infra_user": cli_config.get_user_email() or "",
343
499
  }
344
500
 
345
- with open(yaml_path, "w") as f:
501
+ with open(yaml_path, "w") as f: # noqa: ASYNC230
346
502
  f.write(new_content.lstrip())
347
503
 
348
- click.echo(f"Created Kubernetes configuration in {yaml_path}")
504
+ click.echo(FeedbackManager.info(message=f"** Created Kubernetes configuration in {yaml_path}"))
349
505
 
350
506
  # Apply Terraform configuration if user confirms
351
507
  if not skip_apply:
352
508
  # Initialize Terraform
353
- click.echo("Initializing Terraform...")
354
- init_result = subprocess.run(["terraform", f"-chdir={infra_dir}", "init"], capture_output=True, text=True)
509
+ click.echo(FeedbackManager.info(message="** Initializing Terraform..."))
510
+ command = f"terraform -chdir={infra_dir} init"
511
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
512
+
513
+ init_result = subprocess.run(["terraform", f"-chdir={infra_dir}", "init"], capture_output=True, text=True) # noqa: ASYNC221
355
514
 
356
515
  if init_result.returncode != 0:
357
- click.echo("Terraform initialization failed:")
516
+ click.echo(FeedbackManager.error(message="✗ Error: Terraform initialization failed:"))
358
517
  click.echo(init_result.stderr)
359
518
  return
360
519
 
361
- # Run terraform plan first
362
- click.echo("\nRunning Terraform plan...\n")
363
- plan_result = subprocess.run(["terraform", f"-chdir={infra_dir}", "plan"], capture_output=True, text=True)
520
+ # Run terraform plan
521
+ click.echo(FeedbackManager.info(message="** Running Terraform plan..."))
522
+ command = f"terraform -chdir={infra_dir} plan"
523
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
524
+ plan_result = subprocess.run(command.split(), capture_output=True, text=True) # noqa: ASYNC221
364
525
 
365
526
  if plan_result.returncode != 0:
366
- click.echo("Terraform plan failed:")
367
- click.echo(plan_result.stderr)
527
+ click.echo(FeedbackManager.error(message="✗ Error: Terraform plan failed:"))
528
+ click.echo(format_terraform_error(plan_result.stderr, "TERRAFORM PLAN"))
368
529
  return
369
530
 
370
- click.echo(plan_result.stdout)
531
+ # Display formatted plan output
532
+ click.echo(format_terraform_output(plan_result.stdout, "TERRAFORM PLAN"))
371
533
 
372
534
  # Apply Terraform configuration if user confirms
373
535
  if auto_apply or click.confirm("Would you like to apply the Terraform configuration now?"):
374
- click.echo("\nApplying Terraform configuration...\n")
375
- apply_result = subprocess.run(
376
- ["terraform", f"-chdir={infra_dir}", "apply", "-auto-approve"], capture_output=True, text=True
377
- )
536
+ click.echo(FeedbackManager.info(message="** Applying Terraform configuration..."))
537
+ command = f"terraform -chdir={infra_dir} apply -auto-approve"
538
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
539
+ apply_result = subprocess.run(command.split(), capture_output=True, text=True) # noqa: ASYNC221
378
540
 
379
541
  if apply_result.returncode != 0:
380
- click.echo("Terraform apply failed:")
381
- click.echo(apply_result.stderr)
542
+ click.echo(FeedbackManager.error(message="✗ Error: Terraform apply failed:"))
543
+ click.echo(format_terraform_error(apply_result.stderr, "TERRAFORM APPLY"))
382
544
  return
383
545
 
384
- click.echo(apply_result.stdout)
546
+ # Display formatted apply output
547
+ click.echo(format_terraform_output(apply_result.stdout, "TERRAFORM APPLY"))
385
548
 
386
549
  # Prompt to apply the k8s configuration
387
550
  if not skip_apply and (
388
551
  auto_apply or click.confirm("Would you like to apply the Kubernetes configuration now?")
389
552
  ):
390
553
  # Get current kubectl context
391
- current_context_result = subprocess.run(
392
- ["kubectl", "config", "current-context"], capture_output=True, text=True
393
- )
554
+ command = "kubectl config current-context"
555
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
556
+ current_context_result = subprocess.run(command.split(), capture_output=True, text=True) # noqa: ASYNC221
394
557
 
395
558
  current_context = (
396
559
  current_context_result.stdout.strip() if current_context_result.returncode == 0 else "unknown"
397
560
  )
561
+
398
562
  # Get available contexts
399
- contexts_result = subprocess.run(
400
- ["kubectl", "config", "get-contexts", "-o", "name"], capture_output=True, text=True
401
- )
563
+ command = "kubectl config get-contexts -o name"
564
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
565
+
566
+ contexts_result = subprocess.run(command.split(), capture_output=True, text=True) # noqa: ASYNC221
402
567
 
403
568
  if contexts_result.returncode != 0:
404
- click.echo("Failed to get kubectl contexts:")
569
+ click.echo(FeedbackManager.error(message="✗ Error: Failed to get kubectl contexts:"))
405
570
  click.echo(contexts_result.stderr)
406
571
  return
407
572
 
@@ -412,7 +577,10 @@ def infra_init(
412
577
  return
413
578
 
414
579
  # Prompt user to select a context
415
- if len(available_contexts) == 1:
580
+ if config["kubernetes_context"]:
581
+ selected_context = config["kubernetes_context"]
582
+ click.echo(f"Using the kubectl context specified in the config: {selected_context}")
583
+ elif len(available_contexts) == 1:
416
584
  selected_context = available_contexts[0]
417
585
  click.echo(f"Using the only available kubectl context: {selected_context}")
418
586
  else:
@@ -435,67 +603,83 @@ def infra_init(
435
603
  selected_context = available_contexts[selected_index - 1]
436
604
  click.echo(f"Selected context: {selected_context}")
437
605
 
606
+ # Update the config with the selected context
607
+ config["kubernetes_context"] = selected_context
608
+ with open(config_path, "w") as f: # noqa: ASYNC230
609
+ json.dump(config, f, indent=2)
610
+ click.echo(f"Updated configuration with selected Kubernetes context: {selected_context}")
611
+
438
612
  # Apply the configuration to the selected context
439
613
  click.echo(f"Applying Kubernetes configuration to context '{selected_context}'...")
440
- apply_result = subprocess.run(
441
- ["kubectl", "--context", selected_context, "apply", "-f", str(yaml_path)],
442
- capture_output=True,
443
- text=True,
444
- )
614
+
615
+ # First show a diff of what will be applied
616
+ command = f"kubectl --context {selected_context} diff -f {str(yaml_path)}"
617
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
618
+
619
+ diff_result = subprocess.run(command.split(), capture_output=True, text=True) # noqa: ASYNC221
620
+
621
+ if diff_result.returncode not in [0, 1]: # kubectl diff returns 1 when there are differences
622
+ if (
623
+ "Error from server (NotFound): namespaces" in diff_result.stderr
624
+ or f'namespace "{kubernetes_namespace}" not found' in diff_result.stderr
625
+ ):
626
+ click.echo("\nThis appears to be the first deployment - namespace doesn't exist yet.")
627
+ click.echo("No diff available for the initial deployment.")
628
+ else:
629
+ click.echo("Failed to get diff for Kubernetes configuration:")
630
+ click.echo(diff_result.stderr)
631
+ else:
632
+ if diff_result.stdout:
633
+ click.echo("\nChanges that will be applied:")
634
+ click.echo(diff_result.stdout)
635
+ else:
636
+ click.echo("\nNo changes detected or resources don't exist yet.")
637
+
638
+ # Now apply the configuration
639
+ command = f"kubectl --context {selected_context} apply -f {str(yaml_path)}"
640
+ click.echo(FeedbackManager.highlight(message=f"» Executing: {command}"))
641
+ apply_result = subprocess.run(command.split(), capture_output=True, text=True) # noqa: ASYNC221
445
642
 
446
643
  if apply_result.returncode != 0:
447
- click.echo("Failed to apply Kubernetes configuration:")
644
+ click.echo(FeedbackManager.error(message="✗ Error: Failed to apply Kubernetes configuration:"))
448
645
  click.echo(apply_result.stderr)
449
646
  else:
450
- click.echo("Kubernetes configuration applied successfully:")
647
+ click.echo(FeedbackManager.success(message="Kubernetes configuration applied successfully:"))
451
648
  click.echo(apply_result.stdout)
452
649
 
453
- # Get the namespace from the applied configuration
454
- namespace = None
455
- with open(yaml_path, "r") as f:
456
- for line in f:
457
- if "namespace:" in line and not namespace:
458
- namespace = line.split("namespace:")[1].strip()
459
- break
460
-
461
- if not namespace:
462
- namespace = "tinybird" # Default namespace
463
-
464
- click.echo("\nWaiting for load balancer and DNS to be provisioned...")
650
+ click.echo(FeedbackManager.info(message="Waiting for load balancer and DNS to be provisioned..."))
465
651
 
466
652
  max_attempts = 30 # 30 attempts * 10 seconds = 5 minutes
467
653
  endpoint_url = f"https://{dns_record}.{dns_zone_name}"
468
654
 
469
- with click.progressbar(
470
- range(max_attempts),
471
- label=f"Checking endpoint availability: {endpoint_url}",
472
- length=max_attempts,
473
- show_eta=True,
474
- show_percent=True,
475
- fill_char="█",
476
- empty_char="░",
477
- ) as bar:
478
- for attempt in bar:
479
- try:
480
- response = requests.get(endpoint_url, allow_redirects=False, timeout=5)
481
- if response.status_code < 400: # Consider any non-error response as success
482
- click.echo(click.style("\n✅ HTTPS endpoint is now accessible!", fg="green", bold=True))
483
- break
484
- except requests.RequestException:
485
- pass
486
-
487
- if attempt == max_attempts - 1:
488
- click.echo(
489
- click.style(
490
- "\n⚠️ HTTPS endpoint not accessible after 5 minutes", fg="yellow", bold=True
491
- )
492
- )
655
+ click.echo(f"Checking endpoint availability: {endpoint_url}")
656
+
657
+ for attempt in range(max_attempts):
658
+ progress = "#" * (attempt + 1)
659
+ click.echo(
660
+ f"\rAttempt {attempt + 1}/{max_attempts}: Checking if endpoint is ready... [{progress}]",
661
+ nl=False,
662
+ )
663
+
664
+ try:
665
+ response = requests.get(endpoint_url, allow_redirects=False, timeout=5) # noqa: ASYNC210
666
+ if response.status_code < 400: # Consider any non-error response as success
493
667
  click.echo(
494
- " This might be due to DNS propagation or the Load Balancer provisioning delays"
668
+ "\n" + click.style("✅ HTTPS endpoint is now accessible!", fg="green", bold=True)
495
669
  )
496
- click.echo(f" Please try accessing {endpoint_url} manually in a few minutes")
497
- else:
498
- time.sleep(10)
670
+ break
671
+ except (requests.RequestException, requests.Timeout):
672
+ pass
673
+
674
+ if attempt == max_attempts - 1:
675
+ click.echo(
676
+ "\n"
677
+ + click.style("⚠️ HTTPS endpoint not accessible after 5 minutes", fg="yellow", bold=True)
678
+ )
679
+ click.echo(" This might be due to DNS propagation or the Load Balancer provisioning delays")
680
+ click.echo(f" Try accessing {endpoint_url} manually in a few minutes")
681
+ else:
682
+ time.sleep(10) # noqa: ASYNC251
499
683
 
500
684
  if not skip_apply:
501
685
  # Print a summary with the endpoint URL
@@ -515,110 +699,154 @@ def infra_init(
515
699
  @click.argument("name")
516
700
  @click.pass_context
517
701
  @coro
518
- async def infra_rm(ctx: click.Context, name: str):
702
+ async def infra_remove(ctx: click.Context, name: str):
519
703
  """Delete an infrastructure from Tinybird"""
520
- try:
521
- click.echo(FeedbackManager.highlight(message=f"\n» Deleting infrastructure '{name}' from Tinybird..."))
522
- client: TinyB = ctx.ensure_object(dict)["client"]
523
- user_workspaces = await client.user_workspaces_with_organization()
524
- admin_org_id = user_workspaces.get("organization_id")
525
- if not admin_org_id:
526
- raise CLIException("No organization associated to this workspace")
527
- infras = await client.infra_list(admin_org_id)
528
- infra_id = next((infra["id"] for infra in infras if infra["name"] == name), None)
529
- if not infra_id:
530
- raise CLIException(f"Infrastructure '{name}' not found")
531
- await client.infra_delete(infra_id, admin_org_id)
532
- click.echo(FeedbackManager.success(message=f"\n✓ Infrastructure '{name}' deleted"))
533
- except Exception as e:
534
- click.echo(FeedbackManager.error(message=f"✗ Error: {e}"))
704
+ infra: Infrastructure = ctx.ensure_object(dict)["infra"]
705
+ await infra.remove_infra(name)
535
706
 
536
707
 
537
708
  @infra.command(name="ls")
538
709
  @click.pass_context
539
710
  @coro
540
- async def infra_ls(ctx: click.Context):
711
+ async def infra_list(ctx: click.Context):
541
712
  """List self-managed infrastructures"""
713
+ infra: Infrastructure = ctx.ensure_object(dict)["infra"]
714
+ await infra.list_infras()
542
715
 
543
- client: TinyB = ctx.ensure_object(dict)["client"]
544
- config = CLIConfig.get_project_config()
545
- user_client = config.get_client(token=config.get_user_token() or "")
546
- user_workspaces = await user_client.user_workspaces_with_organization()
547
- admin_org_id = user_workspaces.get("organization_id")
548
- infras = await client.infra_list(organization_id=admin_org_id)
549
- columns = [
550
- "name",
551
- "host",
552
- ]
553
- table_human_readable = []
554
- table_machine_readable = []
555
716
 
556
- for infra in infras:
557
- name = infra["name"]
558
- host = infra["host"]
717
+ @infra.command(name="add")
718
+ @click.option("--name", type=str, help="Name for identifying the self-managed region in Tinybird")
719
+ @click.option("--host", type=str, help="Host for the self-managed region")
720
+ @click.option("--organization-id", type=str, help="Organization ID for the self-managed region")
721
+ @click.pass_context
722
+ def infra_add(ctx: click.Context, name: str, host: Optional[str] = None, organization_id: Optional[str] = None):
723
+ """Creates a new self-managed region from an existing infrastructure URL."""
724
+ infra: Infrastructure = ctx.ensure_object(dict)["infra"]
559
725
 
560
- table_human_readable.append((name, host))
561
- table_machine_readable.append({"name": name, "host": host})
726
+ organization_id, organization_name = asyncio.run(ask_for_organization(infra.orgs, organization_id))
727
+ if not organization_id:
728
+ return
562
729
 
563
- click.echo(FeedbackManager.info(message="\n** Infras:"))
564
- echo_safe_humanfriendly_tables_format_smart_table(table_human_readable, column_names=columns)
565
- click.echo("\n")
730
+ name = name or click.prompt("Enter name", type=str)
731
+ host = host or click.prompt(
732
+ "Enter host URL (e.g., https://tinybird.example.com) (leave empty for None)",
733
+ type=str,
734
+ default="",
735
+ show_default=False,
736
+ )
566
737
 
738
+ click.echo(FeedbackManager.highlight(message=f"» Adding self-managed region '{name}' in Tinybird..."))
739
+ new_infra = asyncio.run(infra.create_infra(name, host, organization_id))
740
+ infra_token = new_infra["token"]
741
+ click.echo(
742
+ FeedbackManager.success(message=f"\n✓ Self-managed region '{name}' added in '{organization_name}' Organization")
743
+ )
567
744
 
568
- @infra.command(name="create")
569
- @click.option("--name", type=str, help="Name for identifying the self-managed infrastructure in Tinybird")
570
- @click.option("--host", type=str, help="Host for the infrastructure")
571
- @click.pass_context
572
- @coro
573
- async def infra_create(ctx: click.Context, name: str, host: str):
574
- """Create a self-managed infrastructure"""
575
- try:
576
- client: TinyB = ctx.ensure_object(dict)["client"]
577
- user_workspaces = await client.user_workspaces_with_organization()
578
- admin_org_id = user_workspaces.get("organization_id")
579
- if not admin_org_id:
580
- raise CLIException("No organization associated to this workspace")
581
- name = name or click.prompt("Enter name", type=str)
582
- host = host or click.prompt("Enter host", type=str)
583
- click.echo(FeedbackManager.highlight(message=f"\n» Creating infrastructure '{name}' in Tinybird..."))
584
- infra = await client.infra_create(organization_id=admin_org_id, name=name, host=host)
585
- click.echo(FeedbackManager.success(message=f"\n✓ Infrastructure '{name}' created"))
586
- pyperclip.copy(infra["token"])
587
- click.echo(FeedbackManager.info(message="Access token has been copied to your clipboard."))
588
- click.echo(
589
- FeedbackManager.info(message="Pass it as an environment variable in your deployment as TB_INFRA_TOKEN")
590
- )
591
- except Exception as e:
592
- click.echo(FeedbackManager.error(message=f"✗ Error: {str(e)}"))
745
+ # Get CLI config to access workspace and user information
746
+ cli_config = CLIConfig.get_project_config()
747
+
748
+ # Print environment variables needed for self-managed infrastructure
749
+ click.echo(FeedbackManager.highlight(message="» Required environment variables:"))
750
+ click.echo(f"TB_INFRA_TOKEN={infra_token}")
751
+ click.echo(f"TB_INFRA_WORKSPACE={cli_config.get('name', '')}")
752
+ click.echo(f"TB_INFRA_ORGANIZATION={organization_name}")
753
+ click.echo(f"TB_INFRA_USER={cli_config.get_user_email() or ''}")
593
754
 
594
755
 
595
756
  @infra.command(name="update")
596
757
  @click.argument("infra_name")
597
- @click.option("--name", type=str, help="Name for identifying the self-managed infrastructure in Tinybird")
598
- @click.option("--host", type=str, help="Host for the infrastructure")
758
+ @click.option("--name", type=str, help="Name for identifying the self-managed region in Tinybird")
759
+ @click.option("--host", type=str, help="Host for the self-managed region")
599
760
  @click.pass_context
600
761
  @coro
601
762
  async def infra_update(ctx: click.Context, infra_name: str, name: str, host: str):
602
- """Update a self-managed infrastructure"""
603
- try:
604
- client: TinyB = ctx.ensure_object(dict)["client"]
605
- user_workspaces = await client.user_workspaces_with_organization()
606
- admin_org_id = user_workspaces.get("organization_id")
607
- if not admin_org_id:
608
- raise CLIException("No organization associated to this workspace")
609
-
610
- if not name and not host:
763
+ """Updates the URL of an existing self-managed region."""
764
+ infra: Infrastructure = ctx.ensure_object(dict)["infra"]
765
+ if not name and not host:
766
+ click.echo(FeedbackManager.warning(message="No name or host provided. Provide either a name or a host."))
767
+ return
768
+
769
+ if name or host:
770
+ try:
611
771
  click.echo(
612
- FeedbackManager.warning(message="No name or host provided. Please provide either a name or a host.")
772
+ FeedbackManager.highlight(message=f"» Updating self-managed region'{infra_name}' in Tinybird...")
613
773
  )
614
- return
615
-
616
- if name or host:
617
- infras = await client.infra_list(organization_id=admin_org_id)
618
- infra_id = next((infra["id"] for infra in infras if infra["name"] == infra_name), None)
774
+ infra_id = await infra.update(infra_name, name, host)
619
775
  if not infra_id:
620
- raise CLIException(f"Infrastructure '{infra_name}' not found")
621
- click.echo(FeedbackManager.highlight(message=f"\n» Updating infrastructure '{infra_name}' in Tinybird..."))
622
- await client.infra_update(infra_id=infra_id, organization_id=admin_org_id, name=name, host=host)
623
- except Exception as e:
624
- click.echo(FeedbackManager.error(message=f"✗ Error: {str(e)}"))
776
+ raise CLIException(f"Self-managed region '{infra_name}' not found")
777
+ except Exception as e:
778
+ click.echo(FeedbackManager.error(message=f"✗ Error: {str(e)}"))
779
+
780
+
781
+ def format_terraform_output(output, command_name="Terraform"):
782
+ """Format Terraform command output with colors and styling.
783
+
784
+ Args:
785
+ output: The stdout string from the Terraform command
786
+ command_name: The name of the command (e.g., "PLAN", "APPLY")
787
+
788
+ Returns:
789
+ Formatted output ready to be displayed
790
+ """
791
+ if not output:
792
+ return FeedbackManager.info(message=f"No output from {command_name} command.")
793
+
794
+ # Add a header
795
+ formatted_lines = [
796
+ "\n" + "=" * 80,
797
+ FeedbackManager.highlight(message=f"{command_name.upper()} OUTPUT"),
798
+ "=" * 80 + "\n",
799
+ ]
800
+
801
+ # Process and format the output
802
+ for line in output.splitlines():
803
+ # Highlight resource changes
804
+ if line.strip().startswith("+ "):
805
+ formatted_lines.append(click.style(line, fg="green"))
806
+ elif line.strip().startswith("- "):
807
+ formatted_lines.append(click.style(line, fg="red"))
808
+ elif line.strip().startswith("~ "):
809
+ formatted_lines.append(click.style(line, fg="yellow"))
810
+ # Highlight plan/apply summary
811
+ elif any(keyword in line for keyword in ["Plan:", "Apply complete", "Destroy complete"]):
812
+ formatted_lines.append("\n" + click.style(line, bold=True, fg="cyan"))
813
+ else:
814
+ formatted_lines.append(line)
815
+
816
+ # Add a footer
817
+ formatted_lines.append("\n" + "=" * 80)
818
+
819
+ return "\n".join(formatted_lines)
820
+
821
+
822
+ def format_terraform_error(error_output, command_name="Terraform"):
823
+ """Format Terraform command error output with colors and styling.
824
+
825
+ Args:
826
+ error_output: The stderr string from the Terraform command
827
+ command_name: The name of the command (e.g., "PLAN", "APPLY")
828
+
829
+ Returns:
830
+ Formatted error output ready to be displayed
831
+ """
832
+ if not error_output:
833
+ return FeedbackManager.error(message=f"Unknown error in {command_name} command.")
834
+
835
+ # Add a header
836
+ formatted_lines = ["\n" + "=" * 80, FeedbackManager.error(message=f"{command_name.upper()} ERROR"), "=" * 80 + "\n"]
837
+
838
+ # Process and format the error output
839
+ for line in error_output.splitlines():
840
+ # Highlight error messages
841
+ if any(keyword in line.lower() for keyword in ["error", "failed", "fatal"]):
842
+ formatted_lines.append(click.style(line, fg="red", bold=True))
843
+ # Highlight warnings
844
+ elif any(keyword in line.lower() for keyword in ["warning", "deprecated"]):
845
+ formatted_lines.append(click.style(line, fg="yellow"))
846
+ else:
847
+ formatted_lines.append(line)
848
+
849
+ # Add a footer
850
+ formatted_lines.append("\n" + "=" * 80)
851
+
852
+ return "\n".join(formatted_lines)