From 902aba74949ef8b4cebf7033187c354ddb269259 Mon Sep 17 00:00:00 2001
From: Vlad Ivoninskii <vlad.ivoninskii@jetbrains.com>
Date: Wed, 11 Mar 2026 13:21:40 -0300
Subject: [PATCH 01/15] add support for arrow key selection

---
 pyproject.toml                                |  1 +
 .../datasource/add_datasource_config.py       | 26 ++++++++++++++++---
 tests/test_add_datasource.py                  |  9 +++++--
 uv.lock                                       | 25 ++++++++++++++----
 4 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 6011d16c..731b0c2d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,6 +9,7 @@ dependencies = [
     "databao-context-engine[snowflake]~=0.6.0",
     "prettytable>=3.10.0",
     "databao-agent~=0.1.4.dev5",
+    "questionary>=2.1.1",
     "streamlit[snowflake]>=1.53.0",
     "uuid6>=2024.7.10",
     "pyyaml>=6.0",
diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add_datasource_config.py
index b1bb83e2..d05e53bc 100644
--- a/src/databao_cli/commands/datasource/add_datasource_config.py
+++ b/src/databao_cli/commands/datasource/add_datasource_config.py
@@ -1,6 +1,7 @@
 import os
 
 import click
+import questionary
 from databao_context_engine import (
     DatabaoContextDomainManager,
     DatabaoContextPluginLoader,
@@ -11,6 +12,19 @@
 from databao_cli.commands.datasource.check_datasource_connection import print_connection_check_results
 from databao_cli.project.layout import ProjectLayout
 
+DISPLAY_NAMES = {
+    "athena": "Amazon Athena",
+    "bigquery": "BigQuery",
+    "clickhouse": "ClickHouse",
+    "duckdb": "DuckDB",
+    "mssql": "Microsoft SQL Server",
+    "mysql": "MySQL",
+    "parquet": "Parquet",
+    "postgres": "PostgreSQL",
+    "snowflake": "Snowflake",
+    "sqlite": "SQLite",
+}
+
 
 def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain: str) -> None:
     domain_dir = project_layout.domains_dir / domain
@@ -47,11 +61,17 @@ def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain
 
 def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType:
     all_datasource_types = sorted([ds_type.full_type for ds_type in supported_datasource_types])
-    config_type = click.prompt(
+    choices = [
+        questionary.Choice(title=DISPLAY_NAMES.get(datasource_type, datasource_type), value=datasource_type)
+        for datasource_type in all_datasource_types
+    ]
+    config_type = questionary.select(
         "What type of datasource do you want to add?",
-        type=click.Choice(all_datasource_types),
+        choices=choices,
         default=all_datasource_types[0] if len(all_datasource_types) == 1 else None,
-    )
+    ).ask()
+    if config_type is None:
+        raise click.Abort()
     click.echo(f"Selected type: {config_type}")
 
     return DatasourceType(full_type=config_type)
diff --git a/tests/test_add_datasource.py b/tests/test_add_datasource.py
index 6265fe3a..947e5f71 100644
--- a/tests/test_add_datasource.py
+++ b/tests/test_add_datasource.py
@@ -1,10 +1,12 @@
 from pathlib import Path
+from unittest.mock import Mock
 
 import duckdb
 import pytest
 from click.testing import CliRunner
 
 from databao_cli.__main__ import cli
+from databao_cli.commands.datasource import add_datasource_config
 from databao_cli.commands.init import init_impl as init_databao_project
 
 
@@ -31,11 +33,14 @@ def temp_parquet_file(request: pytest.FixtureRequest, tmp_path: Path) -> Path:
     return parquet_file
 
 
-def test_databao_datasource_add(tmp_path: Path, temp_parquet_file: Path) -> None:
+def test_databao_datasource_add(tmp_path: Path, temp_parquet_file: Path, monkeypatch: pytest.MonkeyPatch) -> None:
     init_databao_project(tmp_path)
 
+    prompt = Mock()
+    prompt.ask.return_value = "parquet"
+    monkeypatch.setattr(add_datasource_config.questionary, "select", Mock(return_value=prompt))
+
     inputs = [
-        "parquet",
         "resources/my_parq",
         str(temp_parquet_file),
         "N",  # No. skip adding duckdb secret
diff --git a/uv.lock b/uv.lock
index fa85f4c9..489df7cb 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1036,6 +1036,7 @@ dependencies = [
     { name = "nh3" },
     { name = "prettytable" },
     { name = "pyyaml" },
+    { name = "questionary" },
     { name = "streamlit", extra = ["snowflake"] },
     { name = "uuid6" },
 ]
@@ -1091,6 +1092,7 @@ requires-dist = [
     { name = "nh3", specifier = ">=0.2.15" },
     { name = "prettytable", specifier = ">=3.10.0" },
     { name = "pyyaml", specifier = ">=6.0" },
+    { name = "questionary", specifier = ">=2.1.1" },
     { name = "streamlit", extras = ["snowflake"], specifier = ">=1.53.0" },
     { name = "uuid6", specifier = ">=2024.7.10" },
 ]
@@ -1925,7 +1927,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/e8/2e1462c8fdbe0f210feb5ac7ad2d9029af8be3bf45bd9fa39765f821642f/greenlet-3.3.1-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:5fd23b9bc6d37b563211c6abbb1b3cab27db385a4449af5c32e932f93017080c", size = 274974, upload-time = "2026-01-23T15:31:02.891Z" },
     { url = "https://files.pythonhosted.org/packages/7e/a8/530a401419a6b302af59f67aaf0b9ba1015855ea7e56c036b5928793c5bd/greenlet-3.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:09f51496a0bfbaa9d74d36a52d2580d1ef5ed4fdfcff0a73730abfbbbe1403dd", size = 577175, upload-time = "2026-01-23T16:00:56.213Z" },
     { url = "https://files.pythonhosted.org/packages/8e/89/7e812bb9c05e1aaef9b597ac1d0962b9021d2c6269354966451e885c4e6b/greenlet-3.3.1-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb0feb07fe6e6a74615ee62a880007d976cf739b6669cce95daa7373d4fc69c5", size = 590401, upload-time = "2026-01-23T16:05:26.365Z" },
-    { url = "https://files.pythonhosted.org/packages/70/ae/e2d5f0e59b94a2269b68a629173263fa40b63da32f5c231307c349315871/greenlet-3.3.1-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:67ea3fc73c8cd92f42467a72b75e8f05ed51a0e9b1d15398c913416f2dafd49f", size = 601161, upload-time = "2026-01-23T16:15:53.456Z" },
     { url = "https://files.pythonhosted.org/packages/5c/ae/8d472e1f5ac5efe55c563f3eabb38c98a44b832602e12910750a7c025802/greenlet-3.3.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:39eda9ba259cc9801da05351eaa8576e9aa83eb9411e8f0c299e05d712a210f2", size = 590272, upload-time = "2026-01-23T15:32:49.411Z" },
     { url = "https://files.pythonhosted.org/packages/a8/51/0fde34bebfcadc833550717eade64e35ec8738e6b097d5d248274a01258b/greenlet-3.3.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e2e7e882f83149f0a71ac822ebf156d902e7a5d22c9045e3e0d1daf59cee2cc9", size = 1550729, upload-time = "2026-01-23T16:04:20.867Z" },
     { url = "https://files.pythonhosted.org/packages/16/c9/2fb47bee83b25b119d5a35d580807bb8b92480a54b68fef009a02945629f/greenlet-3.3.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:80aa4d79eb5564f2e0a6144fcc744b5a37c56c4a92d60920720e99210d88db0f", size = 1615552, upload-time = "2026-01-23T15:33:45.743Z" },
@@ -1934,7 +1935,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f9/c8/9d76a66421d1ae24340dfae7e79c313957f6e3195c144d2c73333b5bfe34/greenlet-3.3.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:7e806ca53acf6d15a888405880766ec84721aa4181261cd11a457dfe9a7a4975", size = 276443, upload-time = "2026-01-23T15:30:10.066Z" },
     { url = "https://files.pythonhosted.org/packages/81/99/401ff34bb3c032d1f10477d199724f5e5f6fbfb59816ad1455c79c1eb8e7/greenlet-3.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d842c94b9155f1c9b3058036c24ffb8ff78b428414a19792b2380be9cecf4f36", size = 597359, upload-time = "2026-01-23T16:00:57.394Z" },
     { url = "https://files.pythonhosted.org/packages/2b/bc/4dcc0871ed557792d304f50be0f7487a14e017952ec689effe2180a6ff35/greenlet-3.3.1-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:20fedaadd422fa02695f82093f9a98bad3dab5fcda793c658b945fcde2ab27ba", size = 607805, upload-time = "2026-01-23T16:05:28.068Z" },
-    { url = "https://files.pythonhosted.org/packages/3b/cd/7a7ca57588dac3389e97f7c9521cb6641fd8b6602faf1eaa4188384757df/greenlet-3.3.1-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c620051669fd04ac6b60ebc70478210119c56e2d5d5df848baec4312e260e4ca", size = 622363, upload-time = "2026-01-23T16:15:54.754Z" },
     { url = "https://files.pythonhosted.org/packages/cf/05/821587cf19e2ce1f2b24945d890b164401e5085f9d09cbd969b0c193cd20/greenlet-3.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14194f5f4305800ff329cbf02c5fcc88f01886cadd29941b807668a45f0d2336", size = 609947, upload-time = "2026-01-23T15:32:51.004Z" },
     { url = "https://files.pythonhosted.org/packages/a4/52/ee8c46ed9f8babaa93a19e577f26e3d28a519feac6350ed6f25f1afee7e9/greenlet-3.3.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7b2fe4150a0cf59f847a67db8c155ac36aed89080a6a639e9f16df5d6c6096f1", size = 1567487, upload-time = "2026-01-23T16:04:22.125Z" },
     { url = "https://files.pythonhosted.org/packages/8f/7c/456a74f07029597626f3a6db71b273a3632aecb9afafeeca452cfa633197/greenlet-3.3.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:49f4ad195d45f4a66a0eb9c1ba4832bb380570d361912fa3554746830d332149", size = 1636087, upload-time = "2026-01-23T15:33:47.486Z" },
@@ -1943,7 +1943,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/ab/d26750f2b7242c2b90ea2ad71de70cfcd73a948a49513188a0fc0d6fc15a/greenlet-3.3.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:7ab327905cabb0622adca5971e488064e35115430cec2c35a50fd36e72a315b3", size = 275205, upload-time = "2026-01-23T15:30:24.556Z" },
     { url = "https://files.pythonhosted.org/packages/10/d3/be7d19e8fad7c5a78eeefb2d896a08cd4643e1e90c605c4be3b46264998f/greenlet-3.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:65be2f026ca6a176f88fb935ee23c18333ccea97048076aef4db1ef5bc0713ac", size = 599284, upload-time = "2026-01-23T16:00:58.584Z" },
     { url = "https://files.pythonhosted.org/packages/ae/21/fe703aaa056fdb0f17e5afd4b5c80195bbdab701208918938bd15b00d39b/greenlet-3.3.1-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7a3ae05b3d225b4155bda56b072ceb09d05e974bc74be6c3fc15463cf69f33fd", size = 610274, upload-time = "2026-01-23T16:05:29.312Z" },
-    { url = "https://files.pythonhosted.org/packages/06/00/95df0b6a935103c0452dad2203f5be8377e551b8466a29650c4c5a5af6cc/greenlet-3.3.1-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:12184c61e5d64268a160226fb4818af4df02cfead8379d7f8b99a56c3a54ff3e", size = 624375, upload-time = "2026-01-23T16:15:55.915Z" },
     { url = "https://files.pythonhosted.org/packages/cb/86/5c6ab23bb3c28c21ed6bebad006515cfe08b04613eb105ca0041fecca852/greenlet-3.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6423481193bbbe871313de5fd06a082f2649e7ce6e08015d2a76c1e9186ca5b3", size = 612904, upload-time = "2026-01-23T15:32:52.317Z" },
     { url = "https://files.pythonhosted.org/packages/c2/f3/7949994264e22639e40718c2daf6f6df5169bf48fb038c008a489ec53a50/greenlet-3.3.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:33a956fe78bbbda82bfc95e128d61129b32d66bcf0a20a1f0c08aa4839ffa951", size = 1567316, upload-time = "2026-01-23T16:04:23.316Z" },
     { url = "https://files.pythonhosted.org/packages/8d/6e/d73c94d13b6465e9f7cd6231c68abde838bb22408596c05d9059830b7872/greenlet-3.3.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4b065d3284be43728dd280f6f9a13990b56470b81be20375a207cdc814a983f2", size = 1636549, upload-time = "2026-01-23T15:33:48.643Z" },
@@ -1952,7 +1951,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ae/fb/011c7c717213182caf78084a9bea51c8590b0afda98001f69d9f853a495b/greenlet-3.3.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:bd59acd8529b372775cd0fcbc5f420ae20681c5b045ce25bd453ed8455ab99b5", size = 275737, upload-time = "2026-01-23T15:32:16.889Z" },
     { url = "https://files.pythonhosted.org/packages/41/2e/a3a417d620363fdbb08a48b1dd582956a46a61bf8fd27ee8164f9dfe87c2/greenlet-3.3.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b31c05dd84ef6871dd47120386aed35323c944d86c3d91a17c4b8d23df62f15b", size = 646422, upload-time = "2026-01-23T16:01:00.354Z" },
     { url = "https://files.pythonhosted.org/packages/b4/09/c6c4a0db47defafd2d6bab8ddfe47ad19963b4e30f5bed84d75328059f8c/greenlet-3.3.1-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:02925a0bfffc41e542c70aa14c7eda3593e4d7e274bfcccca1827e6c0875902e", size = 658219, upload-time = "2026-01-23T16:05:30.956Z" },
-    { url = "https://files.pythonhosted.org/packages/e2/89/b95f2ddcc5f3c2bc09c8ee8d77be312df7f9e7175703ab780f2014a0e781/greenlet-3.3.1-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3e0f3878ca3a3ff63ab4ea478585942b53df66ddde327b59ecb191b19dbbd62d", size = 671455, upload-time = "2026-01-23T16:15:57.232Z" },
     { url = "https://files.pythonhosted.org/packages/80/38/9d42d60dffb04b45f03dbab9430898352dba277758640751dc5cc316c521/greenlet-3.3.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34a729e2e4e4ffe9ae2408d5ecaf12f944853f40ad724929b7585bca808a9d6f", size = 660237, upload-time = "2026-01-23T15:32:53.967Z" },
     { url = "https://files.pythonhosted.org/packages/96/61/373c30b7197f9e756e4c81ae90a8d55dc3598c17673f91f4d31c3c689c3f/greenlet-3.3.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:aec9ab04e82918e623415947921dea15851b152b822661cce3f8e4393c3df683", size = 1615261, upload-time = "2026-01-23T16:04:25.066Z" },
     { url = "https://files.pythonhosted.org/packages/fd/d3/ca534310343f5945316f9451e953dcd89b36fe7a19de652a1dc5a0eeef3f/greenlet-3.3.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:71c767cf281a80d02b6c1bdc41c9468e1f5a494fb11bc8688c360524e273d7b1", size = 1683719, upload-time = "2026-01-23T15:33:50.61Z" },
@@ -1961,7 +1959,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/28/24/cbbec49bacdcc9ec652a81d3efef7b59f326697e7edf6ed775a5e08e54c2/greenlet-3.3.1-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:3e63252943c921b90abb035ebe9de832c436401d9c45f262d80e2d06cc659242", size = 282706, upload-time = "2026-01-23T15:33:05.525Z" },
     { url = "https://files.pythonhosted.org/packages/86/2e/4f2b9323c144c4fe8842a4e0d92121465485c3c2c5b9e9b30a52e80f523f/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:76e39058e68eb125de10c92524573924e827927df5d3891fbc97bd55764a8774", size = 651209, upload-time = "2026-01-23T16:01:01.517Z" },
     { url = "https://files.pythonhosted.org/packages/d9/87/50ca60e515f5bb55a2fbc5f0c9b5b156de7d2fc51a0a69abc9d23914a237/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c9f9d5e7a9310b7a2f416dd13d2e3fd8b42d803968ea580b7c0f322ccb389b97", size = 654300, upload-time = "2026-01-23T16:05:32.199Z" },
-    { url = "https://files.pythonhosted.org/packages/7c/25/c51a63f3f463171e09cb586eb64db0861eb06667ab01a7968371a24c4f3b/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4b9721549a95db96689458a1e0ae32412ca18776ed004463df3a9299c1b257ab", size = 662574, upload-time = "2026-01-23T16:15:58.364Z" },
     { url = "https://files.pythonhosted.org/packages/1d/94/74310866dfa2b73dd08659a3d18762f83985ad3281901ba0ee9a815194fb/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:92497c78adf3ac703b57f1e3813c2d874f27f71a178f9ea5887855da413cd6d2", size = 653842, upload-time = "2026-01-23T15:32:55.671Z" },
     { url = "https://files.pythonhosted.org/packages/97/43/8bf0ffa3d498eeee4c58c212a3905dd6146c01c8dc0b0a046481ca29b18c/greenlet-3.3.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ed6b402bc74d6557a705e197d47f9063733091ed6357b3de33619d8a8d93ac53", size = 1614917, upload-time = "2026-01-23T16:04:26.276Z" },
     { url = "https://files.pythonhosted.org/packages/89/90/a3be7a5f378fc6e84abe4dcfb2ba32b07786861172e502388b4c90000d1b/greenlet-3.3.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:59913f1e5ada20fde795ba906916aea25d442abcc0593fba7e26c92b7ad76249", size = 1676092, upload-time = "2026-01-23T15:33:52.176Z" },
@@ -5142,6 +5139,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
 ]
 
+[[package]]
+name = "questionary"
+version = "2.1.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "prompt-toolkit" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f6/45/eafb0bba0f9988f6a2520f9ca2df2c82ddfa8d67c95d6625452e97b204a5/questionary-2.1.1.tar.gz", hash = "sha256:3d7e980292bb0107abaa79c68dd3eee3c561b83a0f89ae482860b181c8bd412d", size = 25845, upload-time = "2025-08-28T19:00:20.851Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3c/26/1062c7ec1b053db9e499b4d2d5bc231743201b74051c973dadeac80a8f43/questionary-2.1.1-py3-none-any.whl", hash = "sha256:a51af13f345f1cdea62347589fbb6df3b290306ab8930713bfae4d475a7d4a59", size = 36753, upload-time = "2025-08-28T19:00:19.56Z" },
+]
+
 [[package]]
 name = "rapidocr"
 version = "3.6.0"
@@ -6196,6 +6205,12 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" },
     { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" },
     { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" },
+    { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" },
+    { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" },
+    { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" },
+    { url = "https://files.pythonhosted.org/packages/f4/39/590742415c3030551944edc2ddc273ea1fdfe8ffb2780992e824f1ebee98/torch-2.10.0-3-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b1d5e2aba4eb7f8e87fbe04f86442887f9167a35f092afe4c237dfcaaef6e328", size = 915632474, upload-time = "2026-03-11T14:15:13.666Z" },
+    { url = "https://files.pythonhosted.org/packages/b6/8e/34949484f764dde5b222b7fe3fede43e4a6f0da9d7f8c370bb617d629ee2/torch-2.10.0-3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0228d20b06701c05a8f978357f657817a4a63984b0c90745def81c18aedfa591", size = 915523882, upload-time = "2026-03-11T14:14:46.311Z" },
     { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" },
     { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" },
     { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" },

From 4b4f451b4d2c1b265645ecc7e035c5f28d61e85e Mon Sep 17 00:00:00 2001
From: Vlad Ivoninskii <vlad.ivoninskii@jetbrains.com>
Date: Sun, 15 Mar 2026 10:51:05 -0300
Subject: [PATCH 02/15] add support for arrow key selection

---
 src/databao_cli/__main__.py                   | 11 +++-
 src/databao_cli/commands/app.py               |  6 +-
 src/databao_cli/commands/ask.py               | 22 +++----
 .../datasource/add_datasource_config.py       | 36 +++--------
 src/databao_cli/commands/status.py            | 16 ++---
 src/databao_cli/labels.py                     | 35 +++++++++++
 .../mcp/tools/databao_ask/agent_factory.py    |  8 +--
 src/databao_cli/utils.py                      | 63 +++++++++++++++++++
 8 files changed, 140 insertions(+), 57 deletions(-)
 create mode 100644 src/databao_cli/labels.py
 create mode 100644 src/databao_cli/utils.py

diff --git a/src/databao_cli/__main__.py b/src/databao_cli/__main__.py
index 3fee1138..47bf91c1 100644
--- a/src/databao_cli/__main__.py
+++ b/src/databao_cli/__main__.py
@@ -7,6 +7,9 @@
 from databao_cli.log.logging import configure_logging
 from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project
 
+from databao_cli.labels import LABELS
+from databao_cli.utils import ask_confirm, register_labels
+
 
 @click.group()
 @click.option(
@@ -25,6 +28,8 @@ def cli(ctx: Context, project_dir: Path | None) -> None:
     ctx.ensure_object(dict)
     ctx.obj["project_dir"] = project_path
 
+    register_labels(LABELS)
+
 
 @cli.command()
 @click.pass_context
@@ -48,7 +53,7 @@ def init(ctx: Context) -> None:
     try:
         project_layout = init_impl(project_dir)
     except ProjectDirDoesnotExistError:
-        if click.confirm(
+        if ask_confirm(
             f"The directory {project_dir.resolve()} does not exist. Do you want to create it?",
             default=True,
         ):
@@ -68,12 +73,12 @@ def init(ctx: Context) -> None:
     # except RuntimeError as e:
     #     click.echo(str(e), err=True)
 
-    if not click.confirm("\nDo you want to configure a domain now?"):
+    if not ask_confirm("Do you want to configure a domain now?", default=False):
         return
 
     add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN)
 
-    while click.confirm("\nDo you want to add more datasources?"):
+    while ask_confirm("Do you want to add more datasources?", default=False):
         add_datasource_config_interactive_impl(project_layout, ROOT_DOMAIN)
 
 
diff --git a/src/databao_cli/commands/app.py b/src/databao_cli/commands/app.py
index 866dc6a9..aae132bd 100644
--- a/src/databao_cli/commands/app.py
+++ b/src/databao_cli/commands/app.py
@@ -9,7 +9,7 @@
 
 
 def app_impl(ctx: click.Context) -> None:
-    click.echo("Starting Databao UI...")
+    click.echo("Starting Databao web interface...")
 
     try:
         bootstrap_streamlit_app(
@@ -18,7 +18,7 @@ def app_impl(ctx: click.Context) -> None:
             read_only_domain=ctx.obj.get("read_only_domain", False),
         )
     except subprocess.CalledProcessError as e:
-        click.echo(f"Error running Streamlit: {e}", err=True)
+        click.echo(f"Error starting web interface: {e}", err=True)
         sys.exit(1)
     except KeyboardInterrupt:
-        click.echo("\nShutting down Databao...")
+        click.echo("\nShutting down...")
diff --git a/src/databao_cli/commands/ask.py b/src/databao_cli/commands/ask.py
index 97e8747f..f6c5e870 100644
--- a/src/databao_cli/commands/ask.py
+++ b/src/databao_cli/commands/ask.py
@@ -36,8 +36,8 @@ def dataframe_to_prettytable(df: pd.DataFrame, max_rows: int = DEFAULT_MAX_DISPL
 
 
 def initialize_agent_from_dce(project_path: Path, model: str | None, temperature: float) -> Agent:
-    """Initialize the Databao agent using DCE project at the given path."""
-    # Validate DCE project
+    """Initialize the Databao agent using a Context Engine project at the given path."""
+    # Validate a Context Engine project
     project = ProjectLayout(project_path)
 
     status = databao_project_status(project)
@@ -50,12 +50,12 @@ def initialize_agent_from_dce(project_path: Path, model: str | None, temperature
 
     if status == DatabaoProjectStatus.NO_DATASOURCES:
         click.echo(
-            f"No datasources configured in project at {project.project_dir}. Add datasources first.",
+            f"No data sources configured in project at {project.project_dir}. Add data sources first",
             err=True,
         )
         sys.exit(1)
 
-    click.echo(f"Using DCE project: {project.project_dir}")
+    click.echo(f"Using Context Engine project: {project.project_dir}")
 
     _domain = create_domain(project.root_domain_dir)
 
@@ -101,18 +101,18 @@ def display_result(thread: Thread) -> None:
 
 def _print_help() -> None:
     """Print help message for interactive mode."""
-    click.echo("Databao REPL")
+    click.echo("Databao interactive chat")
     click.echo("Ask questions about your data in natural language.\n")
     click.echo("Commands:")
-    click.echo("  \\help   - Show this help")
-    click.echo("  \\clear  - Start a new conversation")
+    click.echo("  \\help   - Show this help message")
+    click.echo("  \\clear  - Clear conversation history")
     click.echo("  \\q      - Exit\n")
 
 
 def run_interactive_mode(agent: Agent, show_thinking: bool) -> None:
     """Run the interactive REPL mode."""
-    click.echo("\nDatabao REPL")
-    click.echo("\nType \\help for available commands.\n")
+    click.echo("\nDatabao interactive chat")
+    click.echo("Type \\help for available commands.\n")
 
     writer = _create_cli_writer() if show_thinking else None
 
@@ -148,14 +148,14 @@ def run_interactive_mode(agent: Agent, show_thinking: bool) -> None:
                     stream_ask=show_thinking,
                     writer=writer,
                 )
-                click.echo("Conversation cleared.\n")
+                click.echo("Conversation history cleared.\n")
                 continue
 
             if command == "help":
                 _print_help()
                 continue
 
-            click.echo(f"Unknown command: {user_input}. Type \\help for available commands.\n")
+            click.echo(f"Unknown command: {user_input}\nType \\help for available commands\n")
             continue
 
         # Process as a question
diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add_datasource_config.py
index d05e53bc..bff0858e 100644
--- a/src/databao_cli/commands/datasource/add_datasource_config.py
+++ b/src/databao_cli/commands/datasource/add_datasource_config.py
@@ -1,30 +1,17 @@
 import os
 
 import click
-import questionary
 from databao_context_engine import (
     DatabaoContextDomainManager,
     DatabaoContextPluginLoader,
     DatasourceType,
 )
+from databao_cli.utils import ask_confirm, ask_select, ask_text
 
 from databao_cli.commands.context_engine_cli import ClickUserInputCallback
 from databao_cli.commands.datasource.check_datasource_connection import print_connection_check_results
 from databao_cli.project.layout import ProjectLayout
 
-DISPLAY_NAMES = {
-    "athena": "Amazon Athena",
-    "bigquery": "BigQuery",
-    "clickhouse": "ClickHouse",
-    "duckdb": "DuckDB",
-    "mssql": "Microsoft SQL Server",
-    "mysql": "MySQL",
-    "parquet": "Parquet",
-    "postgres": "PostgreSQL",
-    "snowflake": "Snowflake",
-    "sqlite": "SQLite",
-}
-
 
 def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain: str) -> None:
     domain_dir = project_layout.domains_dir / domain
@@ -35,11 +22,11 @@ def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain
 
     datasource_type = _ask_for_datasource_type(plugin_loader.get_all_supported_datasource_types(exclude_file_plugins=True))
 
-    datasource_name = click.prompt("Datasource name?", type=str)
+    datasource_name = ask_text("Datasource name")
 
     datasource_id = domain_manager.datasource_config_exists(datasource_name=datasource_name)
     if datasource_id is not None:
-        click.confirm(
+        ask_confirm(
             f"A config file already exists for this datasource {datasource_id.relative_path_to_config_file()}. "
             f"Do you want to overwrite it?",
             abort=True,
@@ -54,24 +41,17 @@ def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain
         f"{os.linesep}We've created a new config file for your datasource at: "
         f"{domain_manager.get_config_file_path_for_datasource(datasource_id)}"
     )
-    if click.confirm("\nDo you want to check the connection to this new datasource?"):
+    if ask_confirm("Do you want to check the connection to this new datasource?", default=True):
         results = domain_manager.check_datasource_connection(datasource_ids=[datasource_id])
         print_connection_check_results(domain, results)
 
 
 def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) -> DatasourceType:
     all_datasource_types = sorted([ds_type.full_type for ds_type in supported_datasource_types])
-    choices = [
-        questionary.Choice(title=DISPLAY_NAMES.get(datasource_type, datasource_type), value=datasource_type)
-        for datasource_type in all_datasource_types
-    ]
-    config_type = questionary.select(
+    config_type = ask_select(
         "What type of datasource do you want to add?",
-        choices=choices,
-        default=all_datasource_types[0] if len(all_datasource_types) == 1 else None,
-    ).ask()
-    if config_type is None:
-        raise click.Abort()
-    click.echo(f"Selected type: {config_type}")
+        choices=all_datasource_types,
+        default=all_datasource_types[0] if len(all_datasource_types) == 0 else None,
+    )
 
     return DatasourceType(full_type=config_type)
diff --git a/src/databao_cli/commands/status.py b/src/databao_cli/commands/status.py
index 7737396c..c0779d5a 100644
--- a/src/databao_cli/commands/status.py
+++ b/src/databao_cli/commands/status.py
@@ -28,19 +28,19 @@ def status_impl(project_dir: Path) -> str:
 
 def _generate_info_string(command_info: DceInfo, domain_infos: list[DceDomainInfo]) -> str:
     info_lines = [
-        f"Databao context engine version: {command_info.version}",
-        f"Databao agent version: {version('databao-agent')}",
-        f"Databao context engine storage dir: {command_info.dce_path}",
-        f"Databao context engine plugins: {command_info.plugin_ids}",
+        f"Context Engine version: {command_info.version}",
+        f"Agent version: {version('databao-agent')}",
+        f"Context Engine storage directory: {command_info.dce_path}",
+        f"Context Engine plugins: {command_info.plugin_ids}",
         "",
-        f"OS name: {sys.platform}",
-        f"OS architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}",
+        f"OS: {sys.platform}",
+        f"Architecture: {os.uname().machine if hasattr(os, 'uname') else 'unknown'}",
         "",
     ]
 
     for domain_info in domain_infos:
         if domain_info.is_initialized:
-            info_lines.append(f"Databao Domain dir: {domain_info.project_path.resolve()}")
-            info_lines.append(f"Databao Domain ID: {domain_info.project_id!s}")
+            info_lines.append(f"Domain directory: {domain_info.project_path.resolve()}")
+            info_lines.append(f"Domain ID: {domain_info.project_id!s}")
 
     return os.linesep.join(info_lines)
diff --git a/src/databao_cli/labels.py b/src/databao_cli/labels.py
new file mode 100644
index 00000000..3cefe8af
--- /dev/null
+++ b/src/databao_cli/labels.py
@@ -0,0 +1,35 @@
+LABELS = {
+    "athena": "Amazon Athena",
+    "bigquery": "BigQuery",
+    "clickhouse": "ClickHouse",
+    "duckdb": "DuckDB",
+    "mssql": "Microsoft SQL Server",
+    "mysql": "MySQL",
+    "parquet": "Parquet",
+    "postgres": "PostgreSQL",
+    "snowflake": "Snowflake",
+    "sqlite": "SQLite",
+    "BigQueryDefaultAuth": "Default auth",
+    "BigQueryServiceAccountJsonAuth": "Service account JSON credentials",
+    "BigQueryServiceAccountKeyFileAuth": "Service account key file",
+    "SnowflakeKeyPairAuth": "Key pair",
+    "SnowflakePasswordAuth": "Password",
+    "SnowflakeSSOAuth": "SSO",
+    "connection.auth.type": "Authentication type",
+    "connection.host": "Host",
+    "connection.port": "Port",
+    "connection.database": "Database",
+    "connection.schema": "Schema",
+    "connection.username": "Username",
+    "connection.password": "Password",
+    "connection.account": "Account",
+    "connection.warehouse": "Warehouse",
+    "connection.role": "Role",
+    "connection.path": "File path",
+    "connection.project": "Project",
+    "connection.dataset": "Dataset",
+    "connection.location": "Location",
+    "connection.auth.credentials_file": "Credentials file",
+    "connection.auth.key_file": "Key file",
+    "connection.auth.token": "Token",
+}
\ No newline at end of file
diff --git a/src/databao_cli/mcp/tools/databao_ask/agent_factory.py b/src/databao_cli/mcp/tools/databao_ask/agent_factory.py
index 9bba5270..29cb089e 100644
--- a/src/databao_cli/mcp/tools/databao_ask/agent_factory.py
+++ b/src/databao_cli/mcp/tools/databao_ask/agent_factory.py
@@ -19,7 +19,7 @@ def create_agent_for_tool(
     executor: str = "lighthouse",
     cache: Cache | None = None,
 ) -> Agent:
-    """Create a Databao agent from a DCE project, configured for MCP tool use.
+    """Create a Databao agent from a Context Engine project, configured for MCP tool use.
 
     Raises ValueError if the project is not ready (no datasources, no build).
     """
@@ -27,11 +27,11 @@ def create_agent_for_tool(
 
     status = databao_project_status(project)
     if status == DatabaoProjectStatus.NOT_INITIALIZED:
-        raise ValueError("Databao project is not initialized. Run 'databao init' first.")
+        raise ValueError("Databao project is not initialized. Run 'databao init' first")
     if status == DatabaoProjectStatus.NO_DATASOURCES:
-        raise ValueError("No datasources configured. Run 'databao datasource add' first.")
+        raise ValueError("No data sources configured. Run 'databao datasource add' first")
     if not has_build_output(project):
-        raise ValueError("Project has no build output. Run 'databao build' first.")
+        raise ValueError("Project has no build output. Run 'databao build' first")
 
     domain = create_domain(project.root_domain_dir)
 
diff --git a/src/databao_cli/utils.py b/src/databao_cli/utils.py
new file mode 100644
index 00000000..f0c89f55
--- /dev/null
+++ b/src/databao_cli/utils.py
@@ -0,0 +1,63 @@
+import sys
+import click
+import questionary
+
+_labels: dict[str, str] = {}
+
+def register_labels(labels: dict) -> None:
+    _labels.update(labels)
+
+def _resolve(value: str) -> tuple[str, str]:
+    label = _labels.get(value, value)
+    return value, label
+
+def is_interactive() -> bool:
+    """True when running in an interactive terminal."""
+    return sys.stdin.isatty() and sys.stdout.isatty()
+
+
+def ask_select(message: str, choices: list[str], default: str = None) -> str:
+    """Select from a list. Interactive in TTY, plain text otherwise."""
+    if is_interactive():
+        resolved = [_resolve(c) if isinstance(c, str) else c for c in choices]
+        q_choices = [
+            questionary.Choice(title=label, value=value) for value, label in resolved
+        ]
+        result = questionary.select(message, choices=q_choices, default=default).ask()
+        if result is None:
+            raise click.Abort()
+        return result
+    else:
+        click.echo(f"{message}")
+        for i, choice in enumerate(choices, 1):
+            click.echo(f"  {i}. {choice}")
+        value = click.prompt(
+            "Enter a number of value",
+            default=default or choices[0],
+        )
+        if value.isdigit() and 1 <= int(value) <= len(choices):
+            return choices[int(value) - 1]
+        if value in choices:
+            return value
+        raise click.BadParameter(f"Invalid choice: {value}")
+
+
+def ask_confirm(message: str, default: bool = True, abort: bool = False) -> bool:
+    """Yes/no. Fancy in TTY, plain click.confirm otherwise."""
+    if is_interactive():
+        result = questionary.confirm(message, default=default).ask()
+        if result is None:
+            raise click.Abort()
+        if abort and not result:
+            raise click.Abort()
+        return result
+    else:
+        return click.confirm(message, default=default, abort=abort)
+
+
+def ask_text(message: str, default: str = None) -> str:
+    """Text input. Interactive in TTY, plain click.prompt otherwise."""
+    if is_interactive():
+        return questionary.text(message, default=default or "").ask()
+    else:
+        return click.prompt(message, default=default)

From 8d53f0ada26334095cf113267104073af371ff0b Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 11:00:27 +0100
Subject: [PATCH 03/15] e2e-tests: expand Ollama model detection and
 installation steps

(cherry picked from commit 0bbe9d9172fbea5e8501f4f4b44d137fb2c71a89)
---
 e2e-tests/src/project_utils.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/e2e-tests/src/project_utils.py b/e2e-tests/src/project_utils.py
index f60b0867..c4c10fc1 100644
--- a/e2e-tests/src/project_utils.py
+++ b/e2e-tests/src/project_utils.py
@@ -40,7 +40,13 @@ def execute_build(project_dir: Path):
         child = pexpect.spawn("uv run databao build", cwd=project_dir, encoding="utf-8", timeout=30, logfile=logfile)
         try:
             # Wait for Ollama download/installation with extended timeout
-            if child.expect([r"Ollama model .+ not found locally\.", "Found datasource of type"], timeout=5) == 0:
+            if child.expect([
+                r"Ollama model .+ not found locally\.",
+                r"No existing Ollama installation detected",
+                r"We will download and install Ollama.",
+                r"Downloading .*ollama.*\.tgz",
+                r"Found datasource of type",
+            ], timeout=5) == 0:
                 with allure.step("Ollama model not found locally, installing..."):
                     child.expect("Ollama model .+ pulled successfully", timeout=900)
             child.expect(r"Build complete\. Processed \d+ datasources\.")

From 8324e99f659addcf9927b045e5b4eeedf7a54945 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 11:30:04 +0100
Subject: [PATCH 04/15] e2e-tests: adjust timeouts for build execution and
 model installation

(cherry picked from commit a99ed932afa4e58e69595b5b135a7a928bae4adc)
---
 e2e-tests/src/project_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/e2e-tests/src/project_utils.py b/e2e-tests/src/project_utils.py
index c4c10fc1..5dad46f4 100644
--- a/e2e-tests/src/project_utils.py
+++ b/e2e-tests/src/project_utils.py
@@ -37,7 +37,7 @@ def execute_init(project_dir: Path, db: PostgresDB | MysqlDB | SnowflakeDB | Big
 def execute_build(project_dir: Path):
     log_file_path = project_dir / "cli.log"
     with open(log_file_path, "w") as logfile:
-        child = pexpect.spawn("uv run databao build", cwd=project_dir, encoding="utf-8", timeout=30, logfile=logfile)
+        child = pexpect.spawn("uv run databao build", cwd=project_dir, encoding="utf-8", timeout=900, logfile=logfile)
         try:
             # Wait for Ollama download/installation with extended timeout
             if child.expect([
@@ -49,7 +49,7 @@ def execute_build(project_dir: Path):
             ], timeout=5) == 0:
                 with allure.step("Ollama model not found locally, installing..."):
                     child.expect("Ollama model .+ pulled successfully", timeout=900)
-            child.expect(r"Build complete\. Processed \d+ datasources\.")
+            child.expect(r"Build complete\. Processed \d+ datasources\.", timeout=30)
             child.expect(pexpect.EOF)
         finally:
             if log_file_path.exists():

From 341f19e5363ea278f4c100bac6a03c7397f05753 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 11:43:11 +0100
Subject: [PATCH 05/15] e2e-tests: fixed condition for Ollama

(cherry picked from commit d65b198ad8c10fd1b86a8b19e605d87d52eb397d)
---
 e2e-tests/src/project_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/e2e-tests/src/project_utils.py b/e2e-tests/src/project_utils.py
index 5dad46f4..2742e5e9 100644
--- a/e2e-tests/src/project_utils.py
+++ b/e2e-tests/src/project_utils.py
@@ -46,7 +46,7 @@ def execute_build(project_dir: Path):
                 r"We will download and install Ollama.",
                 r"Downloading .*ollama.*\.tgz",
                 r"Found datasource of type",
-            ], timeout=5) == 0:
+            ], timeout=5) < 4:
                 with allure.step("Ollama model not found locally, installing..."):
                     child.expect("Ollama model .+ pulled successfully", timeout=900)
             child.expect(r"Build complete\. Processed \d+ datasources\.", timeout=30)

From 414728c55537565ad1738956a9e760ef8a870156 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 11:58:42 +0100
Subject: [PATCH 06/15] e2e-tests: enhance logging and fixed ruff checks

(cherry picked from commit 24e77989eb174a853146bc1a116835cdcf0aa0c5)
---
 e2e-tests/src/project_utils.py | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/e2e-tests/src/project_utils.py b/e2e-tests/src/project_utils.py
index 2742e5e9..ec92a166 100644
--- a/e2e-tests/src/project_utils.py
+++ b/e2e-tests/src/project_utils.py
@@ -39,18 +39,26 @@ def execute_build(project_dir: Path):
     with open(log_file_path, "w") as logfile:
         child = pexpect.spawn("uv run databao build", cwd=project_dir, encoding="utf-8", timeout=900, logfile=logfile)
         try:
-            # Wait for Ollama download/installation with extended timeout
-            if child.expect([
-                r"Ollama model .+ not found locally\.",
-                r"No existing Ollama installation detected",
-                r"We will download and install Ollama.",
-                r"Downloading .*ollama.*\.tgz",
-                r"Found datasource of type",
-            ], timeout=5) < 4:
-                with allure.step("Ollama model not found locally, installing..."):
-                    child.expect("Ollama model .+ pulled successfully", timeout=900)
-            child.expect(r"Build complete\. Processed \d+ datasources\.", timeout=30)
-            child.expect(pexpect.EOF)
+            with allure.step("Checking for Ollama model"):
+                # Wait for Ollama download/installation with extended timeout
+                if (
+                    child.expect(
+                        [
+                            r"Ollama model .+ not found locally\.",
+                            r"No existing Ollama installation detected",
+                            r"We will download and install Ollama.",
+                            r"Downloading .*ollama.*\.tgz",
+                            r"Found datasource of type",
+                        ],
+                        timeout=5,
+                    )
+                    < 4
+                ):
+                    with allure.step("Ollama model not found locally, installing..."):
+                        child.expect("Ollama model .+ pulled successfully", timeout=900)
+            with allure.step("Waiting for build to complete"):
+                child.expect(r"Build complete\. Processed \d+ datasources\.", timeout=30)
+                child.expect(pexpect.EOF)
         finally:
             if log_file_path.exists():
                 allure.attach.file(log_file_path, name="cli.log", attachment_type=allure.attachment_type.TEXT)

From 22fbbf16576c71dae0877a3584d0f0c16fbfcd9c Mon Sep 17 00:00:00 2001
From: Vlad Ivoninskii <25663530+ivoninskii@users.noreply.github.com>
Date: Mon, 16 Mar 2026 08:51:54 -0300
Subject: [PATCH 07/15] Edit UI texts (#49)

* Edit UI texts

* lint: ignore ambiguous unicode character warnings (RUF001) in Ruff configuration

---------

Co-authored-by: ruslan.golov <ruslan.golov@jetbrains.com>
(cherry picked from commit 23191a0244cf44aeddee5e187a6fe8c0002f7e80)
---
 pyproject.toml                                |  4 ++-
 src/databao_cli/ui/app.py                     |  6 ++--
 .../ui/components/datasource_manager.py       | 16 +++++-----
 src/databao_cli/ui/components/sidebar.py      |  2 +-
 src/databao_cli/ui/pages/agent_settings.py    | 24 +++++++--------
 src/databao_cli/ui/pages/chat.py              | 20 ++++++-------
 src/databao_cli/ui/pages/context_settings.py  | 28 ++++++++---------
 src/databao_cli/ui/pages/general_settings.py  | 30 +++++++++----------
 src/databao_cli/ui/pages/welcome.py           | 22 +++++++-------
 src/databao_cli/ui/services/build_service.py  |  2 +-
 src/databao_cli/ui/services/dce_operations.py |  2 +-
 11 files changed, 79 insertions(+), 77 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index c1caea36..272b080b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -110,7 +110,9 @@ select = [
     # Prefer absolute imports over relative ones
     "TID252",
 ]
-ignore = []
+ignore = [
+    "RUF001",  # Allow ambiguous unicode characters (emojis in UI)
+]
 
 [tool.ruff.lint.per-file-ignores]
 "tests/**" = ["S101"]  # assert is standard in pytest
diff --git a/src/databao_cli/ui/app.py b/src/databao_cli/ui/app.py
index cf316d43..6bcf87d8 100644
--- a/src/databao_cli/ui/app.py
+++ b/src/databao_cli/ui/app.py
@@ -365,13 +365,13 @@ def build_navigation() -> None:
     )
     context_settings_page = st.Page(
         render_context_settings_page,
-        title="Context Settings",
+        title="Context settings",
         icon="📊",
         url_path="context-settings",
     )
     agent_settings_page = st.Page(
         render_agent_settings_page,
-        title="Agent Settings",
+        title="Agent settings",
         icon="⚙️",
         url_path="agent-settings",
     )
@@ -390,7 +390,7 @@ def new_chat_action() -> None:
     chat_pages.append(
         st.Page(
             new_chat_action,
-            title="New Chat",
+            title="New chat",
             icon=":material/add:",
             url_path="new-chat",
         )
diff --git a/src/databao_cli/ui/components/datasource_manager.py b/src/databao_cli/ui/components/datasource_manager.py
index 83f28b06..5044592c 100644
--- a/src/databao_cli/ui/components/datasource_manager.py
+++ b/src/databao_cli/ui/components/datasource_manager.py
@@ -31,8 +31,8 @@ def render_datasource_manager(project_dir: Path, *, read_only: bool = False) ->
     """Render the full datasource management UI.
 
     Shows:
-    - "Add New Datasource" section at top (hidden in read-only mode)
-    - List of existing datasources below, each with Save/Verify/Remove buttons
+    - "Add data source" section at top (hidden in read-only mode)
+    - List of existing data sources below, each with Save/Verify/Remove buttons
       (buttons hidden in read-only mode, fields disabled)
 
     Args:
@@ -50,7 +50,7 @@ def render_datasource_manager(project_dir: Path, *, read_only: bool = False) ->
         st.caption("No datasources configured yet.")
         return
 
-    st.subheader("Configured Datasources")
+    st.subheader("Your data sources")
 
     for idx, ds in enumerate(configured):
         _render_existing_datasource(project_dir, ds, idx, read_only=read_only)
@@ -64,8 +64,8 @@ def _get_form_version() -> int:
 
 
 def _render_add_datasource_section(project_dir: Path) -> None:
-    """Render the 'Add New Datasource' section."""
-    st.subheader("Add New Datasource")
+    """Render the 'Add data source' section."""
+    st.subheader("Add data source")
 
     available_types = get_available_datasource_types()
 
@@ -104,7 +104,7 @@ def _render_add_datasource_section(project_dir: Path) -> None:
     col_add, col_verify_new = st.columns(2)
 
     with col_add:
-        if st.button("Add Datasource", key="add_ds_btn", type="primary", use_container_width=True):
+        if st.button("Add datasource", key="add_ds_btn", type="primary", use_container_width=True):
             if not ds_name or not ds_name.strip():
                 st.error("Please provide a datasource name.")
             elif not selected_type:
@@ -119,7 +119,7 @@ def _render_add_datasource_section(project_dir: Path) -> None:
                     logger.exception("Failed to add datasource")
 
     with col_verify_new:
-        if st.button("Verify Connection", key="verify_new_ds_btn", use_container_width=True):
+        if st.button("Verify connection", key="verify_new_ds_btn", use_container_width=True):
             if not ds_name or not ds_name.strip() or not selected_type:
                 st.error("Please provide a datasource name and type first.")
             else:
@@ -230,7 +230,7 @@ def _render_existing_datasource(project_dir: Path, ds: ConfiguredDatasource, idx
 
 
 def _clear_new_datasource_form() -> None:
-    """Reset the 'Add New Datasource' form by incrementing the version counter.
+    """Reset the 'Add data source' form by incrementing the version counter.
 
     This causes all form widgets to get fresh keys on the next rerun,
     so they render with their default (empty) values.
diff --git a/src/databao_cli/ui/components/sidebar.py b/src/databao_cli/ui/components/sidebar.py
index b53f5974..6d1d417a 100644
--- a/src/databao_cli/ui/components/sidebar.py
+++ b/src/databao_cli/ui/components/sidebar.py
@@ -175,7 +175,7 @@ def render_sidebar_chat_content(project: ProjectLayout | None) -> None:
 
     if current_chat_id and current_chat_id in chats:
         chat = chats[current_chat_id]
-        if st.button("🗑️ Remove Chat", use_container_width=True, type="primary"):
+        if st.button("🗑️ Remove chat", use_container_width=True, type="primary"):
             _confirm_delete_chat(current_chat_id, chat.display_title)
 
     st.markdown("---")
diff --git a/src/databao_cli/ui/pages/agent_settings.py b/src/databao_cli/ui/pages/agent_settings.py
index 33bc3f31..c48d4917 100644
--- a/src/databao_cli/ui/pages/agent_settings.py
+++ b/src/databao_cli/ui/pages/agent_settings.py
@@ -1,4 +1,4 @@
-"""Agent Settings page - Executor and LLM configuration."""
+"""Agent settings page - Executor and LLM configuration."""
 
 from __future__ import annotations
 
@@ -13,13 +13,13 @@
 
 
 def render_agent_settings_page(*, auto_apply: bool = False) -> None:
-    """Render the Agent Settings page."""
+    """Render the Agent settings page."""
     if not auto_apply:
-        st.title("Agent Settings")
+        st.title("Agent settings")
         st.markdown("Configure the AI agent behavior and execution engine.")
         st.markdown("---")
 
-    st.subheader("⚙️ Execution Engine")
+    st.subheader("⚙️ Execution engine")
 
     st.markdown(
         """
@@ -38,30 +38,30 @@ def render_agent_settings_page(*, auto_apply: bool = False) -> None:
     )
 
     if selected == "lighthouse":
-        st.warning(
+        st.info(
             """
             **LighthouseExecutor** uses a sophisticated graph-based approach with planning and validation steps.
             Best for complex queries requiring multiple steps.
             """,
-            icon="⚠️",
+            icon="ℹ️",
         )
     elif selected == "react_duckdb":
-        st.warning(
+        st.info(
             """
             **ReactDuckDBExecutor** is experimental.
             It uses a ReAct-style loop optimized for DuckDB queries.
             May be faster for simple queries but less reliable for complex ones.
             """,
-            icon="⚠️",
+            icon="ℹ️",
         )
     elif selected == "dbt":
-        st.warning(
+        st.info(
             """
             **DbtProjectExecutor** is experimental.
             It relies on initialized dbt project and a connected warehouse.
             Should be used in case you need to automate dbt project changes.
             """,
-            icon="⚠️",
+            icon="ℹ️",
         )
     elif selected == "claude_code":
         st.info(
@@ -87,9 +87,9 @@ def render_agent_settings_page(*, auto_apply: bool = False) -> None:
 
     st.markdown("---")
 
-    st.subheader("🤖 Language Model")
+    st.subheader("🤖 Language model")
 
-    st.markdown("Choose the LLM provider and configure its credentials.")
+    st.markdown("Choose an LLM provider and configure its credentials.")
 
     llm: LLMSettings = st.session_state.get("llm_settings", LLMSettings())
 
diff --git a/src/databao_cli/ui/pages/chat.py b/src/databao_cli/ui/pages/chat.py
index d6f1b77a..3fb87344 100644
--- a/src/databao_cli/ui/pages/chat.py
+++ b/src/databao_cli/ui/pages/chat.py
@@ -36,9 +36,9 @@ def render_chat_page() -> None:
     chat = _get_or_create_current_chat()
 
     if chat is None:
-        st.error("No chat session found. Please start a new chat.")
+        st.error("No chat session found. Start a new chat.")
         welcome_page = st.session_state.get("_page_welcome")
-        if welcome_page and st.button("🏠 Go to Home"):
+        if welcome_page and st.button("🏠 Back to home"):
             st.switch_page(welcome_page)
         return
 
@@ -69,7 +69,7 @@ def render_chat_page() -> None:
 
     if not _get_or_create_thread_for_chat(chat, agent):
         _render_chat_sidebar(project)
-        st.error("Failed to create conversation thread")
+        st.error("Failed to create a conversation thread")
         return
 
     _render_chat_sidebar(project)
@@ -192,12 +192,12 @@ def _render_no_project_state() -> None:
     st.title("💬 Chat")
     st.markdown("---")
 
-    st.warning("No DCE project detected.")
+    st.warning("No Context Engine project found.")
 
     st.markdown(
         """
 
-        To use Databao, you need a DCE (Databao Context Engine) project with configured datasources.
+        Databao requires a Context Engine project with configured data sources.
 
         **Set up a new project:**
         ```bash
@@ -215,22 +215,22 @@ def _render_no_project_state() -> None:
 
 
 def _render_no_datasources_state(project: ProjectLayout) -> None:
-    """Render state when no datasources are configured in the DCE project."""
+    """Render state when no data sources are configured in the DCE project."""
     st.title("💬 Chat")
     st.markdown("---")
 
-    st.warning(f"No datasources configured in project at `{project.project_dir}`.")
+    st.warning(f"No data sources configured in project at `{project.project_dir}`.")
 
     st.markdown(
         """
 
-        Add datasources to your project before using Databao.
+        Add data sources to your project before using Databao.
 
-        See the documentation for how to configure datasources.
+        See the documentation for how to configure data sources.
         """
     )
 
-    if st.button("🔄 Check Again"):
+    if st.button("🔄 Check again"):
         st.session_state.databao_project = None
         st.rerun()
 
diff --git a/src/databao_cli/ui/pages/context_settings.py b/src/databao_cli/ui/pages/context_settings.py
index 631ce44c..1a899156 100644
--- a/src/databao_cli/ui/pages/context_settings.py
+++ b/src/databao_cli/ui/pages/context_settings.py
@@ -1,4 +1,4 @@
-"""Context Settings page - DCE project configuration, datasource management, and build."""
+"""Context settings page - DCE project configuration, datasource management, and build."""
 
 import logging
 
@@ -17,13 +17,13 @@
 
 
 def render_context_settings_page() -> None:
-    """Render the Context Settings page."""
-    st.title("Context Settings")
+    """Render the Context settings page."""
+    st.title("Context settings")
     st.markdown("Configure your data context and sources.")
 
     project: ProjectLayout | None = st.session_state.get("databao_project")
 
-    # ---- DCE Status section ----
+    # ---- DCE status section ----
     st.markdown("---")
     st.subheader("📋 Status")
 
@@ -52,9 +52,9 @@ def render_context_settings_page() -> None:
         invalidate_agent("Reloading project...")
         st.rerun()
 
-    # ---- Datasource Management section ----
+    # ---- Data sources section ----
     st.markdown("---")
-    st.subheader("🔗 Datasource Management")
+    st.subheader("🔗 Data sources")
 
     if project is not None:
         status = databao_project_status(project)
@@ -63,29 +63,29 @@ def render_context_settings_page() -> None:
         else:
             render_datasource_manager(project.root_domain_dir, read_only=is_read_only_domain())
     else:
-        st.caption("Configure a project to manage datasources.")
+        st.caption("Configure a project to manage data sources.")
 
     # ---- Build section ----
     st.markdown("---")
-    st.subheader("🔨 Build Context")
+    st.subheader("🔨 Build context")
 
     if project is not None:
         status = databao_project_status(project)
         if status in (DatabaoProjectStatus.NOT_INITIALIZED, DatabaoProjectStatus.NO_DATASOURCES):
-            st.caption("Add at least one datasource before building.")
+            st.caption("Add at least one data source before building context.")
         else:
             if not is_read_only_domain():
                 st.markdown(
-                    "Build indexes your datasources so Databao can understand your data "
+                    "Build creates indexes of your datasources so Databao can understand your data "
                     "structure and answer questions about it."
                 )
             render_build_section(project.root_domain_dir, read_only=is_read_only_domain())
     else:
         st.caption("Configure a project to build context.")
 
-    # ---- Connected Sources (from Agent) ----
+    # ---- Connected sources (from Agent) ----
     st.markdown("---")
-    st.subheader("📡 Active Agent Sources")
+    st.subheader("📡 Active data sources")
 
     agent = st.session_state.get("agent")
     if agent is None:
@@ -109,7 +109,7 @@ def _render_project_info(project: ProjectLayout) -> bool:
     elif status == DatabaoProjectStatus.NOT_INITIALIZED:
         st.error("Project not initialized", icon="❌")
     elif status == DatabaoProjectStatus.NO_DATASOURCES:
-        st.warning("No datasources configured", icon="⚠️")
+        st.warning("No data sources configured", icon="⚠️")
 
     reload_clicked = st.button("🔄 Reload")
 
@@ -122,7 +122,7 @@ def _render_sources(agent: Agent) -> None:
     dfs = agent.dfs
 
     if not dbs and not dfs:
-        st.caption("No sources configured in this project.")
+        st.caption("No data sources configured in this project.")
         return
 
     if dbs:
diff --git a/src/databao_cli/ui/pages/general_settings.py b/src/databao_cli/ui/pages/general_settings.py
index 05928e5a..29d22a2c 100644
--- a/src/databao_cli/ui/pages/general_settings.py
+++ b/src/databao_cli/ui/pages/general_settings.py
@@ -1,4 +1,4 @@
-"""General Settings page - Storage and app-wide configuration."""
+"""General settings page - Storage and app-wide configuration."""
 
 import logging
 
@@ -59,42 +59,42 @@ def _confirm_reset_settings() -> None:
 
 
 def render_general_settings_page() -> None:
-    """Render the General Settings page."""
-    st.title("General Settings")
+    """Render the General settings page."""
+    st.title("General settings")
     st.markdown("Configure application storage and manage data.")
 
     st.markdown("---")
 
-    st.subheader("📁 Storage Location")
+    st.subheader("📁 Storage location")
 
     base_path = get_storage_base_path()
     chats_dir = get_chats_dir()
     cache_dir = get_cache_dir()
 
-    st.markdown("**Base Path** (read-only)")
+    st.markdown("**Base path** (read-only)")
     st.code(str(base_path), language=None)
 
-    with st.expander("📂 Storage Details", expanded=False):
-        st.markdown("**Chats Directory**")
+    with st.expander("📂 Storage details", expanded=False):
+        st.markdown("**Chat directory**")
         st.code(str(chats_dir), language=None)
 
-        st.markdown("**Cache Directory**")
+        st.markdown("**Cache directory**")
         st.code(str(cache_dir), language=None)
 
         try:
             chats = st.session_state.get("chats", {})
             num_chats = len(chats)
-            st.metric("Saved Chats", num_chats)
+            st.metric("Saved chats", num_chats)
         except Exception:
             logger.debug("Failed to display storage statistics", exc_info=True)
 
     st.markdown("---")
 
-    st.subheader("⚠️ Data Management")
+    st.subheader("⚠️ Data management")
 
     st.markdown(
         """
-        These actions affect your stored data. Use with caution.
+        These actions affect your stored data. Use them with caution.
         """
     )
 
@@ -102,18 +102,18 @@ def render_general_settings_page() -> None:
 
     with col1:
         if st.button(
-            "🗑️ Clear All Chats",
+            "🗑️ Clear all chats",
             use_container_width=True,
             type="secondary",
-            help="Removes all chat history and cached data.",
+            help="Delete all chat history and cached data. This cannot be undone.",
         ):
             _confirm_clear_chats()
 
     with col2:
         if st.button(
-            "🔄 Reset to Defaults",
+            "🔄 Reset settings",
             use_container_width=True,
             type="secondary",
-            help="Resets settings but keeps chats.",
+            help="Restore default settings. Your chats will remain unchanged.",
         ):
             _confirm_reset_settings()
diff --git a/src/databao_cli/ui/pages/welcome.py b/src/databao_cli/ui/pages/welcome.py
index 3088c81d..bebc8384 100644
--- a/src/databao_cli/ui/pages/welcome.py
+++ b/src/databao_cli/ui/pages/welcome.py
@@ -59,12 +59,12 @@ def render_welcome_page() -> None:
 
         st.markdown("---")
 
-        st.subheader("Quick Actions")
+        st.subheader("Quick actions")
 
         action_col1, action_col2 = st.columns(2)
 
         with action_col1:
-            if st.button("💬 Start New Chat", width="stretch", type="primary"):
+            if st.button("💬 New chat", width="stretch", type="primary"):
                 _create_new_chat()
                 st.rerun()
 
@@ -75,7 +75,7 @@ def render_welcome_page() -> None:
                     st.switch_page(context_settings_page)
 
         st.markdown("---")
-        st.subheader("Getting Started")
+        st.subheader("Get started")
 
         with st.expander("📖 How to use Databao", expanded=False):
             st.markdown(
@@ -94,7 +94,7 @@ def render_welcome_page() -> None:
 
         if chats:
             st.markdown("---")
-            st.subheader("Recent Chats")
+            st.subheader("Recent chats")
 
             sorted_chats = sorted(
                 chats.values(),
@@ -124,11 +124,11 @@ def render_setup_wizard_page() -> None:
 
     The wizard is organized into up to five sections, which are disabled
     based on prerequisites:
-    1. Initialize Project
-    2. Configure Datasources
-    3. Configure Agent
-    4. Build Context (optional; can be hidden via feature flag)
-    5. Start Using Databao
+    1. Initialize project
+    2. Configure datasources
+    3. Configure agent
+    4. Build context (optional; can be hidden via feature flag)
+    5. Start using databao
 
     When read-only-domain mode is active, editing sections are disabled with
     an explanation banner.
@@ -255,11 +255,11 @@ def render_setup_wizard_page() -> None:
 
         st.markdown("---")
 
-        # ---- Section 4: Build Context ----
+        # ---- Section 4: Build context ----
         if not hide_build_context:
             _render_section_header(
                 "4",
-                "Build Context (Optional)",
+                "Build context (optional)",
                 completed=build_started_or_done,
                 enabled=has_datasources,
             )
diff --git a/src/databao_cli/ui/services/build_service.py b/src/databao_cli/ui/services/build_service.py
index eb34a382..a94d8902 100644
--- a/src/databao_cli/ui/services/build_service.py
+++ b/src/databao_cli/ui/services/build_service.py
@@ -197,7 +197,7 @@ def _build_fragment() -> None:
 
         with col_build:
             if build_status == "not_started":
-                if st.button("Build Context", key="build_btn", type="primary"):
+                if st.button("Build context", key="build_btn", type="primary"):
                     start_build(project_dir)
                     st.rerun(scope="fragment")
             elif build_status == "running":
diff --git a/src/databao_cli/ui/services/dce_operations.py b/src/databao_cli/ui/services/dce_operations.py
index 6c0fd4a7..c62588f7 100644
--- a/src/databao_cli/ui/services/dce_operations.py
+++ b/src/databao_cli/ui/services/dce_operations.py
@@ -89,7 +89,7 @@ def remove_datasource(project_dir: Path, datasource_id: DatasourceId) -> None:
 
 
 def list_datasources(project_dir: Path) -> list[ConfiguredDatasource]:
-    """List all configured datasources in the DCE project (reads from disk)."""
+    """List all configured data sources in the DCE project (reads from disk)."""
     try:
         manager = DatabaoContextDomainManager(domain_dir=project_dir)
         return manager.get_configured_datasource_list()

From 3dfb398291dbe9257431ec3c0d47bdf579cace41 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 13:04:14 +0100
Subject: [PATCH 08/15] refactor: improve type annotations & error handling in
 CLI utils and fix logic condition in datasource config

---
 src/databao_cli/__main__.py                   |  3 +-
 .../datasource/add_datasource_config.py       |  4 +--
 src/databao_cli/labels.py                     |  2 +-
 src/databao_cli/utils.py                      | 31 +++++++++++--------
 tests/test_add_datasource.py                  |  4 +--
 5 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/databao_cli/__main__.py b/src/databao_cli/__main__.py
index 600bb580..49a12ecc 100644
--- a/src/databao_cli/__main__.py
+++ b/src/databao_cli/__main__.py
@@ -4,10 +4,9 @@
 import click
 from click import Context
 
+from databao_cli.labels import LABELS
 from databao_cli.log.logging import configure_logging
 from databao_cli.project.layout import ROOT_DOMAIN, ProjectLayout, find_project
-
-from databao_cli.labels import LABELS
 from databao_cli.utils import ask_confirm, register_labels
 
 
diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add_datasource_config.py
index bff0858e..d4d0706e 100644
--- a/src/databao_cli/commands/datasource/add_datasource_config.py
+++ b/src/databao_cli/commands/datasource/add_datasource_config.py
@@ -6,11 +6,11 @@
     DatabaoContextPluginLoader,
     DatasourceType,
 )
-from databao_cli.utils import ask_confirm, ask_select, ask_text
 
 from databao_cli.commands.context_engine_cli import ClickUserInputCallback
 from databao_cli.commands.datasource.check_datasource_connection import print_connection_check_results
 from databao_cli.project.layout import ProjectLayout
+from databao_cli.utils import ask_confirm, ask_select, ask_text
 
 
 def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain: str) -> None:
@@ -51,7 +51,7 @@ def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) ->
     config_type = ask_select(
         "What type of datasource do you want to add?",
         choices=all_datasource_types,
-        default=all_datasource_types[0] if len(all_datasource_types) == 0 else None,
+        default=all_datasource_types[0] if len(all_datasource_types) > 0 else None,
     )
 
     return DatasourceType(full_type=config_type)
diff --git a/src/databao_cli/labels.py b/src/databao_cli/labels.py
index 3cefe8af..227b30a9 100644
--- a/src/databao_cli/labels.py
+++ b/src/databao_cli/labels.py
@@ -32,4 +32,4 @@
     "connection.auth.credentials_file": "Credentials file",
     "connection.auth.key_file": "Key file",
     "connection.auth.token": "Token",
-}
\ No newline at end of file
+}
diff --git a/src/databao_cli/utils.py b/src/databao_cli/utils.py
index f0c89f55..f99e3302 100644
--- a/src/databao_cli/utils.py
+++ b/src/databao_cli/utils.py
@@ -1,37 +1,40 @@
 import sys
+from typing import Any
+
 import click
 import questionary
 
 _labels: dict[str, str] = {}
 
-def register_labels(labels: dict) -> None:
+
+def register_labels(labels: dict[str, str]) -> None:
     _labels.update(labels)
 
+
 def _resolve(value: str) -> tuple[str, str]:
     label = _labels.get(value, value)
     return value, label
 
+
 def is_interactive() -> bool:
     """True when running in an interactive terminal."""
     return sys.stdin.isatty() and sys.stdout.isatty()
 
 
-def ask_select(message: str, choices: list[str], default: str = None) -> str:
+def ask_select(message: str, choices: list[str], default: str | None = None) -> str:
     """Select from a list. Interactive in TTY, plain text otherwise."""
     if is_interactive():
         resolved = [_resolve(c) if isinstance(c, str) else c for c in choices]
-        q_choices = [
-            questionary.Choice(title=label, value=value) for value, label in resolved
-        ]
-        result = questionary.select(message, choices=q_choices, default=default).ask()
+        q_choices = [questionary.Choice(title=label, value=value) for value, label in resolved]
+        result: Any = questionary.select(message, choices=q_choices, default=default).ask()
         if result is None:
             raise click.Abort()
-        return result
+        return str(result)
     else:
         click.echo(f"{message}")
         for i, choice in enumerate(choices, 1):
             click.echo(f"  {i}. {choice}")
-        value = click.prompt(
+        value: str = click.prompt(
             "Enter a number of value",
             default=default or choices[0],
         )
@@ -45,19 +48,21 @@ def ask_select(message: str, choices: list[str], default: str = None) -> str:
 def ask_confirm(message: str, default: bool = True, abort: bool = False) -> bool:
     """Yes/no. Fancy in TTY, plain click.confirm otherwise."""
     if is_interactive():
-        result = questionary.confirm(message, default=default).ask()
+        result: Any = questionary.confirm(message, default=default).ask()
         if result is None:
             raise click.Abort()
         if abort and not result:
             raise click.Abort()
-        return result
+        return bool(result)
     else:
         return click.confirm(message, default=default, abort=abort)
 
 
-def ask_text(message: str, default: str = None) -> str:
+def ask_text(message: str, default: str | None = None) -> str:
     """Text input. Interactive in TTY, plain click.prompt otherwise."""
     if is_interactive():
-        return questionary.text(message, default=default or "").ask()
+        result: Any = questionary.text(message, default=default or "").ask()
+        return str(result)
     else:
-        return click.prompt(message, default=default)
+        value: str = click.prompt(message, default=default)
+        return value
diff --git a/tests/test_add_datasource.py b/tests/test_add_datasource.py
index 947e5f71..b07124fb 100644
--- a/tests/test_add_datasource.py
+++ b/tests/test_add_datasource.py
@@ -3,10 +3,10 @@
 
 import duckdb
 import pytest
+import questionary
 from click.testing import CliRunner
 
 from databao_cli.__main__ import cli
-from databao_cli.commands.datasource import add_datasource_config
 from databao_cli.commands.init import init_impl as init_databao_project
 
 
@@ -38,7 +38,7 @@ def test_databao_datasource_add(tmp_path: Path, temp_parquet_file: Path, monkeyp
 
     prompt = Mock()
     prompt.ask.return_value = "parquet"
-    monkeypatch.setattr(add_datasource_config.questionary, "select", Mock(return_value=prompt))
+    monkeypatch.setattr(questionary, "select", Mock(return_value=prompt))
 
     inputs = [
         "resources/my_parq",

From ad2864b8334103bb985bca07313548d8fca361fb Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 13:11:10 +0100
Subject: [PATCH 09/15] refactor: add `allow_empty` option to `ask_text` for
 flexible input handling

---
 src/databao_cli/utils.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/databao_cli/utils.py b/src/databao_cli/utils.py
index f99e3302..510e96fc 100644
--- a/src/databao_cli/utils.py
+++ b/src/databao_cli/utils.py
@@ -58,11 +58,15 @@ def ask_confirm(message: str, default: bool = True, abort: bool = False) -> bool
         return click.confirm(message, default=default, abort=abort)
 
 
-def ask_text(message: str, default: str | None = None) -> str:
+def ask_text(message: str, default: str | None = None, allow_empty: bool = False) -> str:
     """Text input. Interactive in TTY, plain click.prompt otherwise."""
     if is_interactive():
-        result: Any = questionary.text(message, default=default or "").ask()
-        return str(result)
+        while True:
+            result: Any = questionary.text(message, default=default or "").ask()
+            value = str(result) if result is not None else ""
+            if value.strip() or allow_empty:
+                return value
+            click.echo("Value cannot be empty. Please try again.")
     else:
         value: str = click.prompt(message, default=default)
         return value

From df5d02d19a1a1d856f800f934826c754073a5880 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 13:23:00 +0100
Subject: [PATCH 10/15] Add validation handling for interactive datasource
 creation and improve prompt logic for required fields

---
 .../commands/context_engine_cli.py            | 35 ++++++++++++-------
 .../datasource/add_datasource_config.py       | 17 +++++++--
 src/databao_cli/utils.py                      |  3 +-
 3 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/src/databao_cli/commands/context_engine_cli.py b/src/databao_cli/commands/context_engine_cli.py
index 60f5265c..1a1baa9e 100644
--- a/src/databao_cli/commands/context_engine_cli.py
+++ b/src/databao_cli/commands/context_engine_cli.py
@@ -15,19 +15,28 @@ def prompt(
         show_default: bool = default_value is not None and default_value != ""
         final_type = click.Choice(type.choices) if isinstance(type, Choice) else str
 
-        # click goes infinite loop if user gives emptry string as an input AND default_value is None
-        # in order to exit this loop we need to set default value to '' (so it gets accepted)
-        #
-        # Code snippet from click:
-        # while True:
-        #   value = prompt_func(prompt)
-        #     if value:
-        #       break
-        #     elif default is not None:
-        #       value = default
-        #       break
-        default_value = default_value if default_value else "" if final_type is str else None
-        return click.prompt(text=text, default=default_value, hide_input=is_secret, type=final_type, show_default=show_default)
+        is_optional = "(Optional)" in text
+        if default_value:
+            final_default = default_value
+        elif is_optional and final_type is str:
+            final_default = ""
+        elif final_type is str:
+            final_default = None
+        else:
+            final_default = None
+
+        if final_type is str and not is_optional and final_default is None:
+            while True:
+                value = click.prompt(
+                    text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
+                )
+                if value and value.strip():
+                    return value
+                click.echo("This field is required and cannot be empty. Please try again.")
+        else:
+            return click.prompt(
+                text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
+            )
 
     def confirm(self, text: str) -> bool:
         return click.confirm(text=text)
diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add_datasource_config.py
index d4d0706e..b19480ab 100644
--- a/src/databao_cli/commands/datasource/add_datasource_config.py
+++ b/src/databao_cli/commands/datasource/add_datasource_config.py
@@ -6,6 +6,7 @@
     DatabaoContextPluginLoader,
     DatasourceType,
 )
+from pydantic import ValidationError
 
 from databao_cli.commands.context_engine_cli import ClickUserInputCallback
 from databao_cli.commands.datasource.check_datasource_connection import print_connection_check_results
@@ -32,9 +33,19 @@ def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain
             abort=True,
             default=False,
         )
-    created_datasource = domain_manager.create_datasource_config_interactively(
-        datasource_type, datasource_name, ClickUserInputCallback(), overwrite_existing=True
-    )
+
+    while True:
+        try:
+            created_datasource = domain_manager.create_datasource_config_interactively(
+                datasource_type, datasource_name, ClickUserInputCallback(), overwrite_existing=True
+            )
+            break
+        except ValidationError as e:
+            click.echo(click.style("\nValidation error:", fg="red", bold=True))
+            for error in e.errors():
+                field_path = ".".join(str(loc) for loc in error["loc"])
+                click.echo(click.style(f"  • {field_path}: {error['msg']}", fg="red"))
+            click.echo("\nPlease try again with correct values.\n")
 
     datasource_id = created_datasource.datasource.id
     click.echo(
diff --git a/src/databao_cli/utils.py b/src/databao_cli/utils.py
index 510e96fc..f71c6d9e 100644
--- a/src/databao_cli/utils.py
+++ b/src/databao_cli/utils.py
@@ -68,5 +68,4 @@ def ask_text(message: str, default: str | None = None, allow_empty: bool = False
                 return value
             click.echo("Value cannot be empty. Please try again.")
     else:
-        value: str = click.prompt(message, default=default)
-        return value
+        return str(click.prompt(message, default=default))

From 041a864d600ccd627885d12e3db3e1246276d9c1 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 13:29:05 +0100
Subject: [PATCH 11/15] Add interactive choice handling and improve prompt
 logic for optional inputs in context engine CLI

---
 .../commands/context_engine_cli.py            | 73 +++++++++++++++----
 .../datasource/add_datasource_config.py       |  2 +-
 2 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/src/databao_cli/commands/context_engine_cli.py b/src/databao_cli/commands/context_engine_cli.py
index 1a1baa9e..7d3fb007 100644
--- a/src/databao_cli/commands/context_engine_cli.py
+++ b/src/databao_cli/commands/context_engine_cli.py
@@ -1,6 +1,8 @@
+import sys
 from typing import Any
 
 import click
+import questionary
 from databao_context_engine import Choice, UserInputCallback
 
 
@@ -15,28 +17,73 @@ def prompt(
         show_default: bool = default_value is not None and default_value != ""
         final_type = click.Choice(type.choices) if isinstance(type, Choice) else str
 
+        # Determine if this field is optional
         is_optional = "(Optional)" in text
+
+        if isinstance(type, Choice):
+            is_interactive = sys.stdin.isatty() and sys.stdout.isatty()
+            if is_interactive:
+                from databao_cli.labels import LABELS
+
+                choices = [
+                    questionary.Choice(title=LABELS.get(choice, choice), value=choice) for choice in type.choices
+                ]
+                result = questionary.select(
+                    text, choices=choices, default=default_value if default_value in type.choices else None
+                ).ask()
+                if result is None:
+                    raise click.Abort()
+                return result
+            else:
+                return click.prompt(
+                    text=text,
+                    default=default_value,
+                    hide_input=is_secret,
+                    type=click.Choice(type.choices),
+                    show_default=show_default,
+                )
+
         if default_value:
             final_default = default_value
-        elif is_optional and final_type is str:
+        elif is_optional:
             final_default = ""
-        elif final_type is str:
-            final_default = None
         else:
             final_default = None
 
-        if final_type is str and not is_optional and final_default is None:
-            while True:
-                value = click.prompt(
+        is_interactive = sys.stdin.isatty() and sys.stdout.isatty()
+        if is_interactive and final_type is str:
+            if is_secret:
+                prompt_func = questionary.password
+            else:
+                prompt_func = questionary.text
+
+            if not is_optional and final_default is None:
+                while True:
+                    result = prompt_func(text, default=final_default or "").ask()
+                    if result is None:
+                        raise click.Abort()
+                    value = str(result).strip()
+                    if value:
+                        return value
+                    click.echo("This field is required and cannot be empty. Please try again.")
+            else:
+                result = prompt_func(text, default=final_default or "").ask()
+                if result is None:
+                    raise click.Abort()
+                return str(result)
+        else:
+            if final_type is str and not is_optional and final_default is None:
+                while True:
+                    value = click.prompt(
+                        text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
+                    )
+                    if value and value.strip():
+                        return value
+                    click.echo("This field is required and cannot be empty. Please try again.")
+            else:
+                return click.prompt(
                     text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
                 )
-                if value and value.strip():
-                    return value
-                click.echo("This field is required and cannot be empty. Please try again.")
-        else:
-            return click.prompt(
-                text=text, default=final_default, hide_input=is_secret, type=final_type, show_default=show_default
-            )
 
     def confirm(self, text: str) -> bool:
         return click.confirm(text=text)
diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add_datasource_config.py
index b19480ab..ad6c1c96 100644
--- a/src/databao_cli/commands/datasource/add_datasource_config.py
+++ b/src/databao_cli/commands/datasource/add_datasource_config.py
@@ -23,7 +23,7 @@ def add_datasource_config_interactive_impl(project_layout: ProjectLayout, domain
 
     datasource_type = _ask_for_datasource_type(plugin_loader.get_all_supported_datasource_types(exclude_file_plugins=True))
 
-    datasource_name = ask_text("Datasource name")
+    datasource_name = ask_text("Datasource name?")
 
     datasource_id = domain_manager.datasource_config_exists(datasource_name=datasource_name)
     if datasource_id is not None:

From 260f4d0f69082d86ae8923de25f72d9292a08e4b Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 14:11:24 +0100
Subject: [PATCH 12/15] Fixed ruff and mypy checks

---
 src/databao_cli/commands/context_engine_cli.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/src/databao_cli/commands/context_engine_cli.py b/src/databao_cli/commands/context_engine_cli.py
index 7d3fb007..72c41c88 100644
--- a/src/databao_cli/commands/context_engine_cli.py
+++ b/src/databao_cli/commands/context_engine_cli.py
@@ -25,11 +25,11 @@ def prompt(
             if is_interactive:
                 from databao_cli.labels import LABELS
 
-                choices = [
-                    questionary.Choice(title=LABELS.get(choice, choice), value=choice) for choice in type.choices
-                ]
+                choices = [questionary.Choice(title=LABELS.get(choice, choice), value=choice) for choice in type.choices]
                 result = questionary.select(
-                    text, choices=choices, default=default_value if default_value in type.choices else None
+                    text,
+                    choices=choices,
+                    default=default_value if default_value is not None and default_value in type.choices else None,
                 ).ask()
                 if result is None:
                     raise click.Abort()
@@ -52,10 +52,7 @@ def prompt(
 
         is_interactive = sys.stdin.isatty() and sys.stdout.isatty()
         if is_interactive and final_type is str:
-            if is_secret:
-                prompt_func = questionary.password
-            else:
-                prompt_func = questionary.text
+            prompt_func = questionary.password if is_secret else questionary.text
 
             if not is_optional and final_default is None:
                 while True:

From 8d17a7e35975a7c2c1c6ff84e2f361d3187be1a6 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 14:29:43 +0100
Subject: [PATCH 13/15] Handle `KeyboardInterrupt` in `ask_text` to gracefully
 abort input prompts

---
 src/databao_cli/utils.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/databao_cli/utils.py b/src/databao_cli/utils.py
index f71c6d9e..96cf49e6 100644
--- a/src/databao_cli/utils.py
+++ b/src/databao_cli/utils.py
@@ -62,8 +62,13 @@ def ask_text(message: str, default: str | None = None, allow_empty: bool = False
     """Text input. Interactive in TTY, plain click.prompt otherwise."""
     if is_interactive():
         while True:
-            result: Any = questionary.text(message, default=default or "").ask()
-            value = str(result) if result is not None else ""
+            try:
+                result: Any = questionary.text(message, default=default or "").ask()
+            except KeyboardInterrupt:
+                raise click.Abort()
+            if result is None:
+                raise click.Abort()
+            value = str(result)
             if value.strip() or allow_empty:
                 return value
             click.echo("Value cannot be empty. Please try again.")

From 8cee4669f67e86fe9f1af67c01e324df75f1da61 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 14:48:13 +0100
Subject: [PATCH 14/15] Set default value to `None` for datasource type
 selection

---
 src/databao_cli/commands/datasource/add_datasource_config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/databao_cli/commands/datasource/add_datasource_config.py b/src/databao_cli/commands/datasource/add_datasource_config.py
index ad6c1c96..7974c665 100644
--- a/src/databao_cli/commands/datasource/add_datasource_config.py
+++ b/src/databao_cli/commands/datasource/add_datasource_config.py
@@ -62,7 +62,7 @@ def _ask_for_datasource_type(supported_datasource_types: set[DatasourceType]) ->
     config_type = ask_select(
         "What type of datasource do you want to add?",
         choices=all_datasource_types,
-        default=all_datasource_types[0] if len(all_datasource_types) > 0 else None,
+        default=None,
     )
 
     return DatasourceType(full_type=config_type)

From d72b43b1f98e1c61ceffffd49e314d4116546512 Mon Sep 17 00:00:00 2001
From: "ruslan.golov" <ruslan.golov@jetbrains.com>
Date: Mon, 16 Mar 2026 16:11:28 +0100
Subject: [PATCH 15/15] Handle `KeyboardInterrupt` in `ask_text` without
 traceback for cleaner abort

---
 src/databao_cli/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/databao_cli/utils.py b/src/databao_cli/utils.py
index 96cf49e6..cd319c05 100644
--- a/src/databao_cli/utils.py
+++ b/src/databao_cli/utils.py
@@ -65,7 +65,7 @@ def ask_text(message: str, default: str | None = None, allow_empty: bool = False
             try:
                 result: Any = questionary.text(message, default=default or "").ask()
             except KeyboardInterrupt:
-                raise click.Abort()
+                raise click.Abort() from None
             if result is None:
                 raise click.Abort()
             value = str(result)