feat(adapters): add SSH and database adapters
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user