databao-context-engine 0.7.0__tar.gz → 0.7.1.dev1__tar.gz

{databao_context_engine-0.7.0 → databao_context_engine-0.7.1.dev1}/PKG-INFO

@@ -1,10 +1,9 @@
 Metadata-Version: 2.4
 Name: databao-context-engine
-Version: 0.7.0
+Version: 0.7.1.dev1
 Summary: Semantic context for your LLMs — generated automatically
 License-Expression: Apache-2.0 AND LicenseRef-Additional-Terms
 License-File: LICENSE.md
-Requires-Dist: click>=8.3.0
 Requires-Dist: duckdb>=1.4.3
 Requires-Dist: pyyaml>=6.0.3
 Requires-Dist: requests>=2.32.5
@@ -23,7 +22,7 @@ Requires-Dist: docling>=2.70.0 ; extra == 'pdf'
 Requires-Dist: asyncpg>=0.31.0 ; extra == 'postgresql'
 Requires-Dist: snowflake-connector-python>=4.2.0 ; extra == 'snowflake'
 Requires-Python: >=3.11
-Project-URL: Homepage, https://databao.app/
+Project-URL: Homepage, https://jetbrains.com/databao/
 Project-URL: Source, https://github.com/JetBrains/databao-context-engine
 Provides-Extra: athena
 Provides-Extra: bigquery
@@ -48,10 +47,10 @@ Description-Content-Type: text/markdown
  No more copying schemas. No manual documentation. Just accurate answers.
 </p>
 <p align="center"> 
-  <a href="https://databao.app">Website</a> •
+  <a href="https://jetbrains.com/databao">Website</a> •
   <a href="#quickstart">Quickstart</a> •
-  <a href="#supported-data-sources">Data Sources</a> •
-  <a href="#contributing">Contributing</a>
+ <a href="https://docs.databao.app">Docs</a> •
+ <a href="https://discord.gg/hEUqCcWdVh">Discord</a>
 </p>
 
 ---
@@ -124,22 +123,26 @@ pip install databao-context-engine
 ### 1. Create a domain
 
 ```python
-# Initialize the domain in an existing directory
+# Initialize the domain in a temporary directory
 from databao_context_engine import init_dce_domain
+from pathlib import Path
+import tempfile
 
 domain_manager = init_dce_domain(Path(tempfile.mkdtemp()))
 
 # Or use an existing project
 from databao_context_engine import DatabaoContextDomainManager
 
-domain_manager = DatabaoContextDomainManager(domain_dir=Path("path/to/project"))
+domain_manager = DatabaoContextDomainManager(domain_dir=Path("domain_dir"))
 ```
 
 ### 2. Configure data sources
 
 ```python
 from databao_context_engine import (
+    CheckDatasourceConnectionResult,
     DatasourceConnectionStatus,
+    DatasourceId,
     DatasourceType,
 )
 
@@ -153,11 +156,10 @@ postgres_datasource_id = domain_manager.create_datasource_config(
 ).datasource.id
 
 # Check the connection to the datasource is valid
-check_result = domain_manager.check_datasource_connection()
+check_result: dict[DatasourceId, CheckDatasourceConnectionResult] = domain_manager.check_datasource_connection()
 
 assert len(check_result) == 1
-assert check_result[0].datasource_id == postgres_datasource_id
-assert check_result[0].connection_status == DatasourceConnectionStatus.VALID
+assert check_result[postgres_datasource_id].connection_status == DatasourceConnectionStatus.VALID
 ```
 
 ### 3. Build context
@@ -236,6 +238,7 @@ Apache 2.0 — use it however you want. See the [LICENSE](LICENSE.md) file for d
 </p>
 
 <p align="center">
- <a href="https://databao.app">Website</a> •
+ <a href="https://jetbrains.com/databao">Website</a> •
+ <a href="https://docs.databao.app">Docs</a> •
  <a href="https://discord.gg/hEUqCcWdVh">Discord</a>
 </p>

{databao_context_engine-0.7.0 → databao_context_engine-0.7.1.dev1}/README.md

@@ -11,10 +11,10 @@
  No more copying schemas. No manual documentation. Just accurate answers.
 </p>
 <p align="center"> 
-  <a href="https://databao.app">Website</a> •
+  <a href="https://jetbrains.com/databao">Website</a> •
   <a href="#quickstart">Quickstart</a> •
-  <a href="#supported-data-sources">Data Sources</a> •
-  <a href="#contributing">Contributing</a>
+ <a href="https://docs.databao.app">Docs</a> •
+ <a href="https://discord.gg/hEUqCcWdVh">Discord</a>
 </p>
 
 ---
@@ -87,22 +87,26 @@ pip install databao-context-engine
 ### 1. Create a domain
 
 ```python
-# Initialize the domain in an existing directory
+# Initialize the domain in a temporary directory
 from databao_context_engine import init_dce_domain
+from pathlib import Path
+import tempfile
 
 domain_manager = init_dce_domain(Path(tempfile.mkdtemp()))
 
 # Or use an existing project
 from databao_context_engine import DatabaoContextDomainManager
 
-domain_manager = DatabaoContextDomainManager(domain_dir=Path("path/to/project"))
+domain_manager = DatabaoContextDomainManager(domain_dir=Path("domain_dir"))
 ```
 
 ### 2. Configure data sources
 
 ```python
 from databao_context_engine import (
+    CheckDatasourceConnectionResult,
     DatasourceConnectionStatus,
+    DatasourceId,
     DatasourceType,
 )
 
@@ -116,11 +120,10 @@ postgres_datasource_id = domain_manager.create_datasource_config(
 ).datasource.id
 
 # Check the connection to the datasource is valid
-check_result = domain_manager.check_datasource_connection()
+check_result: dict[DatasourceId, CheckDatasourceConnectionResult] = domain_manager.check_datasource_connection()
 
 assert len(check_result) == 1
-assert check_result[0].datasource_id == postgres_datasource_id
-assert check_result[0].connection_status == DatasourceConnectionStatus.VALID
+assert check_result[postgres_datasource_id].connection_status == DatasourceConnectionStatus.VALID
 ```
 
 ### 3. Build context
@@ -199,6 +202,7 @@ Apache 2.0 — use it however you want. See the [LICENSE](LICENSE.md) file for d
 </p>
 
 <p align="center">
- <a href="https://databao.app">Website</a> •
+ <a href="https://jetbrains.com/databao">Website</a> •
+ <a href="https://docs.databao.app">Docs</a> •
  <a href="https://discord.gg/hEUqCcWdVh">Discord</a>
 </p>

{databao_context_engine-0.7.0 → databao_context_engine-0.7.1.dev1}/pyproject.toml

@@ -1,13 +1,12 @@
 [project]
 name = "databao-context-engine"
-version = "0.7.0"
+version = "0.7.1.dev1"
 description = "Semantic context for your LLMs — generated automatically"
 readme = "README.md"
 requires-python = ">=3.11"
 license = "Apache-2.0 AND LicenseRef-Additional-Terms"
 license-files = ["LICENSE.md"]
 dependencies = [
-    "click>=8.3.0",
     "duckdb>=1.4.3",
     "pyyaml>=6.0.3",
     "requests>=2.32.5",
@@ -24,7 +23,7 @@ default-optional-dependency-keys = [
 ]
 
 [project.urls]
-Homepage = "https://databao.app/"
+Homepage = "https://jetbrains.com/databao/"
 Source = "https://github.com/JetBrains/databao-context-engine"
 
 [project.optional-dependencies]
@@ -59,6 +58,7 @@ build-backend = "uv_build"
 
 [dependency-groups]
 dev = [
+    "click>=8.3.0",
     "mypy>=1.18.2",
     "pre-commit>=4.5.1",
     "pytest>=9.0.2",

{databao_context_engine-0.7.0 → databao_context_engine-0.7.1.dev1}/src/databao_context_engine/__init__.py

@@ -41,6 +41,11 @@ from databao_context_engine.plugins.databases.clickhouse.config_file import (
     ClickhouseConfigFile,
     ClickhouseConnectionProperties,
 )
+from databao_context_engine.plugins.databases.database_context_explorer import (
+    DatabaseSchemaLite,
+    DatabaseTableDetails,
+    DatabaseTableLite,
+)
 from databao_context_engine.plugins.databases.duckdb.config_file import DuckDBConfigFile, DuckDBConnectionConfig
 from databao_context_engine.plugins.databases.mssql.config_file import MSSQLConfigFile, MSSQLConnectionProperties
 from databao_context_engine.plugins.databases.mysql.config_file import MySQLConfigFile, MySQLConnectionProperties
@@ -138,4 +143,7 @@ __all__ = [
     "DbtContextFilterStructuredRule",
     "DbtContextFilterRule",
     "ParquetConfigFile",
+    "DatabaseSchemaLite",
+    "DatabaseTableLite",
+    "DatabaseTableDetails",
 ]

{databao_context_engine-0.7.0 → databao_context_engine-0.7.1.dev1}/src/databao_context_engine/build_sources/build_runner.py

@@ -18,12 +18,11 @@ from databao_context_engine.build_sources.types import (
 from databao_context_engine.datasources.datasource_context import (
     DatasourceContext,
     hash_context_file,
-    read_datasource_type_from_context,
 )
 from databao_context_engine.datasources.datasource_discovery import discover_datasources, prepare_source
-from databao_context_engine.datasources.types import PreparedConfig, PreparedDatasource
+from databao_context_engine.datasources.types import DatasourceId, PreparedConfig, PreparedDatasource
 from databao_context_engine.pluginlib.build_plugin import DatasourceType
-from databao_context_engine.plugins.plugin_loader import DatabaoContextPluginLoader
+from databao_context_engine.plugins.plugin_loader import NoPluginFoundForDatasource
 from databao_context_engine.progress.progress import ProgressCallback, ProgressEmitter, ProgressStep
 from databao_context_engine.project.layout import ProjectLayout
 
@@ -44,6 +43,21 @@ def _build_step_plan(*, should_index: bool, should_enrich_context: bool) -> tupl
     return tuple(steps)
 
 
+def _emit_all_build_step_as_completed(
+    *,
+    progress: ProgressCallback | None,
+    datasource_id: str | DatasourceId,
+    should_index: bool,
+    should_enrich_context: bool,
+) -> None:
+    emitter = ProgressEmitter(progress)
+    for step in _build_step_plan(should_index=should_index, should_enrich_context=should_enrich_context):
+        emitter.datasource_step_completed(
+            datasource_id=str(datasource_id),
+            step=step,
+        )
+
+
 @perf.perf_run(
     operation="build",
     attrs=lambda *, should_index, should_enrich_context, **_: {
@@ -55,8 +69,8 @@ def _build_step_plan(*, should_index: bool, should_enrich_context: bool) -> tupl
 def build(
     *,
     project_layout: ProjectLayout,
-    plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
+    datasource_ids: list[DatasourceId] | None,
     should_index: bool,
     should_enrich_context: bool,
     progress: ProgressCallback | None = None,
@@ -73,7 +87,8 @@ def build(
     Returns:
         A list of per-datasource build results.
     """
-    datasource_ids = discover_datasources(project_layout)
+    if not datasource_ids:
+        datasource_ids = discover_datasources(project_layout)
 
     emitter = ProgressEmitter(progress)
 
@@ -98,7 +113,6 @@ def build(
         try:
             result = _build_one_datasource(
                 project_layout=project_layout,
-                plugin_loader=plugin_loader,
                 build_service=build_service,
                 datasource_id=datasource_id,
                 should_index=should_index,
@@ -151,7 +165,6 @@ def build(
 def _build_one_datasource(
     *,
     project_layout: ProjectLayout,
-    plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
     datasource_id,
     should_index: bool,
@@ -169,15 +182,6 @@ def _build_one_datasource(
         f'Found datasource of type "{prepared_source.datasource_type.full_type}" with name {prepared_source.datasource_id.datasource_path}'
     )
 
-    plugin = plugin_loader.get_plugin_for_datasource_type(prepared_source.datasource_type)
-    if plugin is None:
-        logger.warning(
-            "No plugin for '%s' (datasource=%s) — skipping.",
-            prepared_source.datasource_type.full_type,
-            prepared_source.datasource_id.relative_path_to_config_file(),
-        )
-        return BuildDatasourceResult(datasource_id=datasource_id, status=DatasourceStatus.SKIPPED)
-
     ProgressEmitter(progress).datasource_step_plan_set(
         datasource_id=str(datasource_id),
         step_plan=_build_step_plan(
@@ -186,33 +190,48 @@ def _build_one_datasource(
         ),
     )
 
-    result = build_service.build_context(
-        prepared_source=prepared_source,
-        plugin=plugin,
-        progress=progress,
-    )
+    try:
+        result = build_service.build_context(
+            prepared_source=prepared_source,
+            progress=progress,
+        )
 
-    if should_enrich_context:
-        result = build_service.enrich_built_context(result, plugin, progress)
+        if should_enrich_context:
+            result = build_service.enrich_built_context(built_context=result, progress=progress)
 
-    output_dir = project_layout.output_dir
-    context_file_path = export_build_result(output_dir, result)
+        output_dir = project_layout.output_dir
+        context_file_path = export_build_result(output_dir, result)
 
-    perf.set_attribute("context_size_bytes", context_file_path.stat().st_size)
+        perf.set_attribute("context_size_bytes", context_file_path.stat().st_size)
 
-    if should_index:
-        context_hash = hash_context_file(datasource_id=prepared_source.datasource_id, context_path=context_file_path)
-        build_service.index_built_context(
-            built_context=result, plugin=plugin, context_hash=context_hash, progress=progress
+        if should_index:
+            context_hash = hash_context_file(
+                datasource_id=prepared_source.datasource_id, context_path=context_file_path
+            )
+            build_service.index_built_context(built_context=result, context_hash=context_hash, progress=progress)
+
+        return BuildDatasourceResult(
+            datasource_id=datasource_id,
+            status=DatasourceStatus.OK,
+            datasource_type=DatasourceType(full_type=result.datasource_type),
+            context_built_at=datetime.now(),
+            context_file_path=context_file_path,
         )
-
-    return BuildDatasourceResult(
-        datasource_id=datasource_id,
-        status=DatasourceStatus.OK,
-        datasource_type=DatasourceType(full_type=result.datasource_type),
-        context_built_at=datetime.now(),
-        context_file_path=context_file_path,
-    )
+    except NoPluginFoundForDatasource as e:
+        logger.warning(
+            "No plugin for '%s' (datasource=%s) — skipping.",
+            e.datasource_type.full_type,
+            prepared_source.datasource_id.relative_path_to_config_file(),
+        )
+        # Since the plugin was not found, no build steps were emitted but the plan was set:
+        # we need to emit all steps as completed
+        _emit_all_build_step_as_completed(
+            progress=progress,
+            datasource_id=datasource_id,
+            should_index=should_index,
+            should_enrich_context=should_enrich_context,
+        )
+        return BuildDatasourceResult(datasource_id=datasource_id, status=DatasourceStatus.SKIPPED)
 
 
 def _is_datasource_enabled(prepared_source: PreparedDatasource) -> bool:
@@ -233,7 +252,6 @@ def _is_datasource_enabled(prepared_source: PreparedDatasource) -> bool:
 def run_enrich_context(
     *,
     project_layout: ProjectLayout,
-    plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
     contexts: list[DatasourceContext],
     should_index: bool,
@@ -250,7 +268,6 @@ def run_enrich_context(
             result = _enrich_one_context(
                 project_layout=project_layout,
                 context=context,
-                plugin_loader=plugin_loader,
                 build_service=build_service,
                 should_index=should_index,
             )
@@ -286,46 +303,40 @@ def _enrich_one_context(
     *,
     project_layout: ProjectLayout,
     context: DatasourceContext,
-    plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
     should_index: bool,
 ) -> EnrichContextResult:
     perf.set_attribute("context_size_bytes", len(context.context.encode("utf-8")))
 
-    datasource_type = read_datasource_type_from_context(context)
-    perf.set_attribute("datasource_type", getattr(datasource_type, "full_type", datasource_type))
+    try:
+        enriched_context = build_service.enrich_datasource_context(context=context)
 
-    plugin = plugin_loader.get_plugin_for_datasource_type(datasource_type)
-    if plugin is None:
+        output_dir = project_layout.output_dir
+        context_file_path = export_build_result(output_dir, enriched_context)
+
+        if should_index:
+            context_hash = hash_context_file(datasource_id=context.datasource_id, context_path=context_file_path)
+            build_service.index_built_context(built_context=enriched_context, context_hash=context_hash)
+
+        return EnrichContextResult(
+            datasource_id=context.datasource_id,
+            status=DatasourceStatus.OK,
+            context_file_path=context_file_path,
+        )
+    except NoPluginFoundForDatasource as e:
         logger.warning(
             "No plugin for datasource type '%s' — skipping context enrichment for datasource %s.",
-            getattr(datasource_type, "full_type", datasource_type),
+            e.datasource_type.full_type,
             context.datasource_id,
         )
         return EnrichContextResult(datasource_id=context.datasource_id, status=DatasourceStatus.SKIPPED)
 
-    enriched_context = build_service.enrich_datasource_context(context=context, plugin=plugin)
-
-    output_dir = project_layout.output_dir
-    context_file_path = export_build_result(output_dir, enriched_context)
-
-    if should_index:
-        context_hash = hash_context_file(datasource_id=context.datasource_id, context_path=context_file_path)
-        build_service.index_built_context(built_context=enriched_context, plugin=plugin, context_hash=context_hash)
-
-    return EnrichContextResult(
-        datasource_id=context.datasource_id,
-        status=DatasourceStatus.OK,
-        context_file_path=context_file_path,
-    )
-
 
 @perf.perf_run(operation="index")
 @perf.perf_span("index.total")
 def run_indexing(
     *,
     project_layout: ProjectLayout,
-    plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
     contexts: list[DatasourceContext],
     progress: ProgressCallback | None = None,
@@ -356,9 +367,7 @@ def run_indexing(
         try:
             logger.info(f"Indexing datasource {context.datasource_id}")
 
-            result = _index_one_context(
-                context=context, plugin_loader=plugin_loader, build_service=build_service, progress=progress
-            )
+            result = _index_one_context(context=context, build_service=build_service, progress=progress)
 
             results.append(result)
             if result.status == DatasourceStatus.OK:
@@ -406,28 +415,25 @@ def run_indexing(
 def _index_one_context(
     *,
     context: DatasourceContext,
-    plugin_loader: DatabaoContextPluginLoader,
     build_service: BuildService,
     progress: ProgressCallback | None = None,
 ) -> IndexDatasourceResult:
     perf.set_attribute("context_size_bytes", len(context.context.encode("utf-8")))
 
-    datasource_type = read_datasource_type_from_context(context)
-    perf.set_attribute("datasource_type", getattr(datasource_type, "full_type", datasource_type))
-
-    plugin = plugin_loader.get_plugin_for_datasource_type(datasource_type)
-    if plugin is None:
+    ProgressEmitter(progress).datasource_step_plan_set(
+        datasource_id=str(context.datasource_id),
+        step_plan=BuildService.index_step_plan(),
+    )
+    try:
+        build_service.index_datasource_context(context=context, progress=progress)
+        return IndexDatasourceResult(datasource_id=context.datasource_id, status=DatasourceStatus.OK)
+    except NoPluginFoundForDatasource as e:
         logger.warning(
             "No plugin for datasource type '%s' — skipping indexing for %s.",
-            getattr(datasource_type, "full_type", datasource_type),
+            e.datasource_type.full_type,
             context.datasource_id,
         )
+        # Since the plugin was not found, no index steps were emitted but the plan was set:
+        # we need to emit all steps as completed
+        BuildService.emit_all_index_step_as_completed(progress=progress, datasource_id=context.datasource_id)
         return IndexDatasourceResult(datasource_id=context.datasource_id, status=DatasourceStatus.SKIPPED)
-
-    ProgressEmitter(progress).datasource_step_plan_set(
-        datasource_id=str(context.datasource_id),
-        step_plan=BuildService.index_step_plan(),
-    )
-
-    build_service.index_datasource_context(context=context, plugin=plugin, progress=progress)
-    return IndexDatasourceResult(datasource_id=context.datasource_id, status=DatasourceStatus.OK)