feat: Patch setup_django so we can use an inmemory db in specific commands

This commit is contained in:
Cristian 2020-12-08 18:05:37 -05:00 committed by Cristian Vargas
parent 35a5700c73
commit 9aac09a5e1
6 changed files with 18 additions and 13 deletions

View file

@ -19,6 +19,8 @@ meta_cmds = ('help', 'version')
main_cmds = ('init', 'info', 'config')
archive_cmds = ('add', 'remove', 'update', 'list', 'status')
fake_db = ("oneshot",) + meta_cmds
display_first = (*meta_cmds, *main_cmds, *archive_cmds)
# every imported command module must have these properties in order to be valid
@ -59,7 +61,7 @@ def run_subcommand(subcommand: str,
pwd: Union[Path, str, None]=None) -> None:
"""Run a given ArchiveBox subcommand with the given list of args"""
from ..config import setup_django
setup_django()
setup_django(in_memory_db=subcommand in fake_db)
module = import_module('.archivebox_{}'.format(subcommand), __package__)
module.main(args=subcommand_args, stdin=stdin, pwd=pwd) # type: ignore

View file

@ -991,7 +991,7 @@ def check_data_folder(out_dir: Union[str, Path, None]=None, config: ConfigDict=C
def setup_django(out_dir: Path=None, check_db=False, config: ConfigDict=CONFIG) -> None:
def setup_django(out_dir: Path=None, check_db=False, config: ConfigDict=CONFIG, in_memory_db=False) -> None:
check_system_config()
output_dir = out_dir or Path(config['OUTPUT_DIR'])
@ -1004,7 +1004,15 @@ def setup_django(out_dir: Path=None, check_db=False, config: ConfigDict=CONFIG)
os.environ.setdefault('OUTPUT_DIR', str(output_dir))
assert (config['PACKAGE_DIR'] / 'core' / 'settings.py').exists(), 'settings.py was not found at archivebox/core/settings.py'
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'core.settings')
django.setup()
if in_memory_db:
# Put the db in memory and run migrations in case any command requires it
from django.core.management import call_command
os.environ.setdefault("ARCHIVEBOX_DATABASE_NAME", ":memory:")
django.setup()
call_command("migrate", interactive=False, stdout=False)
else:
django.setup()
if check_db:
sql_index_path = Path(output_dir) / SQL_INDEX_FILENAME

View file

@ -100,10 +100,12 @@ TEMPLATES = [
################################################################################
DATABASE_FILE = Path(OUTPUT_DIR) / SQL_INDEX_FILENAME
DATABASE_NAME = os.environ.get("ARCHIVEBOX_DATABASE_NAME", DATABASE_FILE)
DATABASES = {
'default': {
'ENGINE': 'django.db.backends.sqlite3',
'NAME': str(DATABASE_FILE),
'NAME': DATABASE_NAME,
}
}

View file

@ -353,12 +353,7 @@ class Link:
### Archive Status Helpers
@property
def num_outputs(self) -> int:
try:
return self.as_snapshot().num_outputs
except OperationalError:
return sum(1 for method in self.history.keys()
for result in self.history[method]
if result.status == 'succeeded')
return self.as_snapshot().num_outputs
@property
def num_failures(self) -> int:

View file

@ -524,7 +524,7 @@ def oneshot(url: str, out_dir: Path=OUTPUT_DIR):
)
raise SystemExit(2)
methods = ignore_methods(['title'])
archive_link(oneshot_link[0], out_dir=out_dir, methods=methods, skip_index=True)
archive_link(oneshot_link[0], out_dir=out_dir, methods=methods, skip_index=False)
return oneshot_link
@enforce_types

View file

@ -33,7 +33,6 @@ def write_search_index(link: Link, texts: Union[List[str], None]=None, out_dir:
return
if not skip_text_index and texts:
setup_django(out_dir, check_db=True)
from core.models import Snapshot
snap = Snapshot.objects.filter(url=link.url).first()
@ -91,7 +90,6 @@ def index_links(links: Union[List[Link],None], out_dir: Path=OUTPUT_DIR):
if not links:
return
setup_django(out_dir=out_dir, check_db=True)
from core.models import Snapshot, ArchiveResult
for link in links: