tinybird 0.0.1.dev4__py3-none-any.whl → 0.0.1.dev6__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.

Files changed (43) hide show
  1. tinybird/__cli__.py +7 -8
  2. tinybird/check_pypi.py +1 -1
  3. tinybird/feedback_manager.py +2 -2
  4. tinybird/tb/cli.py +28 -0
  5. tinybird/{tb_cli_modules → tb/modules}/auth.py +5 -5
  6. tinybird/{tb_cli_modules → tb/modules}/branch.py +6 -5
  7. tinybird/{tb_cli_modules → tb/modules}/build.py +8 -26
  8. tinybird/tb/modules/cicd.py +271 -0
  9. tinybird/{tb_cli_modules → tb/modules}/cli.py +24 -91
  10. tinybird/tb/modules/common.py +2098 -0
  11. tinybird/tb/modules/config.py +352 -0
  12. tinybird/{tb_cli_modules → tb/modules}/connection.py +4 -4
  13. tinybird/{tb_cli_modules → tb/modules}/create.py +11 -7
  14. tinybird/{datafile.py → tb/modules/datafile.py} +6 -7
  15. tinybird/{tb_cli_modules → tb/modules}/datasource.py +7 -6
  16. tinybird/tb/modules/exceptions.py +91 -0
  17. tinybird/{tb_cli_modules → tb/modules}/fmt.py +3 -3
  18. tinybird/{tb_cli_modules → tb/modules}/job.py +3 -3
  19. tinybird/{tb_cli_modules → tb/modules}/llm.py +14 -13
  20. tinybird/{tb_cli_modules → tb/modules}/local.py +11 -6
  21. tinybird/tb/modules/mock.py +53 -0
  22. tinybird/{tb_cli_modules → tb/modules}/pipe.py +5 -5
  23. tinybird/{tb_cli_modules → tb/modules}/prompts.py +3 -0
  24. tinybird/tb/modules/regions.py +9 -0
  25. tinybird/tb/modules/table.py +185 -0
  26. tinybird/{tb_cli_modules → tb/modules}/tag.py +2 -2
  27. tinybird/tb/modules/telemetry.py +310 -0
  28. tinybird/{tb_cli_modules → tb/modules}/test.py +5 -5
  29. tinybird/{tb_cli_modules → tb/modules}/tinyunit/tinyunit.py +1 -1
  30. tinybird/{tb_cli_modules → tb/modules}/token.py +3 -3
  31. tinybird/{tb_cli_modules → tb/modules}/workspace.py +5 -5
  32. tinybird/{tb_cli_modules → tb/modules}/workspace_members.py +4 -4
  33. tinybird/tb_cli_modules/common.py +9 -7
  34. tinybird/tb_cli_modules/config.py +0 -8
  35. {tinybird-0.0.1.dev4.dist-info → tinybird-0.0.1.dev6.dist-info}/METADATA +1 -1
  36. tinybird-0.0.1.dev6.dist-info/RECORD +58 -0
  37. tinybird-0.0.1.dev6.dist-info/entry_points.txt +2 -0
  38. tinybird/tb_cli.py +0 -27
  39. tinybird-0.0.1.dev4.dist-info/RECORD +0 -50
  40. tinybird-0.0.1.dev4.dist-info/entry_points.txt +0 -2
  41. /tinybird/{tb_cli_modules → tb/modules}/tinyunit/tinyunit_lib.py +0 -0
  42. {tinybird-0.0.1.dev4.dist-info → tinybird-0.0.1.dev6.dist-info}/WHEEL +0 -0
  43. {tinybird-0.0.1.dev4.dist-info → tinybird-0.0.1.dev6.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,2098 @@
1
+ # This is the common file for our CLI. Please keep it clean (as possible)
2
+ #
3
+ # - Put here any common utility function you consider.
4
+ # - If any function is only called within a specific command, consider moving
5
+ # the function to the proper command file.
6
+ # - Please, **do not** define commands here.
7
+
8
+ import asyncio
9
+ import json
10
+ import os
11
+ import re
12
+ import socket
13
+ import sys
14
+ import uuid
15
+ from contextlib import closing
16
+ from copy import deepcopy
17
+ from enum import Enum
18
+ from functools import wraps
19
+ from os import chmod, environ, getcwd, getenv
20
+ from pathlib import Path
21
+ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
22
+ from urllib.parse import urlparse
23
+
24
+ import aiofiles
25
+ import click
26
+ import click.formatting
27
+ import humanfriendly
28
+ import humanfriendly.tables
29
+ import pyperclip
30
+ from click import Context
31
+ from click._termui_impl import ProgressBar
32
+ from humanfriendly.tables import format_pretty_table
33
+ from packaging.version import Version
34
+
35
+ from tinybird.client import (
36
+ AuthException,
37
+ AuthNoTokenException,
38
+ ConnectorNothingToLoad,
39
+ DoesNotExistException,
40
+ JobException,
41
+ OperationCanNotBePerformed,
42
+ TinyB,
43
+ )
44
+ from tinybird.config import (
45
+ DEFAULT_API_HOST,
46
+ DEFAULT_UI_HOST,
47
+ DEPRECATED_PROJECT_PATHS,
48
+ PROJECT_PATHS,
49
+ SUPPORTED_CONNECTORS,
50
+ VERSION,
51
+ FeatureFlags,
52
+ get_config,
53
+ get_display_host,
54
+ write_config,
55
+ )
56
+ from tinybird.tb.modules.table import format_table
57
+
58
+ if TYPE_CHECKING:
59
+ from tinybird.connectors import Connector
60
+
61
+ from tinybird.feedback_manager import FeedbackManager, warning_message
62
+ from tinybird.git_settings import DEFAULT_TINYENV_FILE
63
+ from tinybird.syncasync import async_to_sync
64
+ from tinybird.tb.modules.cicd import APPEND_FIXTURES_SH, DEFAULT_REQUIREMENTS_FILE, EXEC_TEST_SH
65
+ from tinybird.tb.modules.config import CLIConfig
66
+ from tinybird.tb.modules.exceptions import (
67
+ CLIAuthException,
68
+ CLIConnectionException,
69
+ CLIException,
70
+ CLIReleaseException,
71
+ CLIWorkspaceException,
72
+ )
73
+ from tinybird.tb.modules.regions import Region
74
+ from tinybird.tb.modules.telemetry import (
75
+ add_telemetry_event,
76
+ add_telemetry_sysinfo_event,
77
+ flush_telemetry,
78
+ init_telemetry,
79
+ is_ci_environment,
80
+ )
81
+
82
+ SUPPORTED_FORMATS = ["csv", "ndjson", "json", "parquet"]
83
+ OLDEST_ROLLBACK = "oldest_rollback"
84
+ MAIN_BRANCH = "main"
85
+
86
+
87
+ def obfuscate_token(value: Optional[str]) -> Optional[str]:
88
+ if not value:
89
+ return None
90
+ return f"{value[:4]}...{value[-8:]}"
91
+
92
+
93
+ def create_connector(connector: str, options: Dict[str, Any]):
94
+ # Imported here to improve startup time when the connectors aren't used
95
+ from tinybird.connectors import UNINSTALLED_CONNECTORS
96
+ from tinybird.connectors import create_connector as _create_connector
97
+
98
+ if connector in UNINSTALLED_CONNECTORS:
99
+ raise CLIException(FeedbackManager.error_connector_not_installed(connector=connector))
100
+ return _create_connector(connector, options)
101
+
102
+
103
+ def coro(f):
104
+ @wraps(f)
105
+ def wrapper(*args, **kwargs):
106
+ return asyncio.run(f(*args, **kwargs))
107
+
108
+ return wrapper
109
+
110
+
111
+ async def gather_with_concurrency(n, *tasks):
112
+ semaphore = asyncio.Semaphore(n)
113
+
114
+ async def sem_task(task):
115
+ async with semaphore:
116
+ return await task
117
+
118
+ return await asyncio.gather(*(sem_task(task) for task in tasks))
119
+
120
+
121
+ def echo_safe_humanfriendly_tables_format_smart_table(data: Iterable[Any], column_names: List[str]) -> None:
122
+ """
123
+ There is a bug in the humanfriendly library: it breaks to render the small table for small terminals
124
+ (`format_robust_table`) if we call format_smart_table with an empty dataset. This catches the error and prints
125
+ what we would call an empty "robust_table".
126
+ """
127
+ try:
128
+ click.echo(humanfriendly.tables.format_smart_table(data, column_names=column_names))
129
+ except ValueError as exc:
130
+ if str(exc) == "max() arg is an empty sequence":
131
+ click.echo("------------")
132
+ click.echo("Empty")
133
+ click.echo("------------")
134
+ else:
135
+ raise exc
136
+
137
+
138
+ def echo_safe_format_table(data: Iterable[Any], columns) -> None:
139
+ """
140
+ There is a bug in the humanfriendly library: it breaks to render the small table for small terminals
141
+ (`format_robust_table`) if we call format_smart_table with an empty dataset. This catches the error and prints
142
+ what we would call an empty "robust_table".
143
+ """
144
+ try:
145
+ click.echo(format_table(data, columns))
146
+ except ValueError as exc:
147
+ if str(exc) == "max() arg is an empty sequence":
148
+ click.echo("------------")
149
+ click.echo("Empty")
150
+ click.echo("------------")
151
+ else:
152
+ raise exc
153
+
154
+
155
+ def normalize_datasource_name(s: str) -> str:
156
+ s = re.sub(r"[^0-9a-zA-Z_]", "_", s)
157
+ if s[0] in "0123456789":
158
+ return "c_" + s
159
+ return s
160
+
161
+
162
+ def generate_datafile(
163
+ datafile: str, filename: str, data: Optional[bytes], force: Optional[bool] = False, _format: Optional[str] = "csv"
164
+ ):
165
+ p = Path(filename)
166
+ base = Path("datasources")
167
+ datasource_name = normalize_datasource_name(p.stem)
168
+ if not base.exists():
169
+ base = Path()
170
+ f = base / (datasource_name + ".datasource")
171
+ if not f.exists() or force:
172
+ with open(f"{f}", "w") as ds_file:
173
+ ds_file.write(datafile)
174
+ click.echo(FeedbackManager.info_file_created(file=f))
175
+
176
+ if data and (base / "fixtures").exists():
177
+ # Generating a fixture for Parquet files is not so trivial, since Parquet format
178
+ # is column-based. We would need to add PyArrow as a dependency (which is huge)
179
+ # just to analyze the whole Parquet file to extract one single row.
180
+ if _format == "parquet":
181
+ click.echo(FeedbackManager.warning_parquet_fixtures_not_supported())
182
+ else:
183
+ f = base / "fixtures" / (p.stem + f".{_format}")
184
+ newline = b"\n" # TODO: guess
185
+ with open(f, "wb") as fixture_file:
186
+ fixture_file.write(data[: data.rfind(newline)])
187
+ else:
188
+ click.echo(FeedbackManager.error_file_already_exists(file=f))
189
+
190
+
191
+ async def get_current_workspace(config: CLIConfig) -> Optional[Dict[str, Any]]:
192
+ client = config.get_client()
193
+ workspaces: List[Dict[str, Any]] = (await client.user_workspaces_and_branches()).get("workspaces", [])
194
+ return _get_current_workspace_common(workspaces, config["id"])
195
+
196
+
197
+ def get_workspace_member_email(workspace, member_id) -> str:
198
+ return next((member["email"] for member in workspace["members"] if member["id"] == member_id), "-")
199
+
200
+
201
+ def _get_current_workspace_common(
202
+ workspaces: List[Dict[str, Any]], current_workspace_id: str
203
+ ) -> Optional[Dict[str, Any]]:
204
+ return next((workspace for workspace in workspaces if workspace["id"] == current_workspace_id), None)
205
+
206
+
207
+ async def get_current_environment(client, config):
208
+ workspaces: List[Dict[str, Any]] = (await client.user_workspaces_and_branches()).get("workspaces", [])
209
+ return next((workspace for workspace in workspaces if workspace["id"] == config["id"]), None)
210
+
211
+
212
+ async def get_current_workspace_branches(config: CLIConfig) -> List[Dict[str, Any]]:
213
+ current_main_workspace: Optional[Dict[str, Any]] = await get_current_main_workspace(config)
214
+ if not current_main_workspace:
215
+ raise CLIException(FeedbackManager.error_unable_to_identify_main_workspace())
216
+
217
+ client = config.get_client()
218
+ user_branches: List[Dict[str, Any]] = (await client.user_workspace_branches()).get("workspaces", [])
219
+ all_branches: List[Dict[str, Any]] = (await client.branches()).get("environments", [])
220
+ branches = user_branches + [branch for branch in all_branches if branch not in user_branches]
221
+
222
+ return [branch for branch in branches if branch.get("main") == current_main_workspace["id"]]
223
+
224
+
225
+ class AliasedGroup(click.Group):
226
+ def get_command(self, ctx, cmd_name):
227
+ # Step one: built-in commands as normal
228
+ cm = click.Group.get_command(self, ctx, cmd_name)
229
+ if cm is not None:
230
+ return cm
231
+
232
+ def resolve_command(self, ctx, args):
233
+ # always return the command's name, not the alias
234
+ _, cmd, args = super().resolve_command(ctx, args)
235
+ return cmd.name, cmd, args # type: ignore[union-attr]
236
+
237
+
238
+ class CatchAuthExceptions(AliasedGroup):
239
+ """utility class to get all the auth exceptions"""
240
+
241
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
242
+ init_telemetry()
243
+ add_telemetry_sysinfo_event()
244
+ super().__init__(*args, **kwargs)
245
+
246
+ def format_epilog(self, ctx: Context, formatter: click.formatting.HelpFormatter) -> None:
247
+ super().format_epilog(ctx, formatter)
248
+
249
+ formatter.write_paragraph()
250
+ formatter.write_heading("Telemetry")
251
+ formatter.write_text(
252
+ """
253
+ Tinybird collects anonymous usage data and errors to improve the command
254
+ line experience. To opt-out, set TB_CLI_TELEMETRY_OPTOUT environment
255
+ variable to '1' or 'true'."""
256
+ )
257
+ formatter.write_paragraph()
258
+
259
+ def __call__(self, *args, **kwargs) -> None:
260
+ error_msg: Optional[str] = None
261
+ error_event: str = "error"
262
+
263
+ exit_code: int = 0
264
+
265
+ try:
266
+ self.main(*args, **kwargs)
267
+ except AuthNoTokenException:
268
+ error_msg = FeedbackManager.error_notoken()
269
+ error_event = "auth_error"
270
+ exit_code = 1
271
+ except AuthException as ex:
272
+ error_msg = FeedbackManager.error_exception(error=str(ex))
273
+ error_event = "auth_error"
274
+ exit_code = 1
275
+ except SystemExit as ex:
276
+ exit_code = int(ex.code) if ex.code else 0
277
+ except Exception as ex:
278
+ error_msg = str(ex)
279
+ exit_code = 1
280
+
281
+ if error_msg:
282
+ click.echo(error_msg)
283
+ add_telemetry_event(error_event, error=error_msg)
284
+ flush_telemetry(wait=True)
285
+
286
+ sys.exit(exit_code)
287
+
288
+
289
+ def load_connector_config(ctx: Context, connector_name: str, debug: bool, check_uninstalled: bool = False):
290
+ config_file = Path(getcwd()) / f".tinyb_{connector_name}"
291
+ try:
292
+ if connector_name not in ctx.ensure_object(dict):
293
+ with open(config_file) as file:
294
+ config = json.loads(file.read())
295
+ from tinybird.connectors import UNINSTALLED_CONNECTORS
296
+
297
+ if check_uninstalled and connector_name in UNINSTALLED_CONNECTORS:
298
+ click.echo(FeedbackManager.warning_connector_not_installed(connector=connector_name))
299
+ return
300
+ ctx.ensure_object(dict)[connector_name] = create_connector(connector_name, config)
301
+ except IOError:
302
+ if debug:
303
+ click.echo(f"** {connector_name} connector not configured")
304
+ pass
305
+
306
+
307
+ def getenv_bool(key: str, default: bool) -> bool:
308
+ v: Optional[str] = getenv(key)
309
+ if v is None:
310
+ return default
311
+ return v.lower() == "true" or v == "1"
312
+
313
+
314
+ def _get_tb_client(token: str, host: str, semver: Optional[str] = None) -> TinyB:
315
+ disable_ssl: bool = getenv_bool("TB_DISABLE_SSL_CHECKS", False)
316
+ return TinyB(token, host, version=VERSION, disable_ssl_checks=disable_ssl, send_telemetry=True, semver=semver)
317
+
318
+
319
+ def create_tb_client(ctx: Context) -> TinyB:
320
+ token = ctx.ensure_object(dict)["config"].get("token", "")
321
+ host = ctx.ensure_object(dict)["config"].get("host", DEFAULT_API_HOST)
322
+ semver = ctx.ensure_object(dict)["config"].get("semver", "")
323
+ return _get_tb_client(token, host, semver=semver)
324
+
325
+
326
+ async def _analyze(filename: str, client: TinyB, format: str, connector: Optional["Connector"] = None):
327
+ data: Optional[bytes] = None
328
+ if not connector:
329
+ parsed = urlparse(filename)
330
+ if parsed.scheme in ("http", "https"):
331
+ meta = await client.datasource_analyze(filename)
332
+ else:
333
+ async with aiofiles.open(filename, "rb") as file:
334
+ # We need to read the whole file in binary for Parquet, while for the
335
+ # others we just read 1KiB
336
+ if format == "parquet":
337
+ data = await file.read()
338
+ else:
339
+ data = await file.read(1024 * 1024)
340
+
341
+ meta = await client.datasource_analyze_file(data)
342
+ else:
343
+ meta = connector.datasource_analyze(filename)
344
+ return meta, data
345
+
346
+
347
+ async def _generate_datafile(
348
+ filename: str, client: TinyB, format: str, connector: Optional["Connector"] = None, force: Optional[bool] = False
349
+ ):
350
+ meta, data = await _analyze(filename, client, format, connector=connector)
351
+ schema = meta["analysis"]["schema"]
352
+ schema = schema.replace(", ", ",\n ")
353
+ datafile = f"""DESCRIPTION >\n Generated from {filename}\n\nSCHEMA >\n {schema}"""
354
+ return generate_datafile(datafile, filename, data, force, _format=format)
355
+
356
+
357
+ async def folder_init(
358
+ client: TinyB,
359
+ folder: str,
360
+ generate_datasources: Optional[bool] = False,
361
+ force: Optional[bool] = False,
362
+ generate_releases: Optional[bool] = False,
363
+ ):
364
+ for x in filter(lambda x: x not in DEPRECATED_PROJECT_PATHS, PROJECT_PATHS):
365
+ try:
366
+ f = Path(folder) / x
367
+ f.mkdir()
368
+ click.echo(FeedbackManager.info_path_created(path=x))
369
+ except FileExistsError:
370
+ if not force:
371
+ click.echo(FeedbackManager.info_path_already_exists(path=x))
372
+ pass
373
+
374
+ if generate_datasources:
375
+ for format in SUPPORTED_FORMATS:
376
+ for path in Path(folder).glob(f"*.{format}"):
377
+ await _generate_datafile(str(path), client, format=format, force=force)
378
+
379
+ if generate_releases:
380
+ base = Path(".")
381
+ f = base / (".tinyenv")
382
+ if not f.exists() or force:
383
+ async with aiofiles.open(".tinyenv", "w") as file:
384
+ await file.write(DEFAULT_TINYENV_FILE)
385
+ click.echo(FeedbackManager.info_file_created(file=".tinyenv"))
386
+ else:
387
+ click.echo(FeedbackManager.info_dottinyenv_already_exists())
388
+
389
+ base = Path(".")
390
+ f = base / ("requirements.txt")
391
+ if not f.exists() or force:
392
+ async with aiofiles.open("requirements.txt", "w") as file:
393
+ await file.write(DEFAULT_REQUIREMENTS_FILE)
394
+ click.echo(FeedbackManager.info_file_created(file="requirements.txt"))
395
+
396
+ base = Path("scripts")
397
+ if not base.exists():
398
+ base = Path()
399
+ f = base / ("exec_test.sh")
400
+ if not f.exists() or force:
401
+ async with aiofiles.open(f"{f}", "w") as t_file:
402
+ await t_file.write(EXEC_TEST_SH)
403
+ click.echo(FeedbackManager.info_file_created(file="scripts/exec_test.sh"))
404
+ chmod(f, 0o755)
405
+
406
+ f = base / ("append_fixtures.sh")
407
+ if not f.exists() or force:
408
+ async with aiofiles.open(f"{f}", "w") as t_file:
409
+ await t_file.write(APPEND_FIXTURES_SH)
410
+ click.echo(FeedbackManager.info_file_created(file="scripts/append_fixtures.sh"))
411
+ chmod(f, 0o755)
412
+
413
+ base = Path("tests")
414
+ if not base.exists():
415
+ base = Path()
416
+ f = base / ("example.yml")
417
+ if not base.exists() or force:
418
+ async with aiofiles.open(f"{f}", "w") as t_file:
419
+ await t_file.write(
420
+ """
421
+ ##############################################################################################################################
422
+ ### Visit https://www.tinybird.co/docs/production/implementing-test-strategies.html#data-quality-tests ###
423
+ ### for more details on Data Quality tests ###
424
+ ##############################################################################################################################
425
+
426
+ - example_no_negative_numbers:
427
+ max_bytes_read: null
428
+ max_time: null
429
+ sql: |
430
+ SELECT
431
+ number
432
+ FROM numbers(10)
433
+ WHERE
434
+ number < 0
435
+
436
+ # - example_top_products_params_no_empty_top_10_on_2023:
437
+ # max_bytes_read: null
438
+ # max_time: null
439
+ # sql: |
440
+ # SELECT *
441
+ # FROM top_products_params
442
+ # WHERE empty(top_10)
443
+ # pipe:
444
+ # name: top_products_params
445
+ # params:
446
+ # start: '2023-01-01'
447
+ # end: '2023-12-31'
448
+
449
+ """
450
+ )
451
+
452
+ f = base / ("regression.yaml")
453
+ if not base.exists() or force:
454
+ async with aiofiles.open(f"{f}", "w") as t_file:
455
+ await t_file.write(
456
+ """
457
+ ############################################################################################################################
458
+ ### Visit https://www.tinybird.co/docs/production/implementing-test-strategies.html#regression-tests ###
459
+ ### for more details on Regression tests ###
460
+ ############################################################################################################################
461
+
462
+ ###
463
+ ### New pipes are covered by this rule, rules below this one supersede this setting
464
+ ###
465
+ - pipe: '.*'
466
+ tests:
467
+ - coverage:
468
+
469
+
470
+
471
+ ###
472
+ ### These are rules to customize regression testing by pipe using regular expressions
473
+ ### For instance skip regression tests for the pipes matching `endpoint_name.*`
474
+ ###
475
+ - pipe: 'endpoint_name.*'
476
+ tests:
477
+ - coverage:
478
+ config:
479
+ skip: True
480
+
481
+ """
482
+ )
483
+
484
+
485
+ async def configure_connector(connector):
486
+ if connector not in SUPPORTED_CONNECTORS:
487
+ raise CLIException(FeedbackManager.error_invalid_connector(connectors=", ".join(SUPPORTED_CONNECTORS)))
488
+
489
+ file_name = f".tinyb_{connector}"
490
+ config_file = Path(getcwd()) / file_name
491
+ if connector == "bigquery":
492
+ project = click.prompt("BigQuery project ID")
493
+ service_account = click.prompt(
494
+ "Path to a JSON service account file with permissions to export from BigQuery, write in Storage and sign URLs (leave empty to use GOOGLE_APPLICATION_CREDENTIALS environment variable)",
495
+ default=environ.get("GOOGLE_APPLICATION_CREDENTIALS", ""),
496
+ )
497
+ bucket_name = click.prompt("Name of a Google Cloud Storage bucket to store temporary exported files")
498
+
499
+ try:
500
+ config = {"project_id": project, "service_account": service_account, "bucket_name": bucket_name}
501
+ await write_config(config, file_name)
502
+ except Exception:
503
+ raise CLIException(FeedbackManager.error_file_config(config_file=config_file))
504
+ elif connector == "snowflake":
505
+ sf_account = click.prompt("Snowflake Account (e.g. your-domain.west-europe.azure)")
506
+ sf_warehouse = click.prompt("Snowflake warehouse name")
507
+ sf_database = click.prompt("Snowflake database name")
508
+ sf_schema = click.prompt("Snowflake schema name")
509
+ sf_role = click.prompt("Snowflake role name")
510
+ sf_user = click.prompt("Snowflake user name")
511
+ sf_password = click.prompt("Snowflake password")
512
+ sf_storage_integration = click.prompt(
513
+ "Snowflake GCS storage integration name (leave empty to auto-generate one)", default=""
514
+ )
515
+ sf_stage = click.prompt("Snowflake GCS stage name (leave empty to auto-generate one)", default="")
516
+ project = click.prompt("Google Cloud project ID to store temporary files")
517
+ service_account = click.prompt(
518
+ "Path to a JSON service account file with permissions to write in Storagem, sign URLs and IAM (leave empty to use GOOGLE_APPLICATION_CREDENTIALS environment variable)",
519
+ default=environ.get("GOOGLE_APPLICATION_CREDENTIALS", ""),
520
+ )
521
+ bucket_name = click.prompt("Name of a Google Cloud Storage bucket to store temporary exported files")
522
+
523
+ if not service_account:
524
+ service_account = getenv("GOOGLE_APPLICATION_CREDENTIALS")
525
+
526
+ try:
527
+ config = {
528
+ "account": sf_account,
529
+ "warehouse": sf_warehouse,
530
+ "database": sf_database,
531
+ "schema": sf_schema,
532
+ "role": sf_role,
533
+ "user": sf_user,
534
+ "password": sf_password,
535
+ "storage_integration": sf_storage_integration,
536
+ "stage": sf_stage,
537
+ "service_account": service_account,
538
+ "bucket_name": bucket_name,
539
+ "project_id": project,
540
+ }
541
+ await write_config(config, file_name)
542
+ except Exception:
543
+ raise CLIException(FeedbackManager.error_file_config(config_file=config_file))
544
+
545
+ click.echo(FeedbackManager.success_connector_config(connector=connector, file_name=file_name))
546
+
547
+
548
+ def _compare_region_host(region_name_or_host: str, region: Dict[str, Any]) -> bool:
549
+ if region["name"].lower() == region_name_or_host:
550
+ return True
551
+ if region["host"] == region_name_or_host:
552
+ return True
553
+ if region["api_host"] == region_name_or_host:
554
+ return True
555
+ return False
556
+
557
+
558
+ def ask_for_region_interactively(regions):
559
+ region_index = -1
560
+
561
+ while region_index == -1:
562
+ click.echo(FeedbackManager.info_available_regions())
563
+ for index, region in enumerate(regions):
564
+ provider = f" ({region.get('provider')})" if region.get("provider") else ""
565
+ click.echo(f" [{index + 1}] {region['name'].lower()}{provider} ({region['host']}) ")
566
+ click.echo(" [0] Cancel")
567
+
568
+ region_index = click.prompt("\nUse region", default=1)
569
+
570
+ if region_index == 0:
571
+ click.echo(FeedbackManager.info_auth_cancelled_by_user())
572
+ return None
573
+
574
+ try:
575
+ return regions[int(region_index) - 1]
576
+ except Exception:
577
+ available_options = ", ".join(map(str, range(1, len(regions) + 1)))
578
+ click.echo(FeedbackManager.error_region_index(host_index=region_index, available_options=available_options))
579
+ region_index = -1
580
+
581
+
582
+ def get_region_info(ctx, region=None):
583
+ name = region["name"] if region else "default"
584
+ api_host = format_host(
585
+ region["api_host"] if region else ctx.obj["config"].get("host", DEFAULT_API_HOST), subdomain="api"
586
+ )
587
+ ui_host = format_host(region["host"] if region else ctx.obj["config"].get("host", DEFAULT_UI_HOST), subdomain="ui")
588
+ return name, api_host, ui_host
589
+
590
+
591
+ def format_host(host: str, subdomain: Optional[str] = None) -> str:
592
+ """
593
+ >>> format_host('api.tinybird.co')
594
+ 'https://api.tinybird.co'
595
+ >>> format_host('https://api.tinybird.co')
596
+ 'https://api.tinybird.co'
597
+ >>> format_host('http://localhost:8001')
598
+ 'http://localhost:8001'
599
+ >>> format_host('localhost:8001')
600
+ 'http://localhost:8001'
601
+ >>> format_host('localhost:8001', subdomain='ui')
602
+ 'http://localhost:8001'
603
+ >>> format_host('localhost:8001', subdomain='api')
604
+ 'http://localhost:8001'
605
+ >>> format_host('https://api.tinybird.co', subdomain='ui')
606
+ 'https://ui.tinybird.co'
607
+ >>> format_host('https://api.us-east.tinybird.co', subdomain='ui')
608
+ 'https://ui.us-east.tinybird.co'
609
+ >>> format_host('https://api.us-east.tinybird.co', subdomain='api')
610
+ 'https://api.us-east.tinybird.co'
611
+ >>> format_host('https://ui.us-east.tinybird.co', subdomain='api')
612
+ 'https://api.us-east.tinybird.co'
613
+ >>> format_host('https://inditex-rt-pro.tinybird.co', subdomain='ui')
614
+ 'https://inditex-rt-pro.tinybird.co'
615
+ >>> format_host('https://cluiente-tricky.tinybird.co', subdomain='api')
616
+ 'https://cluiente-tricky.tinybird.co'
617
+ """
618
+ is_localhost = FeatureFlags.is_localhost()
619
+ if subdomain and not is_localhost:
620
+ url_info = urlparse(host)
621
+ current_subdomain = url_info.netloc.split(".")[0]
622
+ if current_subdomain == "api" or current_subdomain == "ui":
623
+ host = host.replace(current_subdomain, subdomain)
624
+ if "localhost" in host or is_localhost:
625
+ host = f"http://{host}" if "http" not in host else host
626
+ elif not host.startswith("http"):
627
+ host = f"https://{host}"
628
+ return host
629
+
630
+
631
+ def region_from_host(region_name_or_host, regions):
632
+ """Returns the region that matches region_name_or_host"""
633
+
634
+ return next((r for r in regions if _compare_region_host(region_name_or_host, r)), None)
635
+
636
+
637
+ def ask_for_user_token(action: str, ui_host: str) -> str:
638
+ return click.prompt(
639
+ f'\nUse the token called "user token" in order to {action}. Copy it from {ui_host}/tokens and paste it here',
640
+ hide_input=True,
641
+ show_default=False,
642
+ default=None,
643
+ )
644
+
645
+
646
+ async def check_user_token(ctx: Context, token: str):
647
+ client: TinyB = ctx.ensure_object(dict)["client"]
648
+ try:
649
+ user_client: TinyB = deepcopy(client)
650
+ user_client.token = token
651
+
652
+ is_authenticated = await user_client.check_auth_login()
653
+ except Exception as e:
654
+ raise CLIWorkspaceException(FeedbackManager.error_exception(error=str(e)))
655
+
656
+ if not is_authenticated.get("is_valid", False):
657
+ raise CLIWorkspaceException(
658
+ FeedbackManager.error_exception(
659
+ error='Invalid token. Please, be sure you are using the "user token" instead of the "admin your@email" token.'
660
+ )
661
+ )
662
+ if is_authenticated.get("is_valid") and not is_authenticated.get("is_user", False):
663
+ raise CLIWorkspaceException(
664
+ FeedbackManager.error_exception(
665
+ error='Invalid user authentication. Please, be sure you are using the "user token" instead of the "admin your@email" token.'
666
+ )
667
+ )
668
+
669
+
670
+ async def get_available_starterkits(ctx: Context) -> List[Dict[str, Any]]:
671
+ ctx_dict = ctx.ensure_object(dict)
672
+ available_starterkits = ctx_dict.get("available_starterkits", None)
673
+ if available_starterkits is not None:
674
+ return available_starterkits
675
+
676
+ try:
677
+ client: TinyB = ctx_dict["client"]
678
+
679
+ available_starterkits = await client.starterkits()
680
+ ctx_dict["available_starterkits"] = available_starterkits
681
+ return available_starterkits
682
+ except Exception as ex:
683
+ raise CLIException(FeedbackManager.error_exception(error=ex))
684
+
685
+
686
+ async def get_starterkit(ctx: Context, name: str) -> Optional[Dict[str, Any]]:
687
+ available_starterkits = await get_available_starterkits(ctx)
688
+ if not available_starterkits:
689
+ return None
690
+ return next((sk for sk in available_starterkits if sk.get("friendly_name", None) == name), None)
691
+
692
+
693
+ async def is_valid_starterkit(ctx: Context, name: str) -> bool:
694
+ return name == "blank" or (await get_starterkit(ctx, name) is not None)
695
+
696
+
697
+ async def ask_for_starterkit_interactively(ctx: Context) -> Optional[str]:
698
+ starterkit = [{"friendly_name": "blank", "description": "Empty workspace"}]
699
+ starterkit.extend(await get_available_starterkits(ctx))
700
+ rows = [(index + 1, sk["friendly_name"], sk["description"]) for index, sk in enumerate(starterkit)]
701
+
702
+ echo_safe_humanfriendly_tables_format_smart_table(rows, column_names=["Idx", "Id", "Description"])
703
+ click.echo("")
704
+ click.echo(" [0] to cancel")
705
+
706
+ sk_index = -1
707
+ while sk_index == -1:
708
+ sk_index = click.prompt("\nUse starter kit", default=1)
709
+ if sk_index < 0 or sk_index > len(starterkit):
710
+ click.echo(FeedbackManager.error_starterkit_index(starterkit_index=sk_index))
711
+ sk_index = -1
712
+
713
+ if sk_index == 0:
714
+ click.echo(FeedbackManager.info_cancelled_by_user())
715
+ return None
716
+
717
+ return starterkit[sk_index - 1]["friendly_name"]
718
+
719
+
720
+ async def fork_workspace(client: TinyB, user_client: TinyB, created_workspace):
721
+ config = CLIConfig.get_project_config()
722
+
723
+ datasources = await client.datasources()
724
+ for datasource in datasources:
725
+ await user_client.datasource_share(datasource["id"], config["id"], created_workspace["id"])
726
+
727
+
728
+ async def create_workspace_non_interactive(
729
+ ctx: Context, workspace_name: str, starterkit: str, user_token: str, fork: bool
730
+ ):
731
+ """Creates a workspace using the provided name and starterkit"""
732
+ client: TinyB = ctx.ensure_object(dict)["client"]
733
+
734
+ try:
735
+ user_client: TinyB = deepcopy(client)
736
+ user_client.token = user_token
737
+
738
+ created_workspace = await user_client.create_workspace(workspace_name, starterkit)
739
+ click.echo(FeedbackManager.success_workspace_created(workspace_name=workspace_name))
740
+
741
+ if fork:
742
+ await fork_workspace(client, user_client, created_workspace)
743
+
744
+ except Exception as e:
745
+ raise CLIWorkspaceException(FeedbackManager.error_exception(error=str(e)))
746
+
747
+
748
+ async def create_workspace_interactive(
749
+ ctx: Context, workspace_name: Optional[str], starterkit: Optional[str], user_token: str, fork: bool
750
+ ):
751
+ if not starterkit and not is_ci_environment():
752
+ click.echo("\n")
753
+ starterkit = await ask_for_starterkit_interactively(ctx)
754
+ if not starterkit: # Cancelled by user
755
+ return
756
+
757
+ if starterkit == "blank": # 'blank' == empty workspace
758
+ starterkit = None
759
+
760
+ if not workspace_name:
761
+ """Creates a workspace guiding the user"""
762
+ click.echo("\n")
763
+ click.echo(FeedbackManager.info_workspace_create_greeting())
764
+ default_name = f"new_workspace_{uuid.uuid4().hex[0:4]}"
765
+ workspace_name = click.prompt("\nWorkspace name", default=default_name, err=True, type=str)
766
+
767
+ await create_workspace_non_interactive(ctx, workspace_name, starterkit, user_token, fork) # type: ignore
768
+
769
+
770
+ async def create_workspace_branch(
771
+ branch_name: Optional[str],
772
+ last_partition: bool,
773
+ all: bool,
774
+ ignore_datasources: Optional[List[str]],
775
+ wait: Optional[bool],
776
+ ) -> None:
777
+ """
778
+ Creates a workspace branch
779
+ """
780
+ config = CLIConfig.get_project_config()
781
+ _ = await try_update_config_with_remote(config)
782
+
783
+ try:
784
+ workspace = await get_current_workspace(config)
785
+ if not workspace:
786
+ raise CLIWorkspaceException(FeedbackManager.error_workspace())
787
+
788
+ if not branch_name:
789
+ click.echo(FeedbackManager.info_workspace_branch_create_greeting())
790
+ default_name = f"{workspace['name']}_{uuid.uuid4().hex[0:4]}"
791
+ branch_name = click.prompt("\Branch name", default=default_name, err=True, type=str)
792
+ assert isinstance(branch_name, str)
793
+
794
+ response = await config.get_client().create_workspace_branch(
795
+ branch_name,
796
+ last_partition,
797
+ all,
798
+ ignore_datasources,
799
+ )
800
+ assert isinstance(response, dict)
801
+
802
+ is_job: bool = "job" in response
803
+ is_summary: bool = "partitions" in response
804
+
805
+ if not is_job and not is_summary:
806
+ raise CLIException(str(response))
807
+
808
+ if all and not is_job:
809
+ raise CLIException(str(response))
810
+
811
+ click.echo(
812
+ FeedbackManager.success_workspace_branch_created(workspace_name=workspace["name"], branch_name=branch_name)
813
+ )
814
+
815
+ job_id: Optional[str] = None
816
+
817
+ if is_job:
818
+ job_id = response["job"]["job_id"]
819
+ job_url = response["job"]["job_url"]
820
+ click.echo(FeedbackManager.info_data_branch_job_url(url=job_url))
821
+
822
+ if wait and is_job:
823
+ assert isinstance(job_id, str)
824
+
825
+ # Await the job to finish and get the result dict
826
+ job_response = await wait_job(config.get_client(), job_id, job_url, "Branch creation")
827
+ if job_response is None:
828
+ raise CLIException(f"Empty job API response (job_id: {job_id}, job_url: {job_url})")
829
+ else:
830
+ response = job_response.get("result", {})
831
+ is_summary = "partitions" in response
832
+
833
+ await switch_workspace(config, branch_name, only_environments=True)
834
+ if is_summary and (bool(last_partition) or bool(all)):
835
+ await print_data_branch_summary(config.get_client(), None, response)
836
+
837
+ except Exception as e:
838
+ raise CLIException(FeedbackManager.error_exception(error=str(e)))
839
+
840
+
841
+ async def print_data_branch_summary(client, job_id, response=None):
842
+ response = await client.job(job_id) if job_id else response or {"partitions": []}
843
+ columns = ["Data Source", "Partition", "Status", "Error"]
844
+ table = []
845
+ for partition in response["partitions"]:
846
+ for p in partition["partitions"]:
847
+ table.append([partition["datasource"]["name"], p["partition"], p["status"], p.get("error", "")])
848
+ echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
849
+
850
+
851
+ async def print_branch_regression_tests_summary(client, job_id, host, response=None):
852
+ def format_metric(metric: Union[str, float], is_percentage: bool = False) -> str:
853
+ if isinstance(metric, float):
854
+ if is_percentage:
855
+ return f"{round(metric, 3):+} %"
856
+ else:
857
+ return f"{round(metric, 3)} seconds"
858
+ else:
859
+ return metric
860
+
861
+ failed = False
862
+ response = await client.job(job_id) if job_id else response or {"progress": []}
863
+ output = "\n"
864
+ for step in response["progress"]:
865
+ run = step["run"]
866
+ if run["output"]:
867
+ # If the output contains an alert emoji, it means that it should be ou
868
+ output += (
869
+ warning_message(run["output"])()
870
+ if isinstance(run["output"], str) and "🚨" in run["output"]
871
+ else "".join(run["output"])
872
+ )
873
+ if not run["was_successfull"]:
874
+ failed = True
875
+ click.echo(output)
876
+
877
+ if failed:
878
+ click.echo("")
879
+ click.echo("")
880
+ click.echo("==== Failures Detail ====")
881
+ click.echo("")
882
+ for step in response["progress"]:
883
+ if not step["run"]["was_successfull"]:
884
+ for failure in step["run"]["failed"]:
885
+ try:
886
+ click.echo(f"❌ {failure['name']}")
887
+ click.echo(FeedbackManager.error_branch_check_pipe(error=failure["error"]))
888
+ click.echo("")
889
+ except Exception:
890
+ pass
891
+
892
+ click.echo("")
893
+ click.echo("")
894
+ click.echo("==== Performance metrics ====")
895
+ click.echo("")
896
+ for step in response["progress"]:
897
+ run = step["run"]
898
+ if run.get("metrics_summary") and run.get("metrics_timing"):
899
+ column_names = [f"{run['pipe_name']}({run['test_type']})", "Origin", "Branch", "Delta"]
900
+
901
+ click.echo(
902
+ format_pretty_table(
903
+ [
904
+ [
905
+ metric,
906
+ format_metric(run["metrics_timing"][metric][0]),
907
+ format_metric(run["metrics_timing"][metric][1]),
908
+ format_metric(run["metrics_timing"][metric][2], is_percentage=True),
909
+ ]
910
+ for metric in [
911
+ "min response time",
912
+ "max response time",
913
+ "mean response time",
914
+ "median response time",
915
+ "p90 response time",
916
+ "min read bytes",
917
+ "max read bytes",
918
+ "mean read bytes",
919
+ "median read bytes",
920
+ "p90 read bytes",
921
+ ]
922
+ ],
923
+ column_names=column_names,
924
+ )
925
+ )
926
+
927
+ click.echo("")
928
+ click.echo("")
929
+ click.echo("==== Results Summary ====")
930
+ click.echo("")
931
+ click.echo(
932
+ format_pretty_table(
933
+ [
934
+ [
935
+ step["run"]["pipe_name"],
936
+ step["run"]["test_type"],
937
+ step["run"]["metrics_summary"].get("run", 0),
938
+ step["run"]["metrics_summary"].get("passed", 0),
939
+ step["run"]["metrics_summary"].get("failed", 0),
940
+ format_metric(
941
+ (
942
+ step["run"]["metrics_timing"]["mean response time"][2]
943
+ if "mean response time" in step["run"]["metrics_timing"]
944
+ else 0.0
945
+ ),
946
+ is_percentage=True,
947
+ ),
948
+ format_metric(
949
+ (
950
+ step["run"]["metrics_timing"]["mean read bytes"][2]
951
+ if "mean read bytes" in step["run"]["metrics_timing"]
952
+ else 0.0
953
+ ),
954
+ is_percentage=True,
955
+ ),
956
+ ]
957
+ for step in response["progress"]
958
+ ],
959
+ column_names=["Endpoint", "Test", "Run", "Passed", "Failed", "Mean response time", "Mean read bytes"],
960
+ )
961
+ )
962
+ click.echo("")
963
+ if failed:
964
+ for step in response["progress"]:
965
+ if not step["run"]["was_successfull"]:
966
+ for failure in step["run"]["failed"]:
967
+ click.echo(f"❌ FAILED {failure['name']}\n")
968
+ if failed:
969
+ raise CLIException(
970
+ "Check Failures Detail above for more information. If the results are expected, skip asserts or increase thresholds, see 💡 Hints above (note skip asserts flags are applied to all regression tests, so use them when it makes sense).\n\nIf you are using the CI template for GitHub or GitLab you can add skip asserts flags as labels to the MR and they are automatically applied. Find available flags to skip asserts and thresholds here => https://www.tinybird.co/docs/production/implementing-test-strategies.html#fixture-tests"
971
+ )
972
+
973
+
974
+ class PlanName(Enum):
975
+ DEV = "Build"
976
+ PRO = "Pro"
977
+ ENTERPRISE = "Enterprise"
978
+
979
+
980
+ def _get_workspace_plan_name(plan):
981
+ """
982
+ >>> _get_workspace_plan_name("dev")
983
+ 'Build'
984
+ >>> _get_workspace_plan_name("pro")
985
+ 'Pro'
986
+ >>> _get_workspace_plan_name("enterprise")
987
+ 'Enterprise'
988
+ >>> _get_workspace_plan_name("branch_enterprise")
989
+ 'Enterprise'
990
+ >>> _get_workspace_plan_name("other_plan")
991
+ 'Custom'
992
+ """
993
+ if plan == "dev":
994
+ return PlanName.DEV.value
995
+ if plan == "pro":
996
+ return PlanName.PRO.value
997
+ if plan in ("enterprise", "branch_enterprise"):
998
+ return PlanName.ENTERPRISE.value
999
+ return "Custom"
1000
+
1001
+
1002
+ def get_format_from_filename_or_url(filename_or_url: str) -> str:
1003
+ """
1004
+ >>> get_format_from_filename_or_url('wadus_parquet.csv')
1005
+ 'csv'
1006
+ >>> get_format_from_filename_or_url('wadus_csv.parquet')
1007
+ 'parquet'
1008
+ >>> get_format_from_filename_or_url('wadus_csv.ndjson')
1009
+ 'ndjson'
1010
+ >>> get_format_from_filename_or_url('wadus_csv.json')
1011
+ 'ndjson'
1012
+ >>> get_format_from_filename_or_url('wadus_parquet.csv?auth=pepe')
1013
+ 'csv'
1014
+ >>> get_format_from_filename_or_url('wadus_csv.parquet?auth=pepe')
1015
+ 'parquet'
1016
+ >>> get_format_from_filename_or_url('wadus_parquet.ndjson?auth=pepe')
1017
+ 'ndjson'
1018
+ >>> get_format_from_filename_or_url('wadus.json?auth=pepe')
1019
+ 'ndjson'
1020
+ >>> get_format_from_filename_or_url('wadus_csv_')
1021
+ 'csv'
1022
+ >>> get_format_from_filename_or_url('wadus_json_csv_')
1023
+ 'csv'
1024
+ >>> get_format_from_filename_or_url('wadus_json_')
1025
+ 'ndjson'
1026
+ >>> get_format_from_filename_or_url('wadus_ndjson_')
1027
+ 'ndjson'
1028
+ >>> get_format_from_filename_or_url('wadus_parquet_')
1029
+ 'parquet'
1030
+ >>> get_format_from_filename_or_url('wadus')
1031
+ 'csv'
1032
+ >>> get_format_from_filename_or_url('https://storage.googleapis.com/tinybird-waduscom/stores_stock__v2_1646741850424_final.csv?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=44444444444-compute@developer.gserviceaccount.com/1234/auto/storage/goog4_request&X-Goog-Date=20220308T121750Z&X-Goog-Expires=86400&X-Goog-SignedHeaders=host&X-Goog-Signature=8888888888888888888888888888888888888888888888888888888')
1033
+ 'csv'
1034
+ """
1035
+ filename_or_url = filename_or_url.lower()
1036
+ if filename_or_url.endswith("json") or filename_or_url.endswith("ndjson"):
1037
+ return "ndjson"
1038
+ if filename_or_url.endswith("parquet"):
1039
+ return "parquet"
1040
+ if filename_or_url.endswith("csv"):
1041
+ return "csv"
1042
+ try:
1043
+ parsed = urlparse(filename_or_url)
1044
+ if parsed.path.endswith("json") or parsed.path.endswith("ndjson"):
1045
+ return "ndjson"
1046
+ if parsed.path.endswith("parquet"):
1047
+ return "parquet"
1048
+ if parsed.path.endswith("csv"):
1049
+ return "csv"
1050
+ except Exception:
1051
+ pass
1052
+ if "csv" in filename_or_url:
1053
+ return "csv"
1054
+ if "json" in filename_or_url:
1055
+ return "ndjson"
1056
+ if "parquet" in filename_or_url:
1057
+ return "parquet"
1058
+ return "csv"
1059
+
1060
+
1061
+ async def push_data(
1062
+ ctx: Context,
1063
+ client: TinyB,
1064
+ datasource_name: str,
1065
+ url,
1066
+ connector: Optional[str],
1067
+ sql: Optional[str],
1068
+ mode: str = "append",
1069
+ sql_condition: Optional[str] = None,
1070
+ replace_options=None,
1071
+ ignore_empty: bool = False,
1072
+ concurrency: int = 1,
1073
+ ):
1074
+ if url and type(url) is tuple:
1075
+ url = url[0]
1076
+
1077
+ if connector and sql:
1078
+ load_connector_config(ctx, connector, False, check_uninstalled=False)
1079
+ if connector not in ctx.obj:
1080
+ raise CLIException(FeedbackManager.error_connector_not_configured(connector=connector))
1081
+ else:
1082
+ _connector: "Connector" = ctx.obj[connector]
1083
+ click.echo(FeedbackManager.info_starting_export_process(connector=connector))
1084
+ try:
1085
+ url = _connector.export_to_gcs(sql, datasource_name, mode)
1086
+ except ConnectorNothingToLoad as e:
1087
+ if ignore_empty:
1088
+ click.echo(str(e))
1089
+ return
1090
+ else:
1091
+ raise e
1092
+
1093
+ def cb(res):
1094
+ if cb.First: # type: ignore[attr-defined]
1095
+ blocks_to_process = len([x for x in res["block_log"] if x["status"] == "idle"])
1096
+ if blocks_to_process:
1097
+ cb.bar = click.progressbar(label=FeedbackManager.info_progress_blocks(), length=blocks_to_process) # type: ignore[attr-defined]
1098
+ cb.bar.update(0) # type: ignore[attr-defined]
1099
+ cb.First = False # type: ignore[attr-defined]
1100
+ cb.blocks_to_process = blocks_to_process # type: ignore[attr-defined]
1101
+ else:
1102
+ done = len([x for x in res["block_log"] if x["status"] == "done"])
1103
+ if done * 2 > cb.blocks_to_process: # type: ignore[attr-defined]
1104
+ cb.bar.label = FeedbackManager.info_progress_current_blocks() # type: ignore[attr-defined]
1105
+ cb.bar.update(done - cb.prev_done) # type: ignore[attr-defined]
1106
+ cb.prev_done = done # type: ignore[attr-defined]
1107
+
1108
+ cb.First = True # type: ignore[attr-defined]
1109
+ cb.prev_done = 0 # type: ignore[attr-defined]
1110
+
1111
+ click.echo(FeedbackManager.gray(message=f"\nImporting data to {datasource_name} Data Source..."))
1112
+
1113
+ if isinstance(url, list):
1114
+ urls = url
1115
+ else:
1116
+ urls = [url]
1117
+
1118
+ async def process_url(
1119
+ datasource_name: str, url: str, mode: str, sql_condition: Optional[str], replace_options: Optional[Set[str]]
1120
+ ):
1121
+ parsed = urlparse(url)
1122
+ # poor man's format detection
1123
+ _format = get_format_from_filename_or_url(url)
1124
+ if parsed.scheme in ("http", "https"):
1125
+ res = await client.datasource_create_from_url(
1126
+ datasource_name,
1127
+ url,
1128
+ mode=mode,
1129
+ status_callback=cb,
1130
+ sql_condition=sql_condition,
1131
+ format=_format,
1132
+ replace_options=replace_options,
1133
+ )
1134
+ else:
1135
+ res = await client.datasource_append_data(
1136
+ datasource_name,
1137
+ file=url,
1138
+ mode=mode,
1139
+ sql_condition=sql_condition,
1140
+ format=_format,
1141
+ replace_options=replace_options,
1142
+ )
1143
+
1144
+ datasource_name = res["datasource"]["name"]
1145
+ try:
1146
+ datasource = await client.get_datasource(datasource_name)
1147
+ except DoesNotExistException:
1148
+ raise CLIException(FeedbackManager.error_datasource_does_not_exist(datasource=datasource_name))
1149
+ except Exception as e:
1150
+ raise CLIException(FeedbackManager.error_exception(error=str(e)))
1151
+
1152
+ total_rows = (datasource.get("statistics", {}) or {}).get("row_count", 0)
1153
+ appended_rows = 0
1154
+ parser = None
1155
+
1156
+ if res.get("error"):
1157
+ raise CLIException(FeedbackManager.error_exception(error=res["error"]))
1158
+ if res.get("errors"):
1159
+ raise CLIException(FeedbackManager.error_exception(error=res["errors"]))
1160
+ if res.get("blocks"):
1161
+ for block in res["blocks"]:
1162
+ if "process_return" in block and block["process_return"] is not None:
1163
+ process_return = block["process_return"][0]
1164
+ parser = process_return["parser"] if process_return.get("parser") else parser
1165
+ if parser and parser != "clickhouse":
1166
+ parser = process_return["parser"]
1167
+ appended_rows += process_return["lines"]
1168
+
1169
+ return parser, total_rows, appended_rows
1170
+
1171
+ try:
1172
+ tasks = [process_url(datasource_name, url, mode, sql_condition, replace_options) for url in urls]
1173
+ output = await gather_with_concurrency(concurrency, *tasks)
1174
+ parser, total_rows, appended_rows = list(output)[-1]
1175
+ except AuthNoTokenException:
1176
+ raise
1177
+ except OperationCanNotBePerformed as e:
1178
+ raise CLIException(FeedbackManager.error_operation_can_not_be_performed(error=e))
1179
+ except Exception as e:
1180
+ raise CLIException(FeedbackManager.error_exception(error=e))
1181
+ else:
1182
+ if mode == "append" and parser and parser != "clickhouse":
1183
+ click.echo(FeedbackManager.success_appended_rows(appended_rows=appended_rows))
1184
+
1185
+ if mode == "replace":
1186
+ click.echo(FeedbackManager.success_replaced_datasource(datasource=datasource_name))
1187
+ else:
1188
+ click.echo(FeedbackManager.highlight(message="» 2.57m rows x 9 cols in 852.04ms"))
1189
+
1190
+ click.echo(FeedbackManager.success_progress_blocks())
1191
+
1192
+ finally:
1193
+ try:
1194
+ for url in urls:
1195
+ _connector.clean(urlparse(url).path.split("/")[-1])
1196
+ except Exception:
1197
+ pass
1198
+
1199
+
1200
+ async def sync_data(ctx, datasource_name: str, yes: bool):
1201
+ client: TinyB = ctx.obj["client"]
1202
+ datasource = await client.get_datasource(datasource_name)
1203
+
1204
+ VALID_DATASOURCES = ["bigquery", "snowflake", "s3", "gcs"]
1205
+ if datasource["type"] not in VALID_DATASOURCES:
1206
+ raise CLIException(FeedbackManager.error_sync_not_supported(valid_datasources=VALID_DATASOURCES))
1207
+
1208
+ warning_message = (
1209
+ FeedbackManager.warning_datasource_sync_bucket(datasource=datasource_name)
1210
+ if datasource["type"] in ["s3", "gcs"]
1211
+ else FeedbackManager.warning_datasource_sync(
1212
+ datasource=datasource_name,
1213
+ )
1214
+ )
1215
+ if yes or click.confirm(warning_message):
1216
+ await client.datasource_sync(datasource["id"])
1217
+ click.echo(FeedbackManager.success_sync_datasource(datasource=datasource_name))
1218
+
1219
+
1220
+ # eval "$(_TB_COMPLETE=source_bash tb)"
1221
+ def autocomplete_topics(ctx: Context, args, incomplete):
1222
+ try:
1223
+ config = async_to_sync(get_config)(None, None)
1224
+ ctx.ensure_object(dict)["config"] = config
1225
+ client = create_tb_client(ctx)
1226
+ topics = async_to_sync(client.kafka_list_topics)(args[2])
1227
+ return [t for t in topics if incomplete in t]
1228
+ except Exception:
1229
+ return []
1230
+
1231
+
1232
+ def validate_datasource_name(name):
1233
+ if not isinstance(name, str) or name == "":
1234
+ raise CLIException(FeedbackManager.error_datasource_name())
1235
+
1236
+
1237
+ def validate_connection_id(connection_id):
1238
+ if not isinstance(connection_id, str) or connection_id == "":
1239
+ raise CLIException(FeedbackManager.error_datasource_connection_id())
1240
+
1241
+
1242
+ def validate_kafka_topic(topic):
1243
+ if not isinstance(topic, str):
1244
+ raise CLIException(FeedbackManager.error_kafka_topic())
1245
+
1246
+
1247
+ def validate_kafka_group(group):
1248
+ if not isinstance(group, str):
1249
+ raise CLIException(FeedbackManager.error_kafka_group())
1250
+
1251
+
1252
+ def validate_kafka_auto_offset_reset(auto_offset_reset):
1253
+ valid_values = {"latest", "earliest", "none"}
1254
+ if auto_offset_reset not in valid_values:
1255
+ raise CLIException(FeedbackManager.error_kafka_auto_offset_reset())
1256
+
1257
+
1258
+ def validate_kafka_schema_registry_url(schema_registry_url):
1259
+ if not is_url_valid(schema_registry_url):
1260
+ raise CLIException(FeedbackManager.error_kafka_registry())
1261
+
1262
+
1263
+ def is_url_valid(url):
1264
+ try:
1265
+ result = urlparse(url)
1266
+ return all([result.scheme, result.netloc])
1267
+ except Exception:
1268
+ return False
1269
+
1270
+
1271
+ def validate_kafka_bootstrap_servers(host_and_port):
1272
+ if not isinstance(host_and_port, str):
1273
+ raise CLIException(FeedbackManager.error_kafka_bootstrap_server())
1274
+ parts = host_and_port.split(":")
1275
+ if len(parts) > 2:
1276
+ raise CLIException(FeedbackManager.error_kafka_bootstrap_server())
1277
+ host = parts[0]
1278
+ port_str = parts[1] if len(parts) == 2 else "9092"
1279
+ try:
1280
+ port = int(port_str)
1281
+ except Exception:
1282
+ raise CLIException(FeedbackManager.error_kafka_bootstrap_server())
1283
+ with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
1284
+ try:
1285
+ sock.settimeout(3)
1286
+ sock.connect((host, port))
1287
+ except socket.timeout:
1288
+ raise CLIException(FeedbackManager.error_kafka_bootstrap_server_conn_timeout())
1289
+ except Exception:
1290
+ raise CLIException(FeedbackManager.error_kafka_bootstrap_server_conn())
1291
+
1292
+
1293
+ def validate_kafka_key(s):
1294
+ if not isinstance(s, str):
1295
+ raise CLIException("Key format is not correct, it should be a string")
1296
+
1297
+
1298
+ def validate_kafka_secret(s):
1299
+ if not isinstance(s, str):
1300
+ raise CLIException("Password format is not correct, it should be a string")
1301
+
1302
+
1303
+ def validate_string_connector_param(param, s):
1304
+ if not isinstance(s, str):
1305
+ raise CLIConnectionException(param + " format is not correct, it should be a string")
1306
+
1307
+
1308
+ async def validate_connection_name(client, connection_name, service):
1309
+ if await client.get_connector(connection_name, service) is not None:
1310
+ raise CLIConnectionException(FeedbackManager.error_connection_already_exists(name=connection_name))
1311
+
1312
+
1313
+ def _get_setting_value(connection, setting, sensitive_settings):
1314
+ if setting in sensitive_settings:
1315
+ return "*****"
1316
+ return connection.get(setting, "")
1317
+
1318
+
1319
+ async def switch_workspace(config: CLIConfig, workspace_name_or_id: str, only_environments: bool = False) -> None:
1320
+ try:
1321
+ if only_environments:
1322
+ workspaces = await get_current_workspace_branches(config)
1323
+ else:
1324
+ response = await config.get_client().user_workspaces()
1325
+ workspaces = response["workspaces"]
1326
+
1327
+ workspace = next(
1328
+ (
1329
+ workspace
1330
+ for workspace in workspaces
1331
+ if workspace["name"] == workspace_name_or_id or workspace["id"] == workspace_name_or_id
1332
+ ),
1333
+ None,
1334
+ )
1335
+
1336
+ if not workspace:
1337
+ if only_environments:
1338
+ raise CLIException(FeedbackManager.error_branch(branch=workspace_name_or_id))
1339
+ else:
1340
+ raise CLIException(FeedbackManager.error_workspace(workspace=workspace_name_or_id))
1341
+
1342
+ config.set_token(workspace["token"])
1343
+ config.set_token_for_host(workspace["token"], config.get_host())
1344
+ _ = await try_update_config_with_remote(config)
1345
+
1346
+ # Set the id and name afterwards.
1347
+ # When working with branches the call to try_update_config_with_remote above
1348
+ # sets the data with the main branch ones
1349
+ config["id"] = workspace["id"]
1350
+ config["name"] = workspace["name"]
1351
+
1352
+ config.persist_to_file()
1353
+
1354
+ click.echo(FeedbackManager.success_now_using_config(name=config["name"], id=config["id"]))
1355
+ except AuthNoTokenException:
1356
+ raise
1357
+ except CLIException:
1358
+ raise
1359
+ except Exception as e:
1360
+ raise CLIException(FeedbackManager.error_exception(error=str(e)))
1361
+
1362
+
1363
+ async def switch_to_workspace_by_user_workspace_data(config: CLIConfig, user_workspace_data: Dict[str, Any]):
1364
+ try:
1365
+ config["id"] = user_workspace_data["id"]
1366
+ config["name"] = user_workspace_data["name"]
1367
+ config.set_token(user_workspace_data["token"])
1368
+ config.set_token_for_host(user_workspace_data["token"], config.get_host())
1369
+ config.persist_to_file()
1370
+
1371
+ click.echo(FeedbackManager.success_now_using_config(name=config["name"], id=config["id"]))
1372
+ except Exception as e:
1373
+ raise CLIException(FeedbackManager.error_exception(error=str(e)))
1374
+
1375
+
1376
+ async def print_current_workspace(config: CLIConfig) -> None:
1377
+ _ = await try_update_config_with_remote(config, only_if_needed=True)
1378
+
1379
+ current_main_workspace = await get_current_main_workspace(config)
1380
+ assert isinstance(current_main_workspace, dict)
1381
+
1382
+ columns = ["name", "id", "role", "plan", "current"]
1383
+
1384
+ table = [
1385
+ (
1386
+ current_main_workspace["name"],
1387
+ current_main_workspace["id"],
1388
+ current_main_workspace["role"],
1389
+ _get_workspace_plan_name(current_main_workspace["plan"]),
1390
+ True,
1391
+ )
1392
+ ]
1393
+
1394
+ click.echo(FeedbackManager.info_current_workspace())
1395
+ echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
1396
+
1397
+
1398
+ async def print_current_branch(config: CLIConfig) -> None:
1399
+ _ = await try_update_config_with_remote(config, only_if_needed=True)
1400
+
1401
+ response = await config.get_client().user_workspaces_and_branches()
1402
+
1403
+ columns = ["name", "id", "workspace"]
1404
+ table = []
1405
+
1406
+ for workspace in response["workspaces"]:
1407
+ if config["id"] == workspace["id"]:
1408
+ click.echo(FeedbackManager.info_current_branch())
1409
+ if workspace.get("is_branch"):
1410
+ name = workspace["name"]
1411
+ main_workspace = await get_current_main_workspace(config)
1412
+ assert isinstance(main_workspace, dict)
1413
+ main_name = main_workspace["name"]
1414
+ else:
1415
+ name = MAIN_BRANCH
1416
+ main_name = workspace["name"]
1417
+ table.append([name, workspace["id"], main_name])
1418
+ break
1419
+
1420
+ echo_safe_humanfriendly_tables_format_smart_table(table, column_names=columns)
1421
+
1422
+
1423
+ class ConnectionReplacements:
1424
+ _PARAMS_REPLACEMENTS: Dict[str, Dict[str, str]] = {
1425
+ "s3": {
1426
+ "service": "service",
1427
+ "connection_name": "name",
1428
+ "key": "s3_access_key_id",
1429
+ "secret": "s3_secret_access_key",
1430
+ "region": "s3_region",
1431
+ },
1432
+ "s3_iamrole": {
1433
+ "service": "service",
1434
+ "connection_name": "name",
1435
+ "role_arn": "s3_iamrole_arn",
1436
+ "region": "s3_iamrole_region",
1437
+ },
1438
+ "gcs_hmac": {
1439
+ "service": "service",
1440
+ "connection_name": "name",
1441
+ "key": "gcs_hmac_access_id",
1442
+ "secret": "gcs_hmac_secret",
1443
+ "region": "gcs_region",
1444
+ },
1445
+ "gcs": {
1446
+ "project_id": "gcs_project_id",
1447
+ "client_id": "gcs_client_id",
1448
+ "client_email": "gcs_client_email",
1449
+ "client_x509_cert_url": "gcs_client_x509_cert_url",
1450
+ "private_key": "gcs_private_key",
1451
+ "private_key_id": "gcs_private_key_id",
1452
+ "connection_name": "name",
1453
+ },
1454
+ "dynamodb": {
1455
+ "service": "service",
1456
+ "connection_name": "name",
1457
+ "role_arn": "dynamodb_iamrole_arn",
1458
+ "region": "dynamodb_iamrole_region",
1459
+ },
1460
+ }
1461
+
1462
+ @staticmethod
1463
+ def map_api_params_from_prompt_params(service: str, **params: Any) -> Dict[str, Any]:
1464
+ """Maps prompt parameters to API parameters."""
1465
+
1466
+ api_params = {}
1467
+ for key in params.keys():
1468
+ try:
1469
+ api_params[ConnectionReplacements._PARAMS_REPLACEMENTS[service][key]] = params[key]
1470
+ except KeyError:
1471
+ api_params[key] = params[key]
1472
+
1473
+ api_params["service"] = service
1474
+ return api_params
1475
+
1476
+
1477
+ # ======
1478
+ # Temporal new functions while we fully merge the new CLIConfig
1479
+ # ======
1480
+
1481
+
1482
+ async def get_host_from_region(
1483
+ config: CLIConfig, region_name_or_host_or_id: str, host: Optional[str] = None
1484
+ ) -> Tuple[List[Region], str]:
1485
+ regions: List[Region]
1486
+ region: Optional[Region]
1487
+
1488
+ host = host or config.get_host(use_defaults_if_needed=True)
1489
+
1490
+ try:
1491
+ regions = await get_regions(config)
1492
+ assert isinstance(regions, list)
1493
+ except Exception:
1494
+ regions = []
1495
+
1496
+ if not regions:
1497
+ assert isinstance(host, str)
1498
+ click.echo(f"No regions available, using host: {host}")
1499
+ return [], host
1500
+
1501
+ try:
1502
+ index = int(region_name_or_host_or_id)
1503
+ try:
1504
+ host = regions[index - 1]["api_host"]
1505
+ except Exception:
1506
+ raise CLIException(FeedbackManager.error_getting_region_by_index())
1507
+ except ValueError:
1508
+ region_name = region_name_or_host_or_id.lower()
1509
+ try:
1510
+ region = get_region_from_host(region_name, regions)
1511
+ host = region["api_host"] if region else None
1512
+ except Exception:
1513
+ raise CLIException(FeedbackManager.error_getting_region_by_name_or_url())
1514
+
1515
+ if not host:
1516
+ raise CLIException(FeedbackManager.error_getting_region_by_name_or_url())
1517
+
1518
+ return regions, host
1519
+
1520
+
1521
+ async def get_regions(config: CLIConfig) -> List[Region]:
1522
+ regions: List[Region] = []
1523
+
1524
+ try:
1525
+ response = await config.get_client().regions()
1526
+ regions = response.get("regions", [])
1527
+ except Exception:
1528
+ pass
1529
+
1530
+ try:
1531
+ if "tokens" not in config:
1532
+ return regions
1533
+
1534
+ for key in config["tokens"]:
1535
+ region = next((region for region in regions if key == region["api_host"] or key == region["host"]), None)
1536
+ if region:
1537
+ region["default_password"] = config["tokens"][key]
1538
+ region["provider"] = region["provider"] or ""
1539
+ else:
1540
+ regions.append(
1541
+ {
1542
+ "api_host": format_host(key, subdomain="api"),
1543
+ "host": get_display_host(key),
1544
+ "name": key,
1545
+ "default_password": config["tokens"][key],
1546
+ "provider": "",
1547
+ }
1548
+ )
1549
+
1550
+ except Exception:
1551
+ pass
1552
+
1553
+ return regions
1554
+
1555
+
1556
+ def get_region_from_host(region_name_or_host: str, regions: List[Region]) -> Optional[Region]:
1557
+ """Returns the region that matches region_name_or_host by name, API host or ui host"""
1558
+ for region in regions:
1559
+ if region_name_or_host in (region["name"].lower(), region["host"], region["api_host"]):
1560
+ return region
1561
+ return None
1562
+
1563
+
1564
+ async def try_update_config_with_remote(
1565
+ config: CLIConfig, raise_on_errors: bool = True, only_if_needed: bool = False, auto_persist: bool = True
1566
+ ) -> bool:
1567
+ response: Dict[str, Any]
1568
+
1569
+ if not config.get_token():
1570
+ if not raise_on_errors:
1571
+ return False
1572
+ raise AuthNoTokenException()
1573
+
1574
+ if "id" in config and only_if_needed:
1575
+ return True
1576
+
1577
+ try:
1578
+ response = await config.get_client().workspace_info()
1579
+ except AuthException:
1580
+ if raise_on_errors:
1581
+ raise CLIAuthException(FeedbackManager.error_invalid_token_for_host(host=config.get_host()))
1582
+ return False
1583
+ except Exception as ex:
1584
+ if raise_on_errors:
1585
+ ex_message = str(ex)
1586
+ if "cannot parse" in ex_message.lower():
1587
+ raise CLIAuthException(FeedbackManager.error_invalid_host(host=config.get_host()))
1588
+
1589
+ raise CLIAuthException(FeedbackManager.error_exception(error=ex_message))
1590
+ return False
1591
+
1592
+ for k in ("id", "name", "user_email", "user_id", "scope"):
1593
+ if k in response:
1594
+ config[k] = response[k]
1595
+
1596
+ config.set_token_for_host(config.get_token(), config.get_host())
1597
+
1598
+ if auto_persist:
1599
+ config.persist_to_file()
1600
+
1601
+ return True
1602
+
1603
+
1604
+ def ask_for_admin_token_interactively(ui_host: str, default_token: Optional[str]) -> str:
1605
+ return (
1606
+ click.prompt(
1607
+ f"\nCopy the \"admin your@email\" token from {ui_host}/tokens and paste it here { 'OR press enter to use the token from .tinyb file' if default_token else ''}",
1608
+ hide_input=True,
1609
+ show_default=False,
1610
+ default=default_token,
1611
+ type=str,
1612
+ )
1613
+ or ""
1614
+ )
1615
+
1616
+
1617
+ async def try_authenticate(
1618
+ config: CLIConfig,
1619
+ regions: Optional[List[Region]] = None,
1620
+ interactive: bool = False,
1621
+ try_all_regions: bool = False,
1622
+ ) -> bool:
1623
+ host: Optional[str] = config.get_host()
1624
+
1625
+ if not regions and interactive:
1626
+ regions = await get_regions(config)
1627
+
1628
+ selected_region: Optional[Region] = None
1629
+ default_password: Optional[str] = None
1630
+
1631
+ if regions:
1632
+ if interactive:
1633
+ selected_region = ask_for_region_interactively(regions)
1634
+ if selected_region is None:
1635
+ return False
1636
+
1637
+ host = selected_region.get("api_host")
1638
+ default_password = selected_region.get("default_password")
1639
+ else:
1640
+ assert isinstance(host, str)
1641
+ selected_region = get_region_from_host(host, regions)
1642
+
1643
+ name: str
1644
+ api_host: str
1645
+ ui_host: str
1646
+ token: Optional[str]
1647
+ if host and not selected_region:
1648
+ name, api_host, ui_host = (host, format_host(host, subdomain="api"), format_host(host, subdomain="ui"))
1649
+ token = config.get_token()
1650
+ else:
1651
+ name, api_host, ui_host = get_region_info(config, selected_region)
1652
+ token = config.get_token_for_host(api_host)
1653
+ config.set_host(api_host)
1654
+
1655
+ if not token:
1656
+ token = ask_for_admin_token_interactively(get_display_host(ui_host), default_token=default_password)
1657
+ config.set_token(token)
1658
+
1659
+ add_telemetry_event("auth_token", token=token)
1660
+ authenticated: bool = await try_update_config_with_remote(config, raise_on_errors=not try_all_regions)
1661
+
1662
+ # No luck? Let's try auth in all other regions
1663
+ if not authenticated and try_all_regions and not interactive:
1664
+ if not regions:
1665
+ regions = await get_regions(config)
1666
+
1667
+ # Check other regions, ignoring the previously tested region
1668
+ for region in [r for r in regions if r is not selected_region]:
1669
+ name, host, ui_host = get_region_info(config, region)
1670
+ config.set_host(host)
1671
+ authenticated = await try_update_config_with_remote(config, raise_on_errors=False)
1672
+ if authenticated:
1673
+ click.echo(FeedbackManager.success_using_host(name=name, host=get_display_host(ui_host)))
1674
+ break
1675
+
1676
+ if not authenticated:
1677
+ raise CLIAuthException(FeedbackManager.error_invalid_token())
1678
+
1679
+ config.persist_to_file()
1680
+
1681
+ click.echo(FeedbackManager.success_auth())
1682
+ click.echo(FeedbackManager.success_remember_api_host(api_host=host))
1683
+
1684
+ if not config.get("scope"):
1685
+ click.echo(FeedbackManager.warning_token_scope())
1686
+
1687
+ add_telemetry_event("auth_success")
1688
+
1689
+ return True
1690
+
1691
+
1692
+ async def wait_job(
1693
+ tb_client: TinyB,
1694
+ job_id: str,
1695
+ job_url: str,
1696
+ label: str,
1697
+ wait_observer: Optional[Callable[[Dict[str, Any], ProgressBar], None]] = None,
1698
+ ) -> Dict[str, Any]:
1699
+ progress_bar: ProgressBar
1700
+ with click.progressbar(
1701
+ label=f"{label} ",
1702
+ length=100,
1703
+ show_eta=False,
1704
+ show_percent=wait_observer is None,
1705
+ fill_char=click.style("█", fg="green"),
1706
+ ) as progress_bar:
1707
+
1708
+ def progressbar_cb(res: Dict[str, Any]):
1709
+ if wait_observer:
1710
+ wait_observer(res, progress_bar)
1711
+ return
1712
+
1713
+ if "progress_percentage" in res:
1714
+ progress_bar.update(int(round(res["progress_percentage"])) - progress_bar.pos)
1715
+ elif res["status"] != "working":
1716
+ progress_bar.update(progress_bar.length if progress_bar.length else 0)
1717
+
1718
+ try:
1719
+ # TODO: Simplify this as it's not needed to use two functions for
1720
+ result = await wait_job_no_ui(tb_client, job_id, progressbar_cb)
1721
+ if result["status"] != "done":
1722
+ raise CLIException(FeedbackManager.error_while_running_job(error=result["error"]))
1723
+ return result
1724
+ except asyncio.TimeoutError:
1725
+ raise CLIException(FeedbackManager.error_while_running_job(error="Reach timeout, job cancelled"))
1726
+ except JobException as e:
1727
+ raise CLIException(FeedbackManager.error_while_running_job(error=str(e)))
1728
+ except Exception as e:
1729
+ raise CLIException(FeedbackManager.error_getting_job_info(error=str(e), url=job_url))
1730
+
1731
+
1732
+ async def wait_job_no_ui(
1733
+ tb_client: TinyB,
1734
+ job_id: str,
1735
+ status_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
1736
+ ) -> Dict[str, Any]:
1737
+ try:
1738
+ result = await asyncio.wait_for(tb_client.wait_for_job(job_id, status_callback=status_callback), None)
1739
+ if result["status"] != "done":
1740
+ raise JobException(result.get("error"))
1741
+ return result
1742
+ except asyncio.TimeoutError:
1743
+ await tb_client.job_cancel(job_id)
1744
+ raise
1745
+
1746
+
1747
+ async def get_current_main_workspace(config: CLIConfig) -> Optional[Dict[str, Any]]:
1748
+ current_workspace = await config.get_client().user_workspaces_and_branches()
1749
+ return _get_current_main_workspace_common(current_workspace, config.get("id", current_workspace["id"]))
1750
+
1751
+
1752
+ def _get_current_main_workspace_common(
1753
+ user_workspace_and_branches: Dict[str, Any], current_workspace_id: str
1754
+ ) -> Optional[Dict[str, Any]]:
1755
+ def get_workspace_by_id(workspaces: List[Dict[str, Any]], id: str) -> Optional[Dict[str, Any]]:
1756
+ return next((ws for ws in workspaces if ws["id"] == id), None)
1757
+
1758
+ workspaces: Optional[List[Dict[str, Any]]] = user_workspace_and_branches.get("workspaces")
1759
+ if not workspaces:
1760
+ return None
1761
+
1762
+ current: Optional[Dict[str, Any]] = get_workspace_by_id(workspaces, current_workspace_id)
1763
+ if current and current.get("is_branch"):
1764
+ current = get_workspace_by_id(workspaces, current["main"])
1765
+
1766
+ return current
1767
+
1768
+
1769
+ def is_post_semver(new_version: Version, current_version: Version) -> bool:
1770
+ """
1771
+ Check if only the post part of the semantic version has changed.
1772
+
1773
+ Args:
1774
+ new_version (Version): The new version to check.
1775
+ current_version (Version): The current version to compare with.
1776
+
1777
+ Returns:
1778
+ bool: True if only the post part of the version has changed, False otherwise.
1779
+
1780
+ Examples:
1781
+ >>> is_post_semver(Version("0.0.0-2"), Version("0.0.0-1"))
1782
+ True
1783
+ >>> is_post_semver(Version("0.0.0-1"), Version("0.0.0-1"))
1784
+ False
1785
+ >>> is_post_semver(Version("0.0.1-1"), Version("0.0.0-1"))
1786
+ False
1787
+ >>> is_post_semver(Version("0.1.0-1"), Version("0.0.0-1"))
1788
+ False
1789
+ >>> is_post_semver(Version("1.0.0-1"), Version("0.0.0-1"))
1790
+ False
1791
+ >>> is_post_semver(Version("0.0.1-1"), Version("0.0.0"))
1792
+ False
1793
+ """
1794
+ if (
1795
+ new_version.major == current_version.major
1796
+ and new_version.minor == current_version.minor
1797
+ and new_version.micro == current_version.micro
1798
+ ):
1799
+ return new_version.post is not None and new_version.post != current_version.post
1800
+
1801
+ return False
1802
+
1803
+
1804
+ def is_major_semver(new_version: Version, current_version: Version) -> bool:
1805
+ """
1806
+ Check if only the major part of the semantic version has changed.
1807
+
1808
+ Args:
1809
+ new_version (Version): The new version to check.
1810
+ current_version (Version): The current version to compare with.
1811
+
1812
+ Returns:
1813
+ bool: True if only the major part of the version has changed, False otherwise.
1814
+
1815
+ Examples:
1816
+ >>> is_major_semver(Version("1.0.0"), Version("0.0.0"))
1817
+ True
1818
+ >>> is_major_semver(Version("0.0.0"), Version("0.0.0"))
1819
+ False
1820
+ >>> is_major_semver(Version("1.0.1"), Version("1.0.0"))
1821
+ False
1822
+ >>> is_major_semver(Version("2.0.0-1"), Version("1.0.1-2"))
1823
+ True
1824
+ """
1825
+
1826
+ return new_version.major != current_version.major
1827
+
1828
+
1829
+ async def print_release_summary(config: CLIConfig, semver: Optional[str], info: bool = False, dry_run=False):
1830
+ if not semver:
1831
+ click.echo(FeedbackManager.info_release_no_rollback())
1832
+ return
1833
+ try:
1834
+ client = config.get_client()
1835
+ response = await client.release_rm(config["id"], semver, confirmation=config["name"], dry_run=True)
1836
+ except Exception as e:
1837
+ raise CLIReleaseException(FeedbackManager.error_exception(error=str(e)))
1838
+ else:
1839
+ columns = ["name", "id", "notes"]
1840
+ if not response:
1841
+ click.echo(FeedbackManager.info_release_no_rollback())
1842
+ return
1843
+
1844
+ if len(response["datasources"]) or len(response["pipes"]):
1845
+ semver = response.get("semver", semver)
1846
+ if info:
1847
+ if dry_run:
1848
+ click.echo(FeedbackManager.info_release_rm_resources_dry_run(semver=semver))
1849
+ else:
1850
+ click.echo(FeedbackManager.info_release_rm_resources())
1851
+ else:
1852
+ click.echo(FeedbackManager.info_release_rollback(semver=semver))
1853
+
1854
+ if len(response["datasources"]):
1855
+ click.echo("\nDatasources:")
1856
+ rows = [
1857
+ [ds, response["datasources"][ds], response["notes"].get(response["datasources"][ds], "")]
1858
+ for ds in response["datasources"]
1859
+ ]
1860
+ echo_safe_humanfriendly_tables_format_smart_table(rows, column_names=columns)
1861
+
1862
+ if len(response["pipes"]):
1863
+ click.echo("\nPipes:")
1864
+ rows = [
1865
+ [pipe, response["pipes"][pipe], response["notes"].get(response["pipes"][pipe], "")]
1866
+ for pipe in response["pipes"]
1867
+ ]
1868
+ echo_safe_humanfriendly_tables_format_smart_table(rows, column_names=columns)
1869
+
1870
+
1871
+ async def get_oldest_rollback(config: CLIConfig, client: TinyB) -> Optional[str]:
1872
+ oldest_rollback_response = await client.release_oldest_rollback(config["id"])
1873
+ return oldest_rollback_response.get("semver") if oldest_rollback_response else None
1874
+
1875
+
1876
+ async def remove_release(
1877
+ dry_run: bool, config: CLIConfig, semver: Optional[str], client: TinyB, force: bool, show_print=True
1878
+ ):
1879
+ if semver == OLDEST_ROLLBACK:
1880
+ semver = await get_oldest_rollback(config, client)
1881
+ if show_print:
1882
+ await print_release_summary(config, semver, info=True, dry_run=True)
1883
+ if not dry_run:
1884
+ if semver:
1885
+ response = await client.release_rm(
1886
+ config["id"], semver, confirmation=config["name"], dry_run=dry_run, force=force
1887
+ )
1888
+ click.echo(FeedbackManager.success_release_delete(semver=response.get("semver")))
1889
+ else:
1890
+ click.echo(FeedbackManager.info_no_release_deleted())
1891
+
1892
+
1893
+ async def validate_aws_iamrole_integration(
1894
+ client: TinyB,
1895
+ service: str,
1896
+ role_arn: Optional[str],
1897
+ region: Optional[str],
1898
+ policy: str = "write",
1899
+ no_validate: Optional[bool] = False,
1900
+ ):
1901
+ if no_validate is False:
1902
+ access_policy, trust_policy, external_id = await get_aws_iamrole_policies(
1903
+ client, service=service, policy=policy
1904
+ )
1905
+
1906
+ if not role_arn:
1907
+ if not click.confirm(
1908
+ FeedbackManager.prompt_s3_iamrole_connection_login_aws(),
1909
+ show_default=False,
1910
+ prompt_suffix="Press y to continue:",
1911
+ ):
1912
+ sys.exit(1)
1913
+
1914
+ access_policy_copied = True
1915
+ try:
1916
+ pyperclip.copy(access_policy)
1917
+ except Exception:
1918
+ access_policy_copied = False
1919
+
1920
+ replacements_dict = {
1921
+ "<bucket>": "<bucket> with your bucket name",
1922
+ "<table_name>": "<table_name> with your DynamoDB table name",
1923
+ }
1924
+
1925
+ replacements = [
1926
+ replacements_dict.get(replacement, "")
1927
+ for replacement in replacements_dict.keys()
1928
+ if replacement in access_policy
1929
+ ]
1930
+
1931
+ if not click.confirm(
1932
+ (
1933
+ FeedbackManager.prompt_s3_iamrole_connection_policy(
1934
+ access_policy=access_policy, replacements=", ".join(replacements)
1935
+ )
1936
+ if access_policy_copied
1937
+ else FeedbackManager.prompt_s3_iamrole_connection_policy_not_copied(access_policy=access_policy)
1938
+ ),
1939
+ show_default=False,
1940
+ prompt_suffix="Press y to continue:",
1941
+ ):
1942
+ sys.exit(1)
1943
+
1944
+ trust_policy_copied = True
1945
+ try:
1946
+ pyperclip.copy(trust_policy)
1947
+ except Exception:
1948
+ trust_policy_copied = False
1949
+
1950
+ if not click.confirm(
1951
+ (
1952
+ FeedbackManager.prompt_s3_iamrole_connection_role(trust_policy=trust_policy)
1953
+ if trust_policy_copied
1954
+ else FeedbackManager.prompt_s3_iamrole_connection_role_not_copied(trust_policy=trust_policy)
1955
+ ),
1956
+ show_default=False,
1957
+ prompt_suffix="Press y to continue:",
1958
+ ):
1959
+ sys.exit(1)
1960
+ else:
1961
+ try:
1962
+ trust_policy = await client.get_trust_policy(service)
1963
+ external_id = trust_policy["Statement"][0]["Condition"]["StringEquals"]["sts:ExternalId"]
1964
+ except Exception:
1965
+ external_id = ""
1966
+
1967
+ if not role_arn:
1968
+ role_arn = click.prompt("Enter the ARN of the role you just created")
1969
+ validate_string_connector_param("Role ARN", role_arn)
1970
+
1971
+ if not region:
1972
+ region_resource = "table" if service == DataConnectorType.AMAZON_DYNAMODB else "bucket"
1973
+ region = click.prompt(f"Enter the region where the {region_resource} is located")
1974
+ validate_string_connector_param("Region", region)
1975
+
1976
+ return role_arn, region, external_id
1977
+
1978
+
1979
+ async def get_aws_iamrole_policies(client: TinyB, service: str, policy: str = "write"):
1980
+ access_policy: Dict[str, Any] = {}
1981
+ if service == DataConnectorType.AMAZON_S3_IAMROLE:
1982
+ service = DataConnectorType.AMAZON_S3
1983
+ try:
1984
+ if policy == "write":
1985
+ access_policy = await client.get_access_write_policy(service)
1986
+ elif policy == "read":
1987
+ access_policy = await client.get_access_read_policy(service)
1988
+ else:
1989
+ raise Exception(f"Access policy {policy} not supported. Choose from 'read' or 'write'")
1990
+ if not len(access_policy) > 0:
1991
+ raise Exception(f"{service.upper()} Integration not supported in this region")
1992
+ except Exception as e:
1993
+ raise CLIConnectionException(FeedbackManager.error_connection_integration_not_available(error=str(e)))
1994
+
1995
+ trust_policy: Dict[str, Any] = {}
1996
+ try:
1997
+ trust_policy = await client.get_trust_policy(service)
1998
+ if not len(trust_policy) > 0:
1999
+ raise Exception(f"{service.upper()} Integration not supported in this region")
2000
+ except Exception as e:
2001
+ raise CLIConnectionException(FeedbackManager.error_connection_integration_not_available(error=str(e)))
2002
+ try:
2003
+ external_id = trust_policy["Statement"][0]["Condition"]["StringEquals"]["sts:ExternalId"]
2004
+ except Exception:
2005
+ external_id = ""
2006
+ return json.dumps(access_policy, indent=4), json.dumps(trust_policy, indent=4), external_id
2007
+
2008
+
2009
+ async def validate_aws_iamrole_connection_name(
2010
+ client: TinyB, connection_name: Optional[str], no_validate: Optional[bool] = False
2011
+ ) -> str:
2012
+ if connection_name and no_validate is False:
2013
+ if await client.get_connector(connection_name, skip_bigquery=True) is not None:
2014
+ raise CLIConnectionException(FeedbackManager.info_connection_already_exists(name=connection_name))
2015
+ else:
2016
+ while not connection_name:
2017
+ connection_name = click.prompt("Enter the name for this connection", default=None, show_default=False)
2018
+ assert isinstance(connection_name, str)
2019
+
2020
+ if no_validate is False and await client.get_connector(connection_name) is not None:
2021
+ click.echo(FeedbackManager.info_connection_already_exists(name=connection_name))
2022
+ connection_name = None
2023
+ assert isinstance(connection_name, str)
2024
+ return connection_name
2025
+
2026
+
2027
+ class DataConnectorType(str, Enum):
2028
+ KAFKA = "kafka"
2029
+ GCLOUD_SCHEDULER = "gcscheduler"
2030
+ SNOWFLAKE = "snowflake"
2031
+ BIGQUERY = "bigquery"
2032
+ GCLOUD_STORAGE = "gcs"
2033
+ GCLOUD_STORAGE_HMAC = "gcs_hmac"
2034
+ GCLOUD_STORAGE_SA = "gcs_service_account"
2035
+ AMAZON_S3 = "s3"
2036
+ AMAZON_S3_IAMROLE = "s3_iamrole"
2037
+ AMAZON_DYNAMODB = "dynamodb"
2038
+
2039
+ def __str__(self) -> str:
2040
+ return self.value
2041
+
2042
+
2043
+ async def create_aws_iamrole_connection(client: TinyB, service: str, connection_name, role_arn, region) -> None:
2044
+ conn_file_name = f"{connection_name}.connection"
2045
+ conn_file_path = Path(getcwd(), conn_file_name)
2046
+
2047
+ if os.path.isfile(conn_file_path):
2048
+ raise CLIConnectionException(FeedbackManager.error_connection_file_already_exists(name=conn_file_name))
2049
+
2050
+ if service == DataConnectorType.AMAZON_S3_IAMROLE:
2051
+ click.echo(FeedbackManager.info_creating_s3_iamrole_connection(connection_name=connection_name))
2052
+ if service == DataConnectorType.AMAZON_DYNAMODB:
2053
+ click.echo(FeedbackManager.info_creating_dynamodb_connection(connection_name=connection_name))
2054
+
2055
+ params = ConnectionReplacements.map_api_params_from_prompt_params(
2056
+ service, connection_name=connection_name, role_arn=role_arn, region=region
2057
+ )
2058
+
2059
+ click.echo("** Creating connection...")
2060
+ try:
2061
+ _ = await client.connection_create(params)
2062
+ except Exception as e:
2063
+ raise CLIConnectionException(
2064
+ FeedbackManager.error_connection_create(connection_name=connection_name, error=str(e))
2065
+ )
2066
+
2067
+ async with aiofiles.open(conn_file_path, "w") as f:
2068
+ await f.write(
2069
+ f"""TYPE {service}
2070
+
2071
+ """
2072
+ )
2073
+ click.echo(FeedbackManager.success_connection_file_created(name=conn_file_name))
2074
+
2075
+
2076
+ def get_ca_pem_content(ca_pem: Optional[str], filename: Optional[str] = None) -> Optional[str]:
2077
+ if not ca_pem:
2078
+ return None
2079
+
2080
+ def is_valid_content(text_content: str) -> bool:
2081
+ return text_content.startswith("-----BEGIN CERTIFICATE-----")
2082
+
2083
+ ca_pem_content = ca_pem
2084
+ base_path = Path(getcwd(), filename).parent if filename else Path(getcwd())
2085
+ ca_pem_path = Path(base_path, ca_pem)
2086
+ path_exists = os.path.exists(ca_pem_path)
2087
+
2088
+ if not path_exists:
2089
+ raise CLIConnectionException(FeedbackManager.error_connection_ca_pem_not_found(ca_pem=ca_pem))
2090
+
2091
+ if ca_pem.endswith(".pem") and path_exists:
2092
+ with open(ca_pem_path, "r") as f:
2093
+ ca_pem_content = f.read()
2094
+
2095
+ if not is_valid_content(ca_pem_content):
2096
+ raise CLIConnectionException(FeedbackManager.error_connection_invalid_ca_pem())
2097
+
2098
+ return ca_pem_content