-
Notifications
You must be signed in to change notification settings - Fork 1
feat(db): Add async versions for all database modules #37
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Changes from all commits
Commits
Show all changes
10 commits
Select commit
Hold shift + click to select a range
acaa43f
feat(db): Add async versions for all database modules
openhands-agent 54bab24
update kafka code
haobibo 0328995
Merge branch 'main' into dev/aio
haobibo 352267d
fix: Update license format to SPDX expression
openhands-agent f17e17d
fix ver
openhands-agent 61bbde5
Merge branch 'main' into dev/aio
haobibo 69d9fd9
Merge branch 'main' into dev/aio
haobibo 5b04d67
code clean
haobibo afe2cc8
remov default import in module
haobibo 2fbdc9f
code ql
haobibo File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,177 @@ | ||
| """ | ||
| aloha.db package - Database and middleware connection helpers. | ||
|
|
||
| Sync modules (blocking): | ||
| from aloha.db import PostgresOperator, MySqlOperator, RedisOperator, ... | ||
| from aloha.db.postgres import PostgresOperator | ||
| from aloha.db.mysql import MySqlOperator | ||
| from aloha.db.redis import RedisOperator | ||
| from aloha.db.mongo import MongoOperator | ||
| from aloha.db.elasticsearch import ElasticSearchOperator | ||
| from aloha.db.kafka import KafkaOperator | ||
| from aloha.db.sqlite import SqliteOperator | ||
| from aloha.db.duckdb import DuckOperator | ||
| from aloha.db.oracle import OracledbOperator | ||
|
|
||
| Async modules (non-blocking): | ||
| from aloha.db import PostgresOperator as PostgresOperatorAio, ... | ||
| from aloha.db.postgres_aio import PostgresOperator | ||
| from aloha.db.mysql_aio import MySqlOperator | ||
| from aloha.db.redis_aio import RedisOperator | ||
| from aloha.db.mongo_aio import MongoOperator | ||
| from aloha.db.elasticsearch_aio import ElasticSearchOperator | ||
| from aloha.db.kafka_aio import KafkaOperator | ||
| from aloha.db.sqlite_aio import SqliteOperator | ||
| from aloha.db.duckdb_aio import DuckOperator | ||
| from aloha.db.oracle_aio import OracledbOperator | ||
|
|
||
| Base utilities: | ||
| from aloha.db.base import PasswordVault | ||
| from aloha.db.base_aio import PasswordVault # async version | ||
|
|
||
| Usage example (sync): | ||
| from aloha.db.postgres import PostgresOperator | ||
|
|
||
| op = PostgresOperator(db_config) | ||
| result = op.execute_query("SELECT * FROM users") | ||
| for row in result: | ||
| print(row) | ||
|
|
||
| Usage example (async): | ||
| from aloha.db.postgres_aio import PostgresOperator | ||
|
|
||
| async def main(): | ||
| op = PostgresOperator(db_config) | ||
| result = await op.execute_query("SELECT * FROM users") | ||
| async for row in op.execute_query_scalars("SELECT * FROM users"): | ||
| print(row) | ||
| await op.close() | ||
|
|
||
| import asyncio | ||
| asyncio.run(main()) | ||
|
|
||
|
|
||
| from .base import PasswordVault | ||
|
|
||
| try: | ||
| from .postgres import PostgresOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .mysql import MySqlOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .redis import RedisOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .mongo import MongoOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .elasticsearch import ElasticSearchOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .kafka import ConsumedMessage, KafkaOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .sqlite import SqliteOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .duckdb import DuckOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .oracle import OracledbOperator | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
|
|
||
| # Async modules (importable as aliases for easy switching) | ||
| from .base_aio import PasswordVault as PasswordVaultAio | ||
|
|
||
| try: | ||
| from .postgres_aio import PostgresOperator as PostgresOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
|
github-code-quality[bot] marked this conversation as resolved.
Fixed
github-code-quality[bot] marked this conversation as resolved.
Fixed
|
||
| pass | ||
|
|
||
| try: | ||
| from .mysql_aio import MySqlOperator as MySqlOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .redis_aio import RedisOperator as RedisOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .mongo_aio import MongoOperator as MongoOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .elasticsearch_aio import ElasticSearchOperator as ElasticSearchOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .kafka_aio import ConsumedMessage as ConsumedMessageAio | ||
| from .kafka_aio import KafkaOperator as KafkaOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .sqlite_aio import SqliteOperator as SqliteOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .duckdb_aio import DuckOperator as DuckOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| try: | ||
| from .oracle_aio import OracledbOperator as OracledbOperatorAio | ||
| except (ImportError, ModuleNotFoundError): | ||
| pass | ||
|
|
||
| __all__ = ( | ||
| "ConsumedMessage", | ||
| "ConsumedMessageAio", | ||
| "DuckOperator", | ||
| "DuckOperatorAio", | ||
| "ElasticSearchOperator", | ||
| "ElasticSearchOperatorAio", | ||
| "KafkaOperator", | ||
| "KafkaOperatorAio", | ||
| "MongoOperator", | ||
| "MongoOperatorAio", | ||
| "MySqlOperator", | ||
| "MySqlOperatorAio", | ||
| "OracledbOperator", | ||
| "OracledbOperatorAio", | ||
| "PasswordVault", | ||
| "PasswordVaultAio", | ||
| # Sync operators | ||
| "PostgresOperator", | ||
| # Async operators (aliased) | ||
| "PostgresOperatorAio", | ||
| "RedisOperator", | ||
| "RedisOperatorAio", | ||
| "SqliteOperator", | ||
| "SqliteOperatorAio", | ||
| ) | ||
| """ | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,128 @@ | ||
| """ | ||
| Async DuckDB connection helpers. | ||
| """ | ||
|
|
||
| from pathlib import Path | ||
|
|
||
| import duckdb | ||
| import duckdb_engine | ||
| from sqlalchemy import text | ||
| from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine | ||
|
|
||
| from aloha.logger import LOG | ||
|
|
||
| __all__ = ("DuckOperator",) | ||
|
|
||
| LOG.debug("duckdb_aio version = %s, duckdb_engine = %s (async)", (duckdb.__version__, duckdb_engine.__version__)) | ||
|
|
||
|
|
||
| class DuckOperator: | ||
| """Create and use an async DuckDB connection through SQLAlchemy.""" | ||
|
|
||
| def __init__(self, db_config, **kwargs): | ||
| """Build an async DuckDB engine, creating the database file if necessary.""" | ||
| """db_config example: | ||
| { | ||
| "path": "/path/to/db.duckdb", # file path of duckdb, use ":memory:" for in-memory mode | ||
| "schema": "sales", # optional, 'main' by default | ||
| "read_only": True, # optional, False by default, (will set to False if in in-memory mode) | ||
| "config": {"memory_limit": "500mb"}, # optional, duckdb connection configs | ||
| } | ||
| """ | ||
| self._config = { | ||
| "path": db_config.get("path", ":memory:"), | ||
| "schema": db_config.get("schema", "main"), | ||
| "read_only": bool(db_config.get("read_only", False)), | ||
| "config": db_config.get("config", {}), | ||
| "auto_commit": db_config.get("auto_commit", True), | ||
| } | ||
|
|
||
| if not self._config["path"] or self._config["path"] == ":memory:": | ||
| self._config["path"] = ":memory:" | ||
|
|
||
| if self._config["read_only"]: | ||
| LOG.warning("In-memory database cannot be read-only. Setting read_only=False.") | ||
| self._config["read_only"] = False | ||
|
|
||
| else: | ||
| self._prepare_database() | ||
|
|
||
| try: | ||
| str_connection = f"duckdb+aioduckdb:///{self._config['path']}" | ||
| self.engine: AsyncEngine = create_async_engine( | ||
| str_connection, | ||
| connect_args={"read_only": self._config["read_only"], "config": self._config["config"]}, | ||
| **kwargs, | ||
| ) | ||
|
|
||
| LOG.debug("DuckDB (async) connected: {path} [schema={schema}, read_only={read_only}]".format(**self._config)) | ||
| except Exception as e: | ||
| LOG.exception(e) | ||
| raise RuntimeError("Failed to connect to DuckDB (async)") from e | ||
|
|
||
| def _prepare_database(self): | ||
| """Prepare the database file and its parent directory.""" | ||
| path = self._config["path"] | ||
| path_obj = Path(path) | ||
|
|
||
| parent_dir = path_obj.parent | ||
| if not parent_dir.exists(): | ||
| if self._config["read_only"]: | ||
| raise RuntimeError(f"Directory '{parent_dir}' does not exist and read_only=True") | ||
| try: | ||
| parent_dir.mkdir(parents=True, exist_ok=True) | ||
| LOG.debug(f"Created directory: {parent_dir}") | ||
| except Exception as e: | ||
| raise RuntimeError(f"Failed to create directory '{parent_dir}': {e}") from e | ||
|
|
||
| if not path_obj.exists(): | ||
| if self._config["read_only"]: | ||
| raise RuntimeError(f"DuckDB file '{path}' does not exist and read_only=True") | ||
| try: | ||
| LOG.debug(f"Database file not found, creating: {path}") | ||
| duckdb.connect(path).close() | ||
| except Exception as e: | ||
| raise RuntimeError(f"Failed to create database file '{path}': {e}") from e | ||
|
|
||
| @property | ||
| def connection(self): | ||
| return self.engine | ||
|
|
||
| @property | ||
| def conn(self): | ||
| """Alias for connection property.""" | ||
| return self.engine | ||
|
|
||
| async def execute_query(self, sql, *args, **kwargs): | ||
| """Execute a SQL statement asynchronously and return the cursor result.""" | ||
| async with self.engine.connect() as conn: | ||
| cur = await conn.execute(text(sql), *args, **kwargs) | ||
| if self._config.get("auto_commit", True): | ||
| await conn.commit() | ||
| return cur | ||
|
|
||
| async def execute_query_scalars(self, sql, *args, **kwargs): | ||
| """Execute a SQL statement and return all scalar results.""" | ||
| async with self.engine.connect() as conn: | ||
| cur = await conn.stream_scalars(text(sql), *args, **kwargs) | ||
| async for row in cur: | ||
| yield row | ||
|
|
||
| @property | ||
| def connection_str(self) -> str: | ||
| """Return a human-readable connection string.""" | ||
| return ( | ||
| f"duckdb:///{self._config['path']} [schema={self._config['schema']}, read_only={self._config['read_only']}] (async)" | ||
| ) | ||
|
|
||
| async def close(self): | ||
| """Close the async engine and all connections.""" | ||
| await self.engine.dispose() | ||
|
|
||
| async def __aenter__(self): | ||
| """Async context manager entry.""" | ||
| return self | ||
|
|
||
| async def __aexit__(self, exc_type, exc_val, exc_tb): | ||
| """Async context manager exit.""" | ||
| await self.close() |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.