Feature Request: Custom Database Driver Plugin System

[platinummonkey]

Summary

Add support for pluggable custom database drivers to allow organizations with specialized database infrastructure to integrate their custom connection management, credential handling, and database interaction patterns.

Problem Statement

Our organization uses a managed PostgreSQL platform with custom credential management, connection pooling, and security requirements that differ from standard PostgreSQL connections. The current DBOS implementation hardcodes the PostgreSQL driver (postgresql+psycopg) and connection logic, making it difficult to integrate with our internal platform team's database infrastructure.

We need a way to plug in our custom database driver that handles:

• Custom credential management and rotation
• Specialized connection settings and security requirements
• Internal monitoring and observability integration
• Platform-specific connection pooling strategies

This would of course need to be replicated pattern in other languages, however the primary use case I currently have is python and golang.

Proposed Solution

1. Database Driver Interface

Create an abstract base class that custom drivers can implement:

  from abc import ABC, abstractmethod
  import sqlalchemy as sa
  from typing import Dict, Any

  class DatabaseDriver(ABC):
      @abstractmethod
      def create_engine(self, database_url: str, engine_kwargs: Dict[str, Any]) -> sa.Engine:
          """Create and return a configured SQLAlchemy engine."""
          pass

      @abstractmethod
      def validate_url(self, database_url: str) -> bool:
          """Validate that the provided URL is compatible with this driver."""
          pass

      @abstractmethod
      def get_dialect_name(self) -> str:
          """Return the SQLAlchemy dialect name for this driver."""
          pass

2. Driver Registry System

Implement a registry system similar to the existing custom PostgreSQL implementation pattern:

  class DatabaseDriverRegistry:
      _drivers: Dict[str, DatabaseDriver] = {}

      @classmethod
      def register_driver(cls, name: str, driver: DatabaseDriver) -> None:
          """Register a custom database driver."""
          cls._drivers[name] = driver

      @classmethod
      def get_driver(cls, name: str) -> Optional[DatabaseDriver]:
          """Retrieve a registered driver by name."""
          return cls._drivers.get(name)

3. Configuration Integration

Extend the existing configuration system to support custom drivers:

  # dbos-config.yaml
  name: my-app
  system_database_url: "custom://internal-postgres"
  database:
    custom_driver: "internal_postgres_driver"

  # DBOSConfig type extension
  class DBOSConfig(TypedDict, total=False):
      # ... existing fields ...
      custom_database_driver: Optional[str]

4. Usage Example

Organizations could implement their custom driver:

  from dbos import DatabaseDriver, DatabaseDriverRegistry
  import sqlalchemy as sa

  class InternalPostgresDriver(DatabaseDriver):
      def create_engine(self, database_url: str, engine_kwargs: Dict[str, Any]) -> sa.Engine:
          # Custom credential fetching
          credentials = self._fetch_internal_credentials()

          # Transform URL with internal connection details
          internal_url = self._transform_to_internal_url(database_url, credentials)

          # Apply internal connection settings
          internal_kwargs = {**engine_kwargs, **self._get_internal_engine_config()}

          return sa.create_engine(internal_url, **internal_kwargs)

      def validate_url(self, database_url: str) -> bool:
          return database_url.startswith("custom://")

      def get_dialect_name(self) -> str:
          return "postgresql"  # Still PostgreSQL, but with custom connection logic

Register the driver

  DatabaseDriverRegistry.register_driver("internal_postgres_driver", InternalPostgresDriver())

Benefits

1. Platform Integration: Organizations can integrate DBOS with their existing database infrastructure without forking the codebase
2. Security Compliance: Custom drivers can implement organization-specific security requirements and credential management
3. Observability: Internal drivers can integrate with existing monitoring and logging systems
4. Future-Proof: Extensible system that can accommodate various database platforms and connection patterns
5. Backward Compatibility: Existing configurations continue to work unchanged

Implementation Notes

• The system should fall back to default driver when no custom driver is specified.
• Custom drivers should be loaded and registered before DBOS initialization. But ideally also load after DBOS(...) initialization call.
• Error handling should provide clear feedback when custom drivers fail to load or validate
• Documentation should include examples for common use cases (credential rotation, connection proxies like pgbouncer, pg-proxy, etc.)

Alternative Approaches Considered

1. Environment Variable Overrides: Limited flexibility, doesn't address complex connection logic
2. Configuration File Templating: Doesn't handle dynamic credential management or complex connection patterns
3. Monkey Patching: Fragile and makes upgrades difficult, although this is how I've currently hacked a proof-of-concept together.

This feature would enable organizations with specialized database infrastructure to adopt DBOS while maintaining their existing security, monitoring, and operational practices.

[kraftp]

Thanks for opening this!

Follow-up question: in Golang, you can optionally configure DBOS with a pgx pool which it uses for all system database operations instead of creating its own pool from your system database URL. Here's what it looks like.

Does passing in a custom system database pool work for you in Golang? Would something similar work for you in Python (instead of a custom driver class, you configure DBOS with a sa.Engine that is used for the system database)?

[platinummonkey]

It's possible, however there is some additional observability and connection commit/rollback defaults we've encoded into our internal managed postgres platform. I might be able to play around with a branch you have in mind. This would be similar in go, but I've mainly been experimenting with python so far.

[kraftp]

Great! I'm working on it, will let you know when there's a PR/branch that's ready to look at!

[kraftp]

Here's a PR for supporting a custom database engine: #469

Does this address the issue?

[platinummonkey]

Thanks! I’ll give this a try!