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