databao-agent 0.1.4.dev4__tar.gz → 0.1.4.dev8__tar.gz

{databao_agent-0.1.4.dev4 → databao_agent-0.1.4.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databao-agent
-Version: 0.1.4.dev4
+Version: 0.1.4.dev8
 Summary: databao-agent: NL queries for data
 Project-URL: Homepage, https://databao.app/
 Project-URL: Source, https://github.com/JetBrains/databao-agent
@@ -9,7 +9,8 @@ License-File: LICENSE.md
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: <3.15,>=3.11
-Requires-Dist: databao-context-engine[mysql,postgresql,snowflake]~=0.5.0
+Requires-Dist: claude-agent-sdk>=0.1.48
+Requires-Dist: databao-context-engine[mysql,postgresql,snowflake]~=0.6.0
 Requires-Dist: dbt-core~=1.9.0
 Requires-Dist: dbt-duckdb>=1.10.0
 Requires-Dist: diskcache>=5.6.3
@@ -137,19 +138,21 @@ engine = create_engine(
 ### 2. Create a Databao agent and register sources
 
 ```python
-import databao
-from databao import LLMConfig
+import databao.agent as bao
 
 # Option A - Local: install and run any compatible local LLM
 # For list of compatible models, see "Local Models" below
-# llm_config = LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
+# llm_config = bao.LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
 
 # Option B - Cloud (requires an API key, e.g. OPENAI_API_KEY)
-llm_config = LLMConfig(name="gpt-4o-mini", temperature=0)
-agent = databao.new_agent(name="demo", llm_config=llm_config)
+llm_config = bao.LLMConfig(name="gpt-4o-mini", temperature=0)
 
 # Add your database to the agent
-agent.add_db(engine)
+domain = bao.domain()
+domain.add_db(engine)
+
+agent = bao.agent(domain, name="demo", llm_config=llm_config)
+
 ```
 
 ### 3. Ask questions and materialize results
@@ -201,17 +204,17 @@ Databao agent works great with local LLMs — your data never leaves your machin
 ### Ollama
 
 1. Install [Ollama](https://ollama.com/download) for your OS and make sure it’s running
-2. Use an `LLMConfig` with `name` of the form `"ollama:<model_name>"`:
+2. Use a `bao.LLMConfig` with `name` of the form `"ollama:<model_name>"`:
 
    ```python
-   llm_config = LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
+   llm_config = bao.LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
    ```
 
    The model will be downloaded automatically if it doesn't exist. Or run `ollama pull <model_name>` to download manually.
 
 ### OpenAI-compatible servers
 
-You can use any OpenAI-compatible server by setting `api_base_url` in the `LLMConfig`.
+You can use any OpenAI-compatible server by setting `api_base_url` in the `bao.LLMConfig`.
 
 For an example, see `examples/configs/qwen3-8b-oai.yaml`.
 

{databao_agent-0.1.4.dev4 → databao_agent-0.1.4.dev8}/README.md

@@ -93,19 +93,21 @@ engine = create_engine(
 ### 2. Create a Databao agent and register sources
 
 ```python
-import databao
-from databao import LLMConfig
+import databao.agent as bao
 
 # Option A - Local: install and run any compatible local LLM
 # For list of compatible models, see "Local Models" below
-# llm_config = LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
+# llm_config = bao.LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
 
 # Option B - Cloud (requires an API key, e.g. OPENAI_API_KEY)
-llm_config = LLMConfig(name="gpt-4o-mini", temperature=0)
-agent = databao.new_agent(name="demo", llm_config=llm_config)
+llm_config = bao.LLMConfig(name="gpt-4o-mini", temperature=0)
 
 # Add your database to the agent
-agent.add_db(engine)
+domain = bao.domain()
+domain.add_db(engine)
+
+agent = bao.agent(domain, name="demo", llm_config=llm_config)
+
 ```
 
 ### 3. Ask questions and materialize results
@@ -157,17 +159,17 @@ Databao agent works great with local LLMs — your data never leaves your machin
 ### Ollama
 
 1. Install [Ollama](https://ollama.com/download) for your OS and make sure it’s running
-2. Use an `LLMConfig` with `name` of the form `"ollama:<model_name>"`:
+2. Use a `bao.LLMConfig` with `name` of the form `"ollama:<model_name>"`:
 
    ```python
-   llm_config = LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
+   llm_config = bao.LLMConfig(name="ollama:gpt-oss:20b", temperature=0)
    ```
 
    The model will be downloaded automatically if it doesn't exist. Or run `ollama pull <model_name>` to download manually.
 
 ### OpenAI-compatible servers
 
-You can use any OpenAI-compatible server by setting `api_base_url` in the `LLMConfig`.
+You can use any OpenAI-compatible server by setting `api_base_url` in the `bao.LLMConfig`.
 
 For an example, see `examples/configs/qwen3-8b-oai.yaml`.
 

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/__init__.py

@@ -1,14 +1,14 @@
 import importlib.metadata
 
 try:
-    __version__ = importlib.metadata.version(__name__)
+    __version__ = importlib.metadata.version("databao-agent")
 except importlib.metadata.PackageNotFoundError:
     __version__ = "0.0.0"  # Fallback for development mode
 
 
-from databao.api import agent, domain
-from databao.configs.llm import LLMConfig
-from databao.core import (
+from databao.agent.api import agent, domain
+from databao.agent.configs.llm import LLMConfig
+from databao.agent.core import (
     Agent,
     Domain,
     ExecutionResult,
@@ -18,7 +18,7 @@ from databao.core import (
     VisualisationResult,
     Visualizer,
 )
-from databao.databases import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases import DBConnection, DBConnectionConfig, DBConnectionRuntime
 
 __all__ = [
     "Agent",

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/api.py

@@ -3,16 +3,19 @@ from typing import TextIO
 
 from typing_extensions import deprecated
 
-from databao.caches.in_mem_cache import InMemCache
-from databao.configs.agent import DEFAULT_AGENT_CONFIG, AgentConfig
-from databao.configs.llm import LLMConfig, LLMConfigDirectory
-from databao.core import Agent, Cache, Executor, Visualizer
-from databao.core.domain import Domain, _DCEProjectDomain, _InMemoryDomain
-from databao.executors import ReactDuckDBExecutor
-from databao.executors.dbt.config import DbtConfig
-from databao.executors.dbt.executor import DbtProjectExecutor
-from databao.executors.lighthouse.executor import LighthouseExecutor
-from databao.visualizers.vega_chat import VegaChatVisualizer
+from databao.agent.caches.in_mem_cache import InMemCache
+from databao.agent.configs.agent import DEFAULT_AGENT_CONFIG, AgentConfig
+from databao.agent.configs.llm import LLMConfig, LLMConfigDirectory
+from databao.agent.core import Agent, Cache, Executor, Visualizer
+from databao.agent.core.domain import Domain, _DCEProjectDomain, _InMemoryDomain
+from databao.agent.executors import (
+    ClaudeCodeExecutor,
+    DbtProjectExecutor,
+    LighthouseExecutor,
+    ReactDuckDBExecutor,
+)
+from databao.agent.executors.dbt.config import DbtConfig
+from databao.agent.visualizers.vega_chat import VegaChatVisualizer
 
 
 def agent(
@@ -49,6 +52,8 @@ def agent(
                 data_executor = DbtProjectExecutor(dbt_config=dbt_config, writer=writer)
             case "react_duckdb":
                 data_executor = ReactDuckDBExecutor(writer=writer)
+            case "claude":
+                data_executor = ClaudeCodeExecutor(writer=writer)
             case _:
                 raise ValueError(f"Invalid executor type: {executor_type}")
 

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/caches/disk_cache.py

@@ -6,7 +6,7 @@ from typing import Any
 
 import diskcache  # type: ignore[import-untyped]
 
-from databao.core import Cache
+from databao.agent.core import Cache
 
 
 @dataclass(kw_only=True)

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/caches/in_mem_cache.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from databao.core.cache import Cache
+from databao.agent.core.cache import Cache
 
 
 class InMemCache(Cache):

databao_agent-0.1.4.dev8/databao/agent/configs/__init__.py

@@ -0,0 +1,3 @@
+from databao.agent.configs.llm import LLMConfig, LLMConfigDirectory
+
+__all__ = ["LLMConfig", "LLMConfigDirectory"]

databao_agent-0.1.4.dev8/databao/agent/core/__init__.py

@@ -0,0 +1,20 @@
+from databao.agent.core.agent import Agent
+from databao.agent.core.cache import Cache
+from databao.agent.core.domain import Domain
+from databao.agent.core.executor import ExecutionResult, Executor
+from databao.agent.core.opa import Opa
+from databao.agent.core.thread import Thread
+from databao.agent.core.visualizer import HistoryMode, VisualisationResult, Visualizer
+
+__all__ = [
+    "Agent",
+    "Cache",
+    "Domain",
+    "ExecutionResult",
+    "Executor",
+    "HistoryMode",
+    "Opa",
+    "Thread",
+    "VisualisationResult",
+    "Visualizer",
+]

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/agent.py

@@ -7,18 +7,18 @@ from langchain_core.language_models.chat_models import BaseChatModel
 from pandas import DataFrame
 from typing_extensions import deprecated
 
-from databao.core.data_source import DBDataSource, DFDataSource, Sources
-from databao.core.domain import Domain, _Domain
-from databao.core.thread import Thread
-from databao.databases import DBConnection
-from databao.mcp.manager import McpManager
+from databao.agent.core.data_source import DBDataSource, DFDataSource, Sources
+from databao.agent.core.domain import Domain, _Domain
+from databao.agent.core.thread import Thread
+from databao.agent.databases import DBConnection
+from databao.agent.mcp.manager import McpManager
 
 if TYPE_CHECKING:
-    from databao.configs.agent import AgentConfig
-    from databao.configs.llm import LLMConfig
-    from databao.core.cache import Cache
-    from databao.core.executor import Executor
-    from databao.core.visualizer import Visualizer
+    from databao.agent.configs.agent import AgentConfig
+    from databao.agent.configs.llm import LLMConfig
+    from databao.agent.core.cache import Cache
+    from databao.agent.core.executor import Executor
+    from databao.agent.core.visualizer import Visualizer
 
 logger = logging.getLogger(__name__)
 
@@ -55,7 +55,7 @@ class Agent:
         self.__llm_config = llm
         self.__agent_config = agent_config
 
-        self.__executor = data_executor
+        self.__executor: Executor = data_executor
         self.__visualizer = visualizer
         self.__cache = cache
         self.__mcp: McpManager = McpManager()
@@ -67,16 +67,6 @@ class Agent:
         self.__stream_ask = stream_ask
         self.__stream_plot = stream_plot
 
-        self._init_executor()
-
-    def _init_executor(self) -> None:
-        for db_source in self.sources.dbs.values():
-            self.executor.register_db(db_source)
-        for df_source in self.sources.dfs.values():
-            self.executor.register_df(df_source)
-        for dbt_source in self.sources.dbts.values():
-            self.executor.register_dbt(dbt_source)
-
     @deprecated("Use Domain.add_db() and initialize Agent with Domain instead.")
     def add_db(self, conn: DBConnection, *, name: str | None = None, context: str | Path | None = None) -> None:
         raise NotImplementedError(

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/data_source.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 import pandas as pd
 
-from databao.databases import DBConnectionConfig
+from databao.agent.databases import DBConnectionConfig
 
 
 @dataclass

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/domain.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from abc import ABC
 from pathlib import Path
 from typing import Any, Protocol
@@ -13,9 +14,9 @@ from databao_context_engine import (
 )
 from pandas import DataFrame
 
-from databao.core.data_source import DataSource, DBDataSource, DBTDataSource, DFDataSource, Sources
-from databao.core.sources import SourcesManager
-from databao.databases import (
+from databao.agent.core.data_source import DataSource, DBDataSource, DBTDataSource, DFDataSource, Sources
+from databao.agent.core.sources import SourcesManager
+from databao.agent.databases import (
     DBConnection,
     DBConnectionConfig,
     DBConnectionRuntime,
@@ -23,8 +24,19 @@ from databao.databases import (
     create_db_config_from_runtime,
     try_create_db_config_from_content,
 )
-from databao.dbt import create_dbt_config_file, try_extract_dbt_dir_from_content
-from databao.integrations.dce import DatabaoContextApi
+from databao.agent.dbt import create_dbt_config_file, try_extract_dbt_dir_from_content
+from databao.agent.integrations.dce import DatabaoContextApi
+
+_SOURCE_NAME_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$")
+
+
+def _verify_source_name(name: str | None) -> None:
+    if name is not None and re.match(_SOURCE_NAME_PATTERN, name) is None:
+        # The provided source name should be simple enough for LLMs to use it.
+        # Otherwise, the agent would always have to correctly quote the provided name in SQL queries.
+        raise ValueError(
+            f"Invalid database name: '{name}'. The provided source name must match {_SOURCE_NAME_PATTERN.pattern}"
+        )
 
 
 class Domain(Protocol):
@@ -137,6 +149,7 @@ class _Domain(ABC, Domain):
     def _add_db(
         self, db: DBConnection, name: str | None = None, description: str | Path | None = None
     ) -> DBDataSource | None:
+        _verify_source_name(name)
         if isinstance(db, DBConnectionConfig):
             db_config = db
         elif isinstance(db, DBConnectionRuntime):
@@ -148,11 +161,13 @@ class _Domain(ABC, Domain):
     def _add_df(
         self, df: DataFrame, name: str | None = None, description: str | Path | None = None
     ) -> DFDataSource | None:
+        _verify_source_name(name)
         return self._sources_manager.add_df(df, name=name, description=description)
 
     def _add_dbt(
         self, dbt: str | Path, name: str | None = None, description: str | Path | None = None
     ) -> DBTDataSource | None:
+        _verify_source_name(name)
         if isinstance(dbt, str):
             dbt = Path(dbt)
         return self._sources_manager.add_dbt(dbt, name=name, description=description)
@@ -214,11 +229,7 @@ class _DCEProjectDomain(_Domain):
         return self._dce.search_context(retrieve_text, datasource_ids=datasource_ids)
 
     def is_context_built(self) -> bool:
-        if self._dce.is_context_built():
-            return True
-        # NOTE: (@gas) dbt datasources are configured in DCE but may not appear in the
-        # introspected list; treat them as "built" when present.
-        return len(self._configured) > 0
+        return self._dce.is_context_built()
 
     @property
     def supports_context(self) -> bool:

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/executor.py

@@ -1,20 +1,22 @@
 import base64
+import logging
 import re
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, TextIO
 
-from langchain_core.tools import BaseTool
 from pandas import DataFrame
 from pydantic import BaseModel, ConfigDict
 
-from databao.core.data_source import DBDataSource, DBTDataSource, DFDataSource
-from databao.core.domain import Domain
+from databao.agent.core.domain import Domain
 
 if TYPE_CHECKING:
-    from databao import LLMConfig
-    from databao.configs.agent import AgentConfig
-    from databao.core.cache import Cache
-    from databao.core.opa import Opa
+    from databao.agent import LLMConfig
+    from databao.agent.configs.agent import AgentConfig
+    from databao.agent.core.cache import Cache
+    from databao.agent.core.opa import Opa
+
+
+logger = logging.getLogger(__name__)
 
 
 class OutputModalityHints(BaseModel):
@@ -146,20 +148,8 @@ class Executor(ABC):
     """
 
     @abstractmethod
-    def register_db(self, source: DBDataSource) -> None:
-        pass
-
-    @abstractmethod
-    def register_df(self, source: DFDataSource) -> None:
-        pass
-
-    @abstractmethod
-    def register_dbt(self, source: DBTDataSource) -> None:
-        pass
-
-    @abstractmethod
-    def register_tools(self, tools: list[BaseTool]) -> None:
-        """Register additional LangChain tools to be available during execution."""
+    def register_tools(self, tools: list[Any]) -> None:
+        """Register additional tools to be available during execution."""
 
     @abstractmethod
     def drop_last_opa_group(self, cache: "Cache", n: int = 1) -> None:
@@ -196,7 +186,7 @@ class Executor(ABC):
     @staticmethod
     def prepare_for_execution(domain: "Domain") -> None:
         if domain.supports_context and not domain.is_context_built():
-            print(
+            logger.warning(
                 "Context has not been built yet. Building it now — this may take a while. "
                 "To avoid this delay, call domain.build_context() before starting execution."
             )

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/sources.py

@@ -2,8 +2,8 @@ from pathlib import Path
 
 from pandas import DataFrame
 
-from databao.core.data_source import DBDataSource, DBTDataSource, DFDataSource, Sources
-from databao.databases import DBConnectionConfig
+from databao.agent.core.data_source import DBDataSource, DBTDataSource, DFDataSource, Sources
+from databao.agent.databases import DBConnectionConfig
 
 
 class SourcesManager:

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/thread.py

@@ -4,15 +4,14 @@ from typing import TYPE_CHECKING, Any, Self, TextIO
 
 from pandas import DataFrame
 
-from databao.core.env import in_jupyter_kernel
-from databao.core.executor import ExecutionResult, OutputModalityHints
-from databao.core.opa import Opa
-from databao.core.visualizer import VisualisationResult
+from databao.agent.core.env import in_jupyter_kernel
+from databao.agent.core.executor import ExecutionResult, OutputModalityHints
+from databao.agent.core.opa import Opa
+from databao.agent.core.visualizer import VisualisationResult
 
 if TYPE_CHECKING:
-    from databao.core.agent import Agent
-    from databao.core.visualizer import VisualisationResult
-    from databao.multimodal.jupyter_widget import MultimodalWidget
+    from databao.agent.core.agent import Agent
+    from databao.agent.multimodal.jupyter_widget import MultimodalWidget
 
 
 class Thread:
@@ -198,11 +197,11 @@ class Thread:
             ValueError: If visualization generation fails.
         """
         if in_jupyter_kernel():
-            from databao.multimodal import create_jupyter_widget
+            from databao.agent.multimodal import create_jupyter_widget
 
             return create_jupyter_widget(self)
 
-        from databao.multimodal import open_html_content
+        from databao.agent.multimodal import open_html_content
 
         return open_html_content(self)
 

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/core/visualizer.py

@@ -5,7 +5,7 @@ from typing import Any, ClassVar, Literal
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from databao.core.executor import ExecutionResult
+from databao.agent.core.executor import ExecutionResult
 
 _logger = logging.getLogger(__name__)
 

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/__init__.py

@@ -7,8 +7,8 @@ from databao_context_engine import (
 )
 from databao_context_engine.plugins.databases.bigquery.config_file import BigQueryConnectionProperties
 
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
-from databao.databases.databases import (
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.databases import (
     create_db_config_file,
     create_db_config_from_runtime,
     register_db_in_duckdb,

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/bigquery_adapter.py

@@ -12,8 +12,8 @@ from databao_context_engine.plugins.databases.bigquery.config_file import (
 )
 from sqlalchemy import Connection, Engine, make_url
 
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
 
 PROJECT_KEY = "project"
 DATASET_KEY = "dataset"

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/database_adapter.py

@@ -5,7 +5,7 @@ from _duckdb import DuckDBPyConnection
 from databao_context_engine import DatasourceType
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
 
-from databao.databases.database_connection import (
+from databao.agent.databases.database_connection import (
     DBConnection,
     DBConnectionConfig,
     DBConnectionRuntime,

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/databases.py

@@ -3,17 +3,17 @@ from typing import Any
 from _duckdb import DuckDBPyConnection
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile, DatasourceType
 
-from databao.databases.bigquery_adapter import BigQueryAdapter
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import (
+from databao.agent.databases.bigquery_adapter import BigQueryAdapter
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import (
     DBConnectionConfig,
     DBConnectionRuntime,
 )
-from databao.databases.duckdb_adapter import DuckDBAdapter
-from databao.databases.mysql_adapter import MySQLAdapter
-from databao.databases.postgresql_adapter import PostgreSQLAdapter
-from databao.databases.snowflake_adapter import SnowflakeAdapter
-from databao.databases.sqlite_adapter import SQLiteAdapter
+from databao.agent.databases.duckdb_adapter import DuckDBAdapter
+from databao.agent.databases.mysql_adapter import MySQLAdapter
+from databao.agent.databases.postgresql_adapter import PostgreSQLAdapter
+from databao.agent.databases.snowflake_adapter import SnowflakeAdapter
+from databao.agent.databases.sqlite_adapter import SQLiteAdapter
 
 DATABASE_ADAPTERS: list[DatabaseAdapter] = [
     BigQueryAdapter(),

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/duckdb_adapter.py

@@ -4,8 +4,8 @@ from _duckdb import DuckDBPyConnection
 from databao_context_engine import DatasourceType, DuckDBConfigFile, DuckDBConnectionConfig
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
 
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
 
 
 class DuckDBAdapter(DatabaseAdapter):

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/mysql_adapter.py

@@ -5,9 +5,9 @@ from databao_context_engine import DatasourceType, MySQLConfigFile, MySQLConnect
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
 from sqlalchemy import URL, Connection, Engine, make_url
 
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
-from databao.databases.utils import str_dict
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.utils import str_dict
 
 USER_KEY = "user"
 PASSWORD_KEY = "password"

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/postgresql_adapter.py

@@ -5,8 +5,8 @@ from databao_context_engine import DatasourceType, PostgresConfigFile, PostgresC
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
 from sqlalchemy import URL, Connection, Engine, make_url
 
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
 
 USER_KEY = "user"
 PASSWORD_KEY = "password"

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/snowflake_adapter.py

@@ -11,10 +11,11 @@ from databao_context_engine import (
     SnowflakeSSOAuth,
 )
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
+from snowflake.connector.network import SNOWFLAKE_HOST_SUFFIX
 from sqlalchemy import Connection, Engine, make_url
 
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
 
 WAREHOUSE_KEY = "warehouse"
 ACCOUNT_KEY = "account"
@@ -92,8 +93,13 @@ class SnowflakeAdapter(DatabaseAdapter):
         if "dbname" in content:
             content[DATABASE_KEY] = content.pop("dbname")
 
+        host: str | None = content.pop("host", None)
+        account: str = content.get(ACCOUNT_KEY, "")
+        if host and host.endswith(SNOWFLAKE_HOST_SUFFIX):
+            account = host[: -len(SNOWFLAKE_HOST_SUFFIX)]
+
         return SnowflakeConnectionProperties(
-            account=content.get(ACCOUNT_KEY),
+            account=account,
             warehouse=content.get(WAREHOUSE_KEY),
             database=content.get(DATABASE_KEY),
             user=content.get(USER_KEY),
@@ -127,8 +133,9 @@ class SnowflakeAdapter(DatabaseAdapter):
     def _create_secret_sql(config: SnowflakeConnectionProperties, name: str) -> str:
         params: dict[str, str] = {
             ACCOUNT_KEY: config.account,
-            USER_KEY: config.user,
         }
+        if config.user:
+            params[USER_KEY] = config.user
         if config.database:
             params[DATABASE_KEY] = config.database
         if config.warehouse:

{databao_agent-0.1.4.dev4/databao → databao_agent-0.1.4.dev8/databao/agent}/databases/sqlite_adapter.py

@@ -5,8 +5,8 @@ from databao_context_engine import DatasourceType, SQLiteConfigFile, SQLiteConne
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
 from sqlalchemy import Connection, Engine, make_url
 
-from databao.databases.database_adapter import DatabaseAdapter
-from databao.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
+from databao.agent.databases.database_adapter import DatabaseAdapter
+from databao.agent.databases.database_connection import DBConnection, DBConnectionConfig, DBConnectionRuntime
 
 
 class SQLiteAdapter(DatabaseAdapter):

databao_agent-0.1.4.dev8/databao/agent/dbt/__init__.py

@@ -0,0 +1,3 @@
+from databao.agent.dbt.dbt import create_dbt_config_file, try_extract_dbt_dir_from_content
+
+__all__ = ["create_dbt_config_file", "try_extract_dbt_dir_from_content"]