Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@ When writing code or configurations for this project, you **MUST** strictly adhe

### A. Python Coding & Naming Standards

- **Primary Type/First-Class Identity Prefix**: Place the variable's type, role, or primary characteristics first in its name.
- _Correct_: `name_service`, `port_service`, `svc_ingress`, `cfg_postgres`, `db_mysql`.
- _Incorrect_: `service_name`, `service_port`, `ingress_service`, `postgres_config`, `mysql_db`.
- **Primary Type/First-Class Identity Prefix**: Place the variable's type, role, or primary characteristics first in its name. If multiple variables in a segment of code belong to the same category, type, or semantic group, place the common semantic prefix first.
- _Correct_: `name_service`, `port_service`, `svc_ingress`, `cfg_postgres`, `db_mysql`, `msg_err`, `msg_info`.
- _Incorrect_: `service_name`, `service_port`, `ingress_service`, `postgres_config`, `mysql_db`, `err_msg`, `info_msg`.
- **Logger Naming**: Use lowercase with underscores for logger names, e.g., `db_sync`, `api_router`.
- **Import Conventions**: Use relative imports if possible, especially inside a package.

Expand Down
6 changes: 3 additions & 3 deletions doc/skills/aloha_python/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@ This skill provides coding standards, modular application structures, and usage

When developing Python code in this codebase, adhere to the following naming conventions:

- **Primary Type/First-Class Identity Prefix**: Place the variable's type, role, or primary characteristics first in its name.
- _Correct_: `name_service`, `port_service`, `svc_ingress`, `cfg_postgres`.
- _Incorrect_: `service_name`, `service_port`, `ingress_service`, `postgres_config`.
- **Primary Type/First-Class Identity Prefix**: Place the variable's type, role, or primary characteristics first in its name. If multiple variables in a segment of code belong to the same category, type, or semantic group, place the common semantic prefix first.
- _Correct_: `name_service`, `port_service`, `svc_ingress`, `cfg_postgres`, `db_mysql`, `msg_err`, `msg_info`.
- _Incorrect_: `service_name`, `service_port`, `ingress_service`, `postgres_config`, `mysql_db`, `err_msg`, `info_msg`.
- **Logger naming**: Use lowercase with underscores for logger names, e.g., `db_sync`.

---
Expand Down
177 changes: 177 additions & 0 deletions pkg/aloha/db/__init__.py
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):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .mysql import MySqlOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .redis import RedisOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .mongo import MongoOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .elasticsearch import ElasticSearchOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .kafka import ConsumedMessage, KafkaOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .sqlite import SqliteOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .duckdb import DuckOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
pass

try:
from .oracle import OracledbOperator
except (ImportError, ModuleNotFoundError):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
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):
Comment thread
github-code-quality[bot] marked this conversation as resolved.
Fixed
Comment thread
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",
)
"""
128 changes: 128 additions & 0 deletions pkg/aloha/db/duckdb_aio.py
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()
Loading