75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import os
|
|
import subprocess
|
|
import tempfile
|
|
from typing import List, AsyncIterator, Dict, Any
|
|
from .base import SourceAdapter, FileInfo
|
|
|
|
|
|
class DatabaseAdapter(SourceAdapter):
|
|
async def connect(self) -> None:
|
|
pass
|
|
|
|
async def disconnect(self) -> None:
|
|
pass
|
|
|
|
async def list_files(self, path: str = "") -> List[FileInfo]:
|
|
return []
|
|
|
|
async def read_file(self, path: str) -> AsyncIterator[bytes]:
|
|
raise NotImplementedError("Database adapter does not support file reading")
|
|
|
|
async def get_database_dump(self, config: Dict[str, Any]) -> AsyncIterator[bytes]:
|
|
db_type = config.get("db_type", "postgresql")
|
|
host = config.get("host", "localhost")
|
|
port = config.get("port", 5432 if db_type == "postgresql" else 3306)
|
|
database = config.get("database")
|
|
username = config.get("username")
|
|
password = config.get("password")
|
|
|
|
if not database:
|
|
raise ValueError("Database name is required")
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".sql", delete=False) as tmp:
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
if db_type == "postgresql":
|
|
env = os.environ.copy()
|
|
if password:
|
|
env["PGPASSWORD"] = password
|
|
|
|
cmd = [
|
|
"pg_dump",
|
|
"-h", host,
|
|
"-p", str(port),
|
|
"-U", username or "postgres",
|
|
"-f", tmp_path,
|
|
database
|
|
]
|
|
elif db_type == "mysql":
|
|
env = os.environ.copy()
|
|
if password:
|
|
env["MYSQL_PWD"] = password
|
|
|
|
cmd = [
|
|
"mysqldump",
|
|
"-h", host,
|
|
"-P", str(port),
|
|
"-u", username or "root",
|
|
"--result-file", tmp_path,
|
|
database
|
|
]
|
|
else:
|
|
raise ValueError(f"Unsupported database type: {db_type}")
|
|
|
|
result = subprocess.run(cmd, env=env, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"Database dump failed: {result.stderr}")
|
|
|
|
with open(tmp_path, "rb") as f:
|
|
while chunk := f.read(8192):
|
|
yield chunk
|
|
finally:
|
|
if os.path.exists(tmp_path):
|
|
os.unlink(tmp_path)
|