mirror of
https://github.com/KohakuBlueleaf/KohakuHub.git
synced 2026-04-29 02:48:33 -05:00
update doc and Docker related utils
This commit is contained in:
Kohaku-Blueleaf
272
scripts/sync_lakefs_credentials.py
Normal file
272
scripts/sync_lakefs_credentials.py
Normal file
@@ -0,0 +1,272 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Sync LakeFS Credentials to config.toml
|
||||
|
||||
This script reads LakeFS credentials from credentials.env (auto-generated by Docker)
|
||||
and updates config.toml with the correct values.
|
||||
|
||||
Usage:
|
||||
python scripts/sync_lakefs_credentials.py
|
||||
python scripts/sync_lakefs_credentials.py --credentials-path ./custom/path/credentials.env
|
||||
python scripts/sync_lakefs_credentials.py --config ./custom-config.toml
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def find_credentials_path_from_docker_compose(docker_compose_path: Path) -> Path | None:
|
||||
"""Find credentials.env path by parsing docker-compose.yml.
|
||||
|
||||
Args:
|
||||
docker_compose_path: Path to docker-compose.yml
|
||||
|
||||
Returns:
|
||||
Path to credentials.env or None if not found
|
||||
"""
|
||||
if not docker_compose_path.exists():
|
||||
return None
|
||||
|
||||
try:
|
||||
with open(docker_compose_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
# Look for volume mount pattern: ./path/to/dir:/hub-api-creds
|
||||
# Example: - ./hub-meta/hub-api:/hub-api-creds
|
||||
match = re.search(r"- (\.[\w\-/\\]+):/hub-api-creds", content)
|
||||
if match:
|
||||
host_path = match.group(1)
|
||||
# Resolve relative path
|
||||
base_dir = docker_compose_path.parent
|
||||
full_path = (base_dir / host_path / "credentials.env").resolve()
|
||||
return full_path
|
||||
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"⚠ Failed to parse docker-compose.yml: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def read_credentials_env(filepath: Path) -> dict[str, str]:
|
||||
"""Read credentials from credentials.env file.
|
||||
|
||||
Args:
|
||||
filepath: Path to credentials.env
|
||||
|
||||
Returns:
|
||||
Dict of {key: value}
|
||||
"""
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"Credentials file not found: {filepath}")
|
||||
|
||||
credentials = {}
|
||||
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
|
||||
# Parse KEY=value
|
||||
match = re.match(r"^([A-Z_]+)=(.+)$", line)
|
||||
if match:
|
||||
key, value = match.groups()
|
||||
credentials[key] = value.strip()
|
||||
|
||||
return credentials
|
||||
|
||||
|
||||
def update_config_toml(
|
||||
config_path: Path, lakefs_access_key: str, lakefs_secret_key: str
|
||||
):
|
||||
"""Update config.toml with LakeFS credentials.
|
||||
|
||||
Args:
|
||||
config_path: Path to config.toml
|
||||
lakefs_access_key: LakeFS access key
|
||||
lakefs_secret_key: LakeFS secret key
|
||||
"""
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {config_path}")
|
||||
|
||||
# Read existing config
|
||||
try:
|
||||
with open(config_path, "rb") as f:
|
||||
config = tomllib.load(f)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to parse config.toml: {e}")
|
||||
|
||||
# Update lakefs section
|
||||
if "lakefs" not in config:
|
||||
config["lakefs"] = {}
|
||||
|
||||
config["lakefs"]["access_key"] = lakefs_access_key
|
||||
config["lakefs"]["secret_key"] = lakefs_secret_key
|
||||
|
||||
# Write back
|
||||
lines = []
|
||||
|
||||
for section in [
|
||||
"s3",
|
||||
"lakefs",
|
||||
"smtp",
|
||||
"auth",
|
||||
"admin",
|
||||
"app",
|
||||
"quota",
|
||||
"fallback",
|
||||
]:
|
||||
if section not in config:
|
||||
continue
|
||||
|
||||
lines.append(f"[{section}]")
|
||||
|
||||
for key, val in config[section].items():
|
||||
if isinstance(val, bool):
|
||||
lines.append(f"{key} = {str(val).lower()}")
|
||||
elif isinstance(val, int):
|
||||
# Check if it's a large number with underscores
|
||||
if val >= 1000000:
|
||||
# Format with underscores for readability
|
||||
val_str = f"{val:_}"
|
||||
lines.append(f"{key} = {val_str}")
|
||||
else:
|
||||
lines.append(f"{key} = {val}")
|
||||
elif isinstance(val, float):
|
||||
lines.append(f"{key} = {val}")
|
||||
elif isinstance(val, str):
|
||||
lines.append(f'{key} = "{val}"')
|
||||
elif isinstance(val, list):
|
||||
# Format list
|
||||
items = ", ".join(
|
||||
f'"{item}"' if isinstance(item, str) else str(item) for item in val
|
||||
)
|
||||
lines.append(f"{key} = [{items}]")
|
||||
else:
|
||||
lines.append(f'{key} = "{val}"')
|
||||
|
||||
lines.append("") # Blank line after section
|
||||
|
||||
with open(config_path, "w", encoding="utf-8") as f:
|
||||
f.write("\n".join(lines))
|
||||
|
||||
print(f"✓ Updated {config_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Sync LakeFS credentials from credentials.env to config.toml"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--credentials-path",
|
||||
type=Path,
|
||||
help="Path to credentials.env (default: auto-detect from docker-compose.yml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--config",
|
||||
type=Path,
|
||||
default=Path("config.toml"),
|
||||
help="Path to config.toml (default: config.toml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--docker-compose",
|
||||
type=Path,
|
||||
default=Path("docker-compose.yml"),
|
||||
help="Path to docker-compose.yml (default: docker-compose.yml)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print("=" * 60)
|
||||
print("LakeFS Credentials Sync Tool")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
# Determine credentials path
|
||||
credentials_path = args.credentials_path
|
||||
|
||||
if not credentials_path:
|
||||
# Auto-detect from docker-compose.yml
|
||||
print(f"Auto-detecting credentials path from {args.docker_compose}...")
|
||||
credentials_path = find_credentials_path_from_docker_compose(
|
||||
args.docker_compose
|
||||
)
|
||||
|
||||
if not credentials_path:
|
||||
print("\n✗ Could not auto-detect credentials path from docker-compose.yml")
|
||||
print("\n💡 Trying default path: ./hub-meta/hub-api/credentials.env")
|
||||
credentials_path = Path("hub-meta/hub-api/credentials.env")
|
||||
|
||||
print(f"Credentials file: {credentials_path}")
|
||||
print(f"Config file: {args.config}")
|
||||
print()
|
||||
|
||||
# Check if files exist
|
||||
if not credentials_path.exists():
|
||||
print(f"✗ Credentials file not found: {credentials_path}")
|
||||
print("\n💡 Make sure docker-compose is running and LakeFS has initialized:")
|
||||
print(" docker-compose up -d")
|
||||
print(" # Wait for LakeFS to start and create credentials.env")
|
||||
sys.exit(1)
|
||||
|
||||
if not args.config.exists():
|
||||
print(f"✗ Config file not found: {args.config}")
|
||||
print("\n💡 Generate config.toml first:")
|
||||
print(" python scripts/generate_docker_compose.py")
|
||||
sys.exit(1)
|
||||
|
||||
# Read credentials
|
||||
print("Reading LakeFS credentials...")
|
||||
try:
|
||||
credentials = read_credentials_env(credentials_path)
|
||||
|
||||
lakefs_access_key = credentials.get("KOHAKU_HUB_LAKEFS_ACCESS_KEY")
|
||||
lakefs_secret_key = credentials.get("KOHAKU_HUB_LAKEFS_SECRET_KEY")
|
||||
|
||||
if not lakefs_access_key or not lakefs_secret_key:
|
||||
print("✗ Missing LakeFS credentials in credentials.env")
|
||||
print(f" Found keys: {list(credentials.keys())}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f" ✓ Access Key: {lakefs_access_key}")
|
||||
print(f" ✓ Secret Key: {lakefs_secret_key[:8]}..." + "*" * 20)
|
||||
print()
|
||||
|
||||
except FileNotFoundError as e:
|
||||
print(f"✗ {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to read credentials: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Update config.toml
|
||||
print(f"Updating {args.config}...")
|
||||
try:
|
||||
update_config_toml(args.config, lakefs_access_key, lakefs_secret_key)
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("✓ Sync Complete!")
|
||||
print("=" * 60)
|
||||
print()
|
||||
print("📋 Updated fields:")
|
||||
print(f" • lakefs.access_key = {lakefs_access_key}")
|
||||
print(f" • lakefs.secret_key = {lakefs_secret_key[:8]}***")
|
||||
print()
|
||||
print("💡 Next steps:")
|
||||
print(" 1. Restart dev server if running")
|
||||
print(" 2. Test LakeFS connection: curl http://localhost:28000/_health")
|
||||
print()
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to update config.toml: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user