databao-context-engine 0.1.4.dev1__py3-none-any.whl → 0.1.6__py3-none-any.whl

databao_context_engine/llm/api.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+
+from databao_context_engine.llm.errors import OllamaError
+from databao_context_engine.llm.factory import (
+    DEFAULT_DESCRIPTION_GENERATOR_MODEL,
+    DEFAULT_EMBED_MODEL_ID,
+    create_ollama_service,
+)
+from databao_context_engine.llm.install import resolve_ollama_bin
+
+
+def install_ollama_if_needed() -> Path:
+    """Install the Ollama CLI locally if needed.
+
+    This will look for any existing installation of Ollama on the system. If none is found, it will install it locally.
+
+    Here is the priority order of how it looks for an installed Ollama CLI binary:
+    1. Look at the path defined in the DCE_OLLAMA_BIN env var, if it is set
+    2. Look for `ollama` in the PATH
+    3. Look for a DCE-managed installation in the global DCE path
+
+    If Ollama is not found, it will get installed as a DCE-managed installation in the global DCE path.
+
+    Returns:
+        The path to the Ollama CLI executable.
+    """
+    return Path(resolve_ollama_bin())
+
+
+def download_ollama_models_if_needed(
+    *, download_embed_model: bool = True, download_description_generator_model: bool = False
+) -> None:
+    """Download the Ollama models required to run DCE if needed.
+
+    If the models were already downloaded, this method will do nothing.
+
+    If no Ollama CLI is found on the system, this method will install one as a DCE-managed installation in the global DCE path.
+
+    Args:
+        download_embed_model: Whether to download the embedding model.
+        download_description_generator_model: Whether to download the description generator model.
+
+    Raises:
+        OllamaError: If there is an error downloading one of the models.
+    """
+    ollama_service = create_ollama_service()
+
+    if download_embed_model:
+        try:
+            ollama_service.pull_model_if_needed(model=DEFAULT_EMBED_MODEL_ID)
+        except OllamaError as e:
+            raise e
+    if download_description_generator_model:
+        try:
+            ollama_service.pull_model_if_needed(model=DEFAULT_DESCRIPTION_GENERATOR_MODEL)
+        except OllamaError as e:
+            raise e

databao_context_engine/llm/descriptions/ollama.py

@@ -16,6 +16,4 @@ class OllamaDescriptionProvider(DescriptionProvider):
         return self._model_id
 
     def describe(self, text: str, context: str) -> str:
-        description = self._service.describe(model=self._model_id, text=text, context=context)
-
-        return description
+        return self._service.describe(model=self._model_id, text=text, context=context)

databao_context_engine/llm/factory.py

@@ -5,6 +5,9 @@ from databao_context_engine.llm.install import resolve_ollama_bin
 from databao_context_engine.llm.runtime import OllamaRuntime
 from databao_context_engine.llm.service import OllamaService
 
+DEFAULT_EMBED_MODEL_ID = "nomic-embed-text:v1.5"
+DEFAULT_DESCRIPTION_GENERATOR_MODEL = "llama3.2:1b"
+
 
 def _create_ollama_service_common(
     *,
@@ -39,7 +42,7 @@ def create_ollama_service(
 def create_ollama_embedding_provider(
     service: OllamaService,
     *,
-    model_id: str = "nomic-embed-text:v1.5",
+    model_id: str = DEFAULT_EMBED_MODEL_ID,
     dim: int = 768,
     pull_if_needed: bool = True,
 ) -> OllamaEmbeddingProvider:
@@ -52,7 +55,7 @@ def create_ollama_embedding_provider(
 def create_ollama_description_provider(
     service: OllamaService,
     *,
-    model_id: str = "llama3.2:1b",
+    model_id: str = DEFAULT_DESCRIPTION_GENERATOR_MODEL,
     pull_if_needed: bool = True,
 ):
     if pull_if_needed:

databao_context_engine/llm/install.py

@@ -95,16 +95,18 @@ def _detect_platform() -> str:
     raise RuntimeError(f"Unsupported OS/arch: os={os_name!r} arch={arch!r}")
 
 
-def _download_to_temp(url: str) -> Path:
+def _download_artifact_to_temp(artifact_version: str, artifact_name: str) -> Path:
     """Download to a temporary file and return its path."""
     import urllib.request
 
+    artifact_url = f"https://github.com/ollama/ollama/releases/download/{artifact_version}/{artifact_name}"
+
     tmp_dir = Path(tempfile.mkdtemp(prefix="ollama-download-"))
-    file_name = url.rsplit("/", 1)[-1]
+    file_name = artifact_url.rsplit("/", 1)[-1]
     dest = tmp_dir / file_name
 
-    logger.info("Downloading %s to %s", url, dest)
-    with urllib.request.urlopen(url) as resp, dest.open("wb") as out:
+    logger.info("Downloading %s to %s", artifact_url, dest)
+    with urllib.request.urlopen(artifact_url) as resp, dest.open("wb") as out:
         shutil.copyfileobj(resp, out)
 
     return dest
@@ -128,10 +130,12 @@ def _extract_archive(archive: Path, target_dir: Path) -> None:
 
     if name.endswith(".zip"):
         with ZipFile(archive, "r") as zf:
-            zf.extractall(target_dir)
+            # There is no built-in protection against zip bombs in ZipFile.
+            # However, we previously checked the sha256 of the downloaded archive and we trust the origin (GitHub repo of Ollama)
+            zf.extractall(target_dir)  # noqa: S202
     elif name.endswith(".tgz") or name.endswith(".tar.gz"):
         with tarfile.open(archive, "r:gz") as tf:
-            tf.extractall(target_dir)
+            tf.extractall(target_dir, filter="data")
     else:
         raise RuntimeError(f"Unsupported archive format: {archive}")
 
@@ -142,7 +146,7 @@ def _ensure_executable(path: Path) -> None:
         mode = path.stat().st_mode
         path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
     except Exception:
-        pass
+        logger.debug("Failed to mark %s as executable", path, exc_info=True, stack_info=True)
 
 
 def install_ollama_to(target: Path) -> None:
@@ -172,8 +176,7 @@ def install_ollama_to(target: Path) -> None:
     except KeyError as e:
         raise RuntimeError(f"Unsupported platform: {platform_key}") from e
 
-    url = f"https://github.com/ollama/ollama/releases/download/{DEFAULT_VERSION}/{artifact.name}"
-    archive_path = _download_to_temp(url)
+    archive_path = _download_artifact_to_temp(DEFAULT_VERSION, artifact.name)
 
     try:
         _verify_sha256(archive_path, artifact.sha256)
@@ -217,4 +220,4 @@ def install_ollama_to(target: Path) -> None:
         try:
             archive_path.unlink(missing_ok=True)
         except Exception:
-            pass
+            logger.debug("Failed to remove temporary archive %s", archive_path, exc_info=True, stack_info=True)

databao_context_engine/llm/runtime.py

@@ -26,7 +26,7 @@ class OllamaRuntime:
 
         stdout = subprocess.DEVNULL
 
-        proc = subprocess.Popen(
+        return subprocess.Popen(  # noqa: S603 We're always running Ollama
             cmd,
             cwd=str(self._config.work_dir) if self._config.work_dir else None,
             env=env,
@@ -36,8 +36,6 @@ class OllamaRuntime:
             close_fds=os.name != "nt",
         )
 
-        return proc
-
     def start_and_await(
         self,
         *,
@@ -62,11 +60,11 @@ class OllamaRuntime:
             try:
                 proc.terminate()
             except Exception:
-                pass
+                logger.debug("Failed to terminate Ollama server", exc_info=True, stack_info=True)
             try:
                 proc.kill()
             except Exception:
-                pass
+                logger.debug("Failed to kill Ollama server", exc_info=True, stack_info=True)
 
         raise TimeoutError(
             f"Timed out waiting for Ollama to become healthy at http://{self._config.host}:{self._config.port}"

databao_context_engine/mcp/mcp_server.py

@@ -47,9 +47,7 @@ class McpServer:
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def retrieve_tool(text: str, limit: int | None):
-            retrieve_results = self._databao_context_engine.search_context(
-                retrieve_text=text, limit=limit, export_to_file=False
-            )
+            retrieve_results = self._databao_context_engine.search_context(retrieve_text=text, limit=limit)
 
             display_results = [context_search_result.context_result for context_search_result in retrieve_results]
 

databao_context_engine/plugin_loader.py

@@ -38,8 +38,8 @@ class DatabaoContextPluginLoader:
                 for (datasource_type, plugin) in self._all_plugins_by_type.items()
                 if not isinstance(plugin, BuildFilePlugin)
             }
-        else:
-            return set(self._all_plugins_by_type.keys())
+
+        return set(self._all_plugins_by_type.keys())
 
     def get_plugin_for_datasource_type(self, datasource_type: DatasourceType) -> BuildPlugin:
         """Return the plugin able to build a context for the given datasource type.
@@ -103,9 +103,8 @@ class DatabaoContextPluginLoader:
 
         if isinstance(plugin, CustomiseConfigProperties):
             return plugin.get_config_file_properties()
-        elif isinstance(plugin, BuildDatasourcePlugin):
+        if isinstance(plugin, BuildDatasourcePlugin):
             return get_property_list_from_type(plugin.config_file_type)
-        else:
-            raise ValueError(
-                f'Impossible to create a config for datasource type "{datasource_type.full_type}". The corresponding plugin is a {type(plugin).__name__} but should be a BuildDatasourcePlugin or CustomiseConfigProperties'
-            )
+        raise ValueError(
+            f'Impossible to create a config for datasource type "{datasource_type.full_type}". The corresponding plugin is a {type(plugin).__name__} but should be a BuildDatasourcePlugin or CustomiseConfigProperties'
+        )

databao_context_engine/pluginlib/build_plugin.py

@@ -88,36 +88,3 @@ class DatasourceType:
     """
 
     full_type: str
-
-    def __post_init__(self):
-        type_segments = self.full_type.split("/")
-        if len(type_segments) != 2:
-            raise ValueError(f"Invalid DatasourceType: {self.full_type}")
-
-    @property
-    def main_type(self) -> str:
-        """The main type of the datasource, aka the folder in which the config or raw file is located."""
-        return self.full_type.split("/")[0]
-
-    @property
-    def config_folder(self) -> str:
-        """The folder in which the config or raw file is located. This is equivalent to `main_type`."""
-        return self.main_type
-
-    @property
-    def subtype(self) -> str:
-        """The subtype of the datasource. This is the actual type declared in the config file or the raw file's extension."""
-        return self.full_type.split("/")[1]
-
-    @staticmethod
-    def from_main_and_subtypes(main_type: str, subtype: str) -> "DatasourceType":
-        """Create a DatasourceType from its main type and subtype.
-
-        Args:
-            main_type: The main type (aka config folder) of the datasource.
-            subtype: The subtype of the datasource.
-
-        Returns:
-            A DatasourceType instance with the specified main type and subtype.
-        """
-        return DatasourceType(full_type=f"{main_type}/{subtype}")

databao_context_engine/plugins/{athena_db_plugin.py → databases/athena/athena_db_plugin.py}

@@ -1,11 +1,11 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.athena_introspector import AthenaConfigFile, AthenaIntrospector
+from databao_context_engine.plugins.databases.athena.athena_introspector import AthenaConfigFile, AthenaIntrospector
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
 
 
 class AthenaDbPlugin(BaseDatabasePlugin[AthenaConfigFile]):
     id = "jetbrains/athena"
     name = "Athena DB Plugin"
-    supported = {"databases/athena"}
+    supported = {"athena"}
     config_file_type = AthenaConfigFile
 
     def __init__(self):

databao_context_engine/plugins/databases/{athena_introspector.py → athena/athena_introspector.py}

@@ -7,7 +7,7 @@ from pyathena.cursor import DictCursor
 from pydantic import BaseModel, Field
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabaseConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
 from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
@@ -67,7 +67,7 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
     }
     supports_catalogs = True
 
-    def _connect(self, file_config: AthenaConfigFile):
+    def _connect(self, file_config: AthenaConfigFile, *, catalog: str | None = None) -> Any:
         return connect(**file_config.connection.to_athena_kwargs(), cursor_class=DictCursor)
 
     def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
@@ -79,9 +79,6 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
         catalog = file_config.connection.catalog or self._resolve_pseudo_catalog_name(file_config)
         return [catalog]
 
-    def _connect_to_catalog(self, file_config: AthenaConfigFile, catalog: str):
-        return self._connect(file_config)
-
     def _sql_list_schemas(self, catalogs: list[str] | None) -> SQLQuery:
         if not catalogs:
             return SQLQuery("SELECT schema_name, catalog_name FROM information_schema.schemata", None)

databao_context_engine/plugins/{base_db_plugin.py → databases/base_db_plugin.py}

@@ -37,9 +37,7 @@ class BaseDatabasePlugin(BuildDatasourcePlugin[T]):
         return self.supported
 
     def build_context(self, full_type: str, datasource_name: str, file_config: T) -> Any:
-        introspection_result = self._introspector.introspect_database(file_config)
-
-        return introspection_result
+        return self._introspector.introspect_database(file_config)
 
     def check_connection(self, full_type: str, datasource_name: str, file_config: T) -> None:
         self._introspector.check_connection(file_config)

databao_context_engine/plugins/databases/base_introspector.py

@@ -40,7 +40,7 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
 
         discovered_schemas_per_catalog: dict[str, list[str]] = {}
         for catalog in catalogs:
-            with self._connect_to_catalog(file_config, catalog) as conn:
+            with self._connect(file_config, catalog=catalog) as conn:
                 discovered_schemas_per_catalog[catalog] = self._list_schemas_for_catalog(conn, catalog)
         scope = scope_matcher.filter_scopes(catalogs, discovered_schemas_per_catalog)
 
@@ -50,7 +50,7 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
             if not schemas_to_introspect:
                 continue
 
-            with self._connect_to_catalog(file_config, catalog) as catalog_connection:
+            with self._connect(file_config, catalog=catalog) as catalog_connection:
                 introspected_schemas = self.collect_catalog_model(catalog_connection, catalog, schemas_to_introspect)
 
                 if not introspected_schemas:
@@ -74,9 +74,9 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
         if self.supports_catalogs:
             sql = "SELECT catalog_name, schema_name FROM information_schema.schemata WHERE catalog_name = ANY(%s)"
             return SQLQuery(sql, (catalogs,))
-        else:
-            sql = "SELECT schema_name FROM information_schema.schemata"
-            return SQLQuery(sql, None)
+
+        sql = "SELECT schema_name FROM information_schema.schemata"
+        return SQLQuery(sql, None)
 
     def _list_schemas_for_catalog(self, connection: Any, catalog: str) -> list[str]:
         sql_query = self._sql_list_schemas([catalog] if self.supports_catalogs else None)
@@ -108,7 +108,12 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
         return samples
 
     @abstractmethod
-    def _connect(self, file_config: T):
+    def _connect(self, file_config: T, *, catalog: str | None = None) -> Any:
+        """Connect to the database.
+
+        If the `catalog` argument is provided, the connection is "scoped" to that catalog. For engines that don’t need a new connection,
+        return a connection with the session set/USE’d to that catalog.
+        """
         raise NotImplementedError
 
     @abstractmethod
@@ -119,14 +124,6 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
     def _get_catalogs(self, connection, file_config: T) -> list[str]:
         raise NotImplementedError
 
-    @abstractmethod
-    def _connect_to_catalog(self, file_config: T, catalog: str):
-        """Return a connection scoped to `catalog`.
-
-        For engines that don’t need a new connection, return a connection with the
-        session set/USE’d to that catalog.
-        """
-
     def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
         raise NotImplementedError
 

databao_context_engine/plugins/{clickhouse_db_plugin.py → databases/clickhouse/clickhouse_db_plugin.py}

@@ -1,5 +1,5 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.clickhouse_introspector import (
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.clickhouse.clickhouse_introspector import (
     ClickhouseConfigFile,
     ClickhouseIntrospector,
 )
@@ -8,7 +8,7 @@ from databao_context_engine.plugins.databases.clickhouse_introspector import (
 class ClickhouseDbPlugin(BaseDatabasePlugin[ClickhouseConfigFile]):
     id = "jetbrains/clickhouse"
     name = "Clickhouse DB Plugin"
-    supported = {"databases/clickhouse"}
+    supported = {"clickhouse"}
     config_file_type = ClickhouseConfigFile
 
     def __init__(self):

databao_context_engine/plugins/databases/{clickhouse_introspector.py → clickhouse/clickhouse_introspector.py}

@@ -6,7 +6,7 @@ import clickhouse_connect
 from pydantic import BaseModel, Field
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabaseConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
 from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
@@ -36,14 +36,11 @@ class ClickhouseIntrospector(BaseIntrospector[ClickhouseConfigFile]):
 
     supports_catalogs = True
 
-    def _connect(self, file_config: ClickhouseConfigFile):
+    def _connect(self, file_config: ClickhouseConfigFile, *, catalog: str | None = None):
         return clickhouse_connect.get_client(
             **file_config.connection.to_clickhouse_kwargs(),
         )
 
-    def _connect_to_catalog(self, file_config: ClickhouseConfigFile, catalog: str):
-        return self._connect(file_config)
-
     def _get_catalogs(self, connection, file_config: ClickhouseConfigFile) -> list[str]:
         return ["clickhouse"]
 

databao_context_engine/plugins/databases/duckdb/duckdb_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.duckdb.duckdb_introspector import DuckDBConfigFile, DuckDBIntrospector
+
+
+class DuckDbPlugin(BaseDatabasePlugin[DuckDBConfigFile]):
+    id = "jetbrains/duckdb"
+    name = "DuckDB Plugin"
+    supported = {"duckdb"}
+    config_file_type = DuckDBConfigFile
+
+    def __init__(self):
+        super().__init__(DuckDBIntrospector())

databao_context_engine/plugins/databases/{duckdb_introspector.py → duckdb/duckdb_introspector.py}

@@ -3,7 +3,7 @@ from __future__ import annotations
 import duckdb
 from pydantic import BaseModel, Field
 
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabaseConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
 from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
@@ -16,7 +16,7 @@ class DuckDBConfigFile(BaseDatabaseConfigFile):
 
 
 class DuckDBConnectionConfig(BaseModel):
-    database: str = Field(description="Path to the DuckDB database file")
+    database_path: str = Field(description="Path to the DuckDB database file")
 
 
 class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
@@ -24,13 +24,10 @@ class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
     _IGNORED_SCHEMAS = {"information_schema", "pg_catalog"}
     supports_catalogs = True
 
-    def _connect(self, file_config: DuckDBConfigFile):
-        database_path = str(file_config.connection.database)
+    def _connect(self, file_config: DuckDBConfigFile, *, catalog: str | None = None):
+        database_path = str(file_config.connection.database_path)
         return duckdb.connect(database=database_path)
 
-    def _connect_to_catalog(self, file_config: DuckDBConfigFile, catalog: str):
-        return self._connect(file_config)
-
     def _get_catalogs(self, connection, file_config: DuckDBConfigFile) -> list[str]:
         rows = self._fetchall_dicts(connection, "SELECT database_name FROM duckdb_databases();", None)
         catalogs = [r["database_name"] for r in rows if r.get("database_name")]

databao_context_engine/plugins/{mssql_db_plugin.py → databases/mssql/mssql_db_plugin.py}

@@ -1,11 +1,11 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.mssql_introspector import MSSQLConfigFile, MSSQLIntrospector
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.mssql.mssql_introspector import MSSQLConfigFile, MSSQLIntrospector
 
 
 class MSSQLDbPlugin(BaseDatabasePlugin[MSSQLConfigFile]):
     id = "jetbrains/mssql"
     name = "MSSQL DB Plugin"
-    supported = {"databases/mssql"}
+    supported = {"mssql"}
     config_file_type = MSSQLConfigFile
 
     def __init__(self):

databao_context_engine/plugins/databases/{mssql_introspector.py → mssql/mssql_introspector.py}

@@ -6,7 +6,7 @@ from mssql_python import connect  # type: ignore[import-untyped]
 from pydantic import BaseModel, Field
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabaseConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
 from databao_context_engine.plugins.databases.databases_types import DatabaseSchema, DatabaseTable
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
@@ -55,15 +55,15 @@ class MSSQLIntrospector(BaseIntrospector[MSSQLConfigFile]):
     )
     supports_catalogs = True
 
-    def _connect(self, file_config: MSSQLConfigFile):
+    def _connect(self, file_config: MSSQLConfigFile, *, catalog: str | None = None):
         connection = file_config.connection
-        connection_string = self._create_connection_string_for_config(connection.to_mssql_kwargs())
-        return connect(connection_string)
 
-    def _connect_to_catalog(self, file_config: MSSQLConfigFile, catalog: str):
-        cfg = file_config.model_copy(deep=True)
-        cfg.connection.database = catalog
-        return self._connect(cfg)
+        connection_kwargs = connection.to_mssql_kwargs()
+        if catalog:
+            connection_kwargs["database"] = catalog
+
+        connection_string = self._create_connection_string_for_config(connection_kwargs)
+        return connect(connection_string)
 
     def _get_catalogs(self, connection, file_config: MSSQLConfigFile) -> list[str]:
         database = file_config.connection.database
@@ -422,8 +422,7 @@ class MSSQLIntrospector(BaseIntrospector[MSSQLConfigFile]):
             "trust_server_certificate": "yes" if file_config.get("trust_server_certificate") else None,
         }
 
-        connection_string = ";".join(f"{k}={v}" for k, v in connection_parts.items() if v is not None)
-        return connection_string
+        return ";".join(f"{k}={v}" for k, v in connection_parts.items() if v is not None)
 
     def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
         with connection.cursor() as cursor:

databao_context_engine/plugins/{mysql_db_plugin.py → databases/mysql/mysql_db_plugin.py}

@@ -1,11 +1,11 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.mysql_introspector import MySQLConfigFile, MySQLIntrospector
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.mysql.mysql_introspector import MySQLConfigFile, MySQLIntrospector
 
 
 class MySQLDbPlugin(BaseDatabasePlugin[MySQLConfigFile]):
     id = "jetbrains/mysql"
     name = "MySQL DB Plugin"
-    supported = {"databases/mysql"}
+    supported = {"mysql"}
     config_file_type = MySQLConfigFile
 
     def __init__(self):

databao_context_engine/plugins/databases/{mysql_introspector.py → mysql/mysql_introspector.py}

@@ -5,7 +5,7 @@ from pydantic import BaseModel, Field
 from pymysql.constants import CLIENT
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabaseConfigFile
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
 from databao_context_engine.plugins.databases.databases_types import DatabaseSchema, DatabaseTable
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
@@ -35,18 +35,18 @@ class MySQLIntrospector(BaseIntrospector[MySQLConfigFile]):
 
     supports_catalogs = True
 
-    def _connect(self, file_config: MySQLConfigFile):
+    def _connect(self, file_config: MySQLConfigFile, *, catalog: str | None = None):
+        connection_kwargs = file_config.connection.to_pymysql_kwargs()
+
+        if catalog:
+            connection_kwargs["database"] = catalog
+
         return pymysql.connect(
-            **file_config.connection.to_pymysql_kwargs(),
+            **connection_kwargs,
             cursorclass=pymysql.cursors.DictCursor,
             client_flag=CLIENT.MULTI_STATEMENTS | CLIENT.MULTI_RESULTS,
         )
 
-    def _connect_to_catalog(self, file_config: MySQLConfigFile, catalog: str):
-        cfg = file_config.model_copy(deep=True)
-        cfg.connection.database = catalog
-        return self._connect(cfg)
-
     def _get_catalogs(self, connection, file_config: MySQLConfigFile) -> list[str]:
         with connection.cursor() as cur:
             cur.execute(

databao_context_engine/plugins/databases/postgresql/postgresql_db_plugin.py

@@ -0,0 +1,15 @@
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.postgresql.postgresql_introspector import (
+    PostgresConfigFile,
+    PostgresqlIntrospector,
+)
+
+
+class PostgresqlDbPlugin(BaseDatabasePlugin[PostgresConfigFile]):
+    id = "jetbrains/postgres"
+    name = "PostgreSQL DB Plugin"
+    supported = {"postgres"}
+    config_file_type = PostgresConfigFile
+
+    def __init__(self):
+        super().__init__(PostgresqlIntrospector())