Skip to content

Commit fb75554

Browse files
Charlie-BUCharlie-BU
committed
fix: db-migration
1 parent 02ee414 commit fb75554

File tree

2 files changed

+114
-55
lines changed

2 files changed

+114
-55
lines changed

‎robyn/cli.py‎

Lines changed: 62 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -83,14 +83,33 @@ def create_robyn_app():
8383
if docker == "N":
8484
os.remove(f"{final_project_dir_path}/Dockerfile")
8585

86-
# If database migration is needed, install the latest version of alembic
86+
# If database migration is needed, install alembic
8787
if db_migration == "Y":
88-
print("Installing the latest version of alembic...")
88+
print("Installing alembic...")
8989
try:
90-
subprocess.run([sys.executable, "-m", "pip", "install", "alembic", "-q"], check=True,
91-
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
92-
except subprocess.CalledProcessError:
93-
print("Failed to install alembic. Please install it manually using 'pip install alembic'.")
90+
# Check if alembic is already installed
91+
import importlib.util
92+
alembic_spec = importlib.util.find_spec('alembic')
93+
94+
if alembic_spec is None:
95+
# Install alembic using pip API
96+
try:
97+
import pip
98+
print("Installing alembic using pip API...")
99+
from pip._internal.cli.main import main as pip_main
100+
pip_main(['install', 'alembic', '--quiet'])
101+
print("Successfully installed alembic.")
102+
except ImportError:
103+
# If pip API is not available, use subprocess
104+
print("Installing alembic using subprocess...")
105+
subprocess.run([sys.executable, "-m", "pip", "install", "alembic", "-q"], check=True,
106+
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
107+
print("Successfully installed alembic.")
108+
else:
109+
print("Alembic is already installed.")
110+
except (subprocess.CalledProcessError, ImportError) as e:
111+
print(f"Failed to install alembic: {str(e)}")
112+
print("Please install it manually using 'pip install alembic'.")
94113

95114
print(f"New Robyn project created in '{final_project_dir_path}' ")
96115

@@ -124,20 +143,45 @@ def start_app_normally(config: Config):
124143

125144
def handle_db_command():
126145
"""Handle database migration commands."""
127-
try:
128-
from robyn.migrate import configure_parser, execute_command
129-
except ImportError:
130-
try:
131-
import importlib.util
132-
if importlib.util.find_spec("alembic") is None:
133-
print("ERROR: Alembic has not been installed. Please run 'pip install alembic' to install it.")
146+
import importlib.util
147+
alembic_spec = importlib.util.find_spec("alembic")
148+
149+
if alembic_spec is None:
150+
print("ERROR: Alembic has not been installed.")
151+
install_choice = input("Would you like to install alembic now? (y/n): ").strip().lower()
152+
153+
if install_choice == 'y':
154+
try:
155+
try:
156+
from pip._internal.cli.main import main as pip_main
157+
print("Installing alembic...")
158+
pip_main(['install', 'alembic', '--quiet'])
159+
print("Successfully installed alembic.")
160+
except ImportError:
161+
print("Installing alembic using subprocess...")
162+
subprocess.run([sys.executable, "-m", "pip", "install", "alembic", "-q"], check=True)
163+
print("Successfully installed alembic.")
164+
165+
importlib.invalidate_caches()
166+
alembic_spec = importlib.util.find_spec("alembic")
167+
if alembic_spec is None:
168+
print("ERROR: Failed to install alembic. Please install it manually using 'pip install alembic'.")
169+
sys.exit(1)
170+
except Exception as e:
171+
print(f"ERROR: Failed to install alembic: {str(e)}")
172+
print("Please install it manually using 'pip install alembic'.")
134173
sys.exit(1)
135-
else:
136-
print("ERROR: Failed to import migrate module.")
137-
sys.exit(1)
138-
except ImportError:
139-
print("ERROR: Fail to import migrate module.")
174+
else:
175+
print("Please install alembic manually using 'pip install alembic' before using database commands.")
140176
sys.exit(1)
177+
178+
try:
179+
from robyn.migrate import configure_parser, execute_command
180+
except ImportError as e:
181+
print(f"ERROR: Failed to import migrate module: {str(e)}")
182+
print("This might be due to an incomplete installation or a version mismatch.")
183+
print("Try reinstalling Robyn or updating your dependencies.")
184+
sys.exit(1)
141185
parser = argparse.ArgumentParser(
142186
usage=argparse.SUPPRESS, # omit usage hint
143187
description='Robyn database migration commands.'

‎robyn/migrate.py‎

Lines changed: 52 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import logging
55
from functools import wraps
66
import argparse
7+
import re
78
from pathlib import Path
89
from typing import Optional, List, Dict, Any, Union, Callable
910

@@ -85,57 +86,64 @@ def list_templates() -> None:
8586

8687

8788
def _auto_configure_migrations(directory: str, db_url: Optional[str] = None, model_path: Optional[str] = None) -> None:
88-
"""自动配置 alembic.ini 和 env.py 文件
89+
"""Automatically configure alembic.ini and env.py
8990
9091
Args:
91-
directory: 迁移文件目录
92-
db_url: 数据库 URL
93-
model_path: 模型文件路径
92+
directory: Directory where migration files are stored
93+
db_url: Database URL
94+
model_path: Path to the model file
9495
"""
9596
# Configure alembic.ini
9697
if db_url:
9798
alembic_ini_path = os.path.join(directory, 'alembic.ini')
9899
if os.path.exists(alembic_ini_path):
100+
99101
with open(alembic_ini_path, 'r') as f:
100102
content = f.read()
101103

102-
# Replace the database URL
103-
content = content.replace('sqlalchemy.url = driver://user:pass@localhost/dbname',
104-
f'sqlalchemy.url = {db_url}')
104+
# Replace the database URL using regex pattern for more flexibility
105+
pattern = r'sqlalchemy\.url\s*=\s*[^\n]+'
106+
replacement = f'sqlalchemy.url = {db_url}'
107+
new_content = re.sub(pattern, replacement, content)
105108

106-
with open(alembic_ini_path, 'w') as f:
107-
f.write(content)
108-
print(f"Successfully configured the database URL: {db_url}")
109+
if new_content != content:
110+
with open(alembic_ini_path, 'w') as f:
111+
f.write(new_content)
112+
print(f"Successfully configured the database URL: {db_url}")
113+
else:
114+
print("Warning: Could not find database URL configuration in alembic.ini")
109115

110116
# Configure env.py
111117
if model_path:
112118
env_py_path = os.path.join(directory, 'env.py')
113119
if os.path.exists(env_py_path):
120+
114121
with open(env_py_path, 'r') as f:
115122
content = f.read()
116123

117124
try:
118125
module_path, class_name = model_path.rsplit('.', 1)
119-
120-
# Replace the import statement and target_metadata setting
121126
# Allow importing from the parent directory
122127
import_statement = f"sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))\nfrom {module_path} import {class_name}\ntarget_metadata = {class_name}.metadata"
123128

124-
# Replace the import statement
125-
content = content.replace(
126-
"# add your model's MetaData object here\n# for 'autogenerate' support\n# from myapp import mymodel\n# target_metadata = mymodel.Base.metadata",
127-
f"# add your model's MetaData object here\n# for 'autogenerate' support\n{import_statement}"
128-
)
129+
# Replace the import statement using regex for more robustness
130+
import_pattern = r"# add your model's MetaData object here\s*\n# for 'autogenerate' support\s*\n# from myapp import mymodel\s*\n# target_metadata = mymodel\.Base\.metadata"
131+
import_replacement = f"# add your model's MetaData object here\n# for 'autogenerate' support\n{import_statement}"
132+
133+
new_content = re.sub(import_pattern, import_replacement, content)
129134

130135
# Replace the target_metadata setting
131-
content = content.replace(
132-
"target_metadata = config.attributes.get('sqlalchemy.metadata', None)",
133-
f"# target_metadata = config.attributes.get('sqlalchemy.metadata', None)\n# Already set by the import above"
134-
)
135-
136-
with open(env_py_path, 'w') as f:
137-
f.write(content)
138-
print(f"Successfully configured the model path: {model_path}")
136+
metadata_pattern = r"target_metadata = config\.attributes\.get\('sqlalchemy\.metadata', None\)"
137+
metadata_replacement = f"# target_metadata = config.attributes.get('sqlalchemy.metadata', None)\n# Already set by the import above"
138+
139+
new_content = re.sub(metadata_pattern, metadata_replacement, new_content)
140+
141+
if new_content != content:
142+
with open(env_py_path, 'w') as f:
143+
f.write(new_content)
144+
print(f"Successfully configured the model path: {model_path}")
145+
else:
146+
print("Warning: Could not find expected patterns in env.py, please manually configure it")
139147
except ValueError:
140148
print(f"Warning: Could not parse the model path {model_path}, please manually configure env.py")
141149

@@ -148,15 +156,21 @@ def _special_configure_for_sqlite(directory: str, model_path: Optional[str] = No
148156
with open(env_py_path, 'r') as f:
149157
content = f.read()
150158
try:
151-
content = content.replace(
152-
" context.configure(\n connection=connection,\n target_metadata=target_metadata,\n process_revision_directives=process_revision_directives,\n )",
153-
" from sqlalchemy.engine import Connection\n def is_sqlite(conn: Connection) -> bool:\n return conn.dialect.name == \"sqlite\"\n context.configure(\n connection=connection,\n target_metadata=target_metadata,\n process_revision_directives=process_revision_directives,\n render_as_batch=is_sqlite(connection),\n )"
154-
)
155-
with open(env_py_path, 'w') as f:
156-
f.write(content)
157-
except ValueError:
159+
# Use regex pattern to match the context.configure block more flexibly
160+
pattern = r'\s+context\.configure\(\s*\n\s+connection=connection,\s*\n\s+target_metadata=target_metadata,\s*\n\s+process_revision_directives=process_revision_directives,\s*\n\s+\)'
161+
replacement = "\n from sqlalchemy.engine import Connection\n def is_sqlite(conn: Connection) -> bool:\n return conn.dialect.name == \"sqlite\"\n context.configure(\n connection=connection,\n target_metadata=target_metadata,\n process_revision_directives=process_revision_directives,\n render_as_batch=is_sqlite(connection),\n )"
162+
163+
new_content = re.sub(pattern, replacement, content)
164+
165+
if new_content != content:
166+
with open(env_py_path, 'w') as f:
167+
f.write(new_content)
168+
else:
169+
print(
170+
"Warning: Could not find context.configure block in env.py, please manually add render_as_batch=True for SQLite support")
171+
except Exception as e:
158172
print(
159-
"Warning: If your database is SQLite, you need to manually add `render_as_batch=True` in run_migrations_online() to avoid migration errors caused by SQLite's limited support for ALTER TABLE.")
173+
f"Warning: Could not configure SQLite support: {str(e)}. If your database is SQLite, you need to manually add `render_as_batch=True` in run_migrations_online() to avoid migration errors caused by SQLite's limited support for ALTER TABLE.")
160174
else:
161175
print(
162176
"Warning: If your database is SQLite, you need to manually add `render_as_batch=True` in run_migrations_online() to avoid migration errors caused by SQLite's limited support for ALTER TABLE.")
@@ -193,8 +207,8 @@ def init(directory: str = 'migrations', multidb: bool = False, template: Optiona
193207
print(
194208
"Cannot find models module.\nPlease provide your database URL with \"--db-url=<YOUR_DB_URL>\".")
195209
return
196-
except Exception as e:
197-
print("Please provide your database URL with \"--db-url=<YOUR_DB_URL>\".")
210+
except ImportError:
211+
print("Cannot find models module.\nPlease provide your database URL with \"--db-url=<YOUR_DB_URL>\".")
198212
return
199213

200214
if not model_path:
@@ -214,15 +228,16 @@ def init(directory: str = 'migrations', multidb: bool = False, template: Optiona
214228
print(
215229
"Cannot find models module.\nPlease provide your model path with \"--model-path=<YOUR_MODEL_PATH>\".")
216230
return
217-
except Exception as e:
218-
print("Please provide your model path with \"--model-path=<YOUR_MODEL_PATH>\".")
231+
except ImportError:
232+
print("Cannot find models module.\nPlease provide your model path with \"--model-path=<YOUR_MODEL_PATH>\".")
219233
return
220234

221235
# Ensure the directory exists
222236
os.makedirs(directory, exist_ok=True)
223237

224238
config = Config(directory)
225239
template_path = config._get_template_path(template) if template is not None else config._get_template_path()
240+
print(template_path)
226241
command.init(config, directory, template=template_path, package=package)
227242

228243
_auto_configure_migrations(directory, db_url, model_path)

0 commit comments

Comments
 (0)