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