Merge pull request #1698 from ranaroussi/hotfix/download-database-error

Fix: OperationalError('unable to open database file')

tests/context.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+import appdirs as _ad
+import datetime as _dt
 import sys
 import os
 _parent_dp = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
@@ -14,10 +16,17 @@
 # logging.basicConfig(level=logging.DEBUG)
 
 
+# Use adjacent cache folder for testing, delete if already exists and older than today
+testing_cache_dirpath = os.path.join(_ad.user_cache_dir(), "py-yfinance-testing")
+yfinance.set_tz_cache_location(testing_cache_dirpath)
+if os.path.isdir(testing_cache_dirpath):
+    mtime = _dt.datetime.fromtimestamp(os.path.getmtime(testing_cache_dirpath))
+    if mtime.date() < _dt.date.today():
+        import shutil
+        shutil.rmtree(testing_cache_dirpath)
+
+
 # Setup a session to rate-limit and cache persistently:
-import datetime as _dt
-import os
-import appdirs as _ad
 from requests import Session
 from requests_cache import CacheMixin, SQLiteCache
 from requests_ratelimiter import LimiterMixin, MemoryQueueBucket
@@ -26,12 +35,7 @@ class CachedLimiterSession(CacheMixin, LimiterMixin, Session):
 from pyrate_limiter import Duration, RequestRate, Limiter
 history_rate = RequestRate(1, Duration.SECOND*2)
 limiter = Limiter(history_rate)
-cache_fp = os.path.join(_ad.user_cache_dir(), "py-yfinance", "unittests-cache")
-if os.path.isfile(cache_fp + '.sqlite'):
-    # Delete local cache if older than 1 day:
-    mod_dt = _dt.datetime.fromtimestamp(os.path.getmtime(cache_fp + '.sqlite'))
-    if mod_dt.date() < _dt.date.today():
-        os.remove(cache_fp + '.sqlite')
+cache_fp = os.path.join(testing_cache_dirpath, "unittests-cache")
 session_gbl = CachedLimiterSession(
     limiter=limiter,
     bucket_class=MemoryQueueBucket,

yfinance/utils.py

@@ -902,74 +902,132 @@ def __str__(self):
 # ---------------------------------
 
 
-_cache_dir = _os.path.join(_ad.user_cache_dir(), "py-yfinance")
-DB_PATH = _os.path.join(_cache_dir, 'tkr-tz.db')
-db = _peewee.SqliteDatabase(DB_PATH, pragmas={'journal_mode': 'wal', 'cache_size': -64})
-_tz_cache = None
+_cache_init_lock = Lock()
 
 
 class _TzCacheException(Exception):
     pass
 
 
-class KV(_peewee.Model):
+class _TzCacheDummy:
+    """Dummy cache to use if tz cache is disabled"""
+
+    def lookup(self, tkr):
+        return None
+
+    def store(self, tkr, tz):
+        pass
+
+    @property
+    def tz_db(self):
+        return None
+
+
+class _TzCacheManager:
+    _tz_cache = None
+
+    @classmethod
+    def get_tz(cls):
+        if cls._tz_cache is None:
+            cls._initialise()
+        return cls._tz_cache
+
+    @classmethod
+    def _initialise(cls, cache_dir=None):
+        try:
+            cls._tz_cache = _TzCache()
+        except _TzCacheException as err:
+            get_yf_logger().info(f"Failed to create TzCache, reason: {err}. "
+                                 "TzCache will not be used. "
+                                 "Tip: You can direct cache to use a different location with 'set_tz_cache_location(mylocation)'")
+            cls._tz_cache = _TzCacheDummy()
+
+
+class _DBManager:
+    _db = None
+    _cache_dir = _os.path.join(_ad.user_cache_dir(), "py-yfinance")
+
+    @classmethod
+    def get_database(cls):
+        if cls._db is None:
+            cls._initialise()
+        return cls._db
+
+    @classmethod
+    def close_db(cls):
+        if cls._db is not None:
+            try:
+                cls._db.close()
+            except Exception as e:
+                # Must discard exceptions because Python trying to quit.
+                pass
+
+
+    @classmethod
+    def _initialise(cls, cache_dir=None):
+        if cache_dir is not None:
+            cls._cache_dir = cache_dir
+
+        if not _os.path.isdir(cls._cache_dir):
+            _os.mkdir(cls._cache_dir)
+        cls._db = _peewee.SqliteDatabase(
+            _os.path.join(cls._cache_dir, 'tkr-tz.db'),
+            pragmas={'journal_mode': 'wal', 'cache_size': -64}
+        )
+
+        old_cache_file_path = _os.path.join(cls._cache_dir, "tkr-tz.csv")
+        if _os.path.isfile(old_cache_file_path):
+            _os.remove(old_cache_file_path)
+
+    @classmethod
+    def change_location(cls, new_cache_dir):
+        cls._db.close()
+        cls._db = None
+        cls._cache_dir = new_cache_dir
+# close DB when Python exists
+_atexit.register(_DBManager.close_db)
+
+
+class _KV(_peewee.Model):
     key = _peewee.CharField(primary_key=True)
     value = _peewee.CharField(null=True)
 
     class Meta:
-        database = db
+        database = _DBManager.get_database()
         without_rowid = True
 
 
 class _TzCache:
     def __init__(self):
+        db = _DBManager.get_database()
         db.connect()
-        db.create_tables([KV])
+        db.create_tables([_KV])
 
-        old_cache_file_path = _os.path.join(_cache_dir, "tkr-tz.csv")
-        if _os.path.isfile(old_cache_file_path):
-            _os.remove(old_cache_file_path)
 
     def lookup(self, key):
         try:
-            return KV.get(KV.key == key).value
-        except KV.DoesNotExist:
+            return _KV.get(_KV.key == key).value
+        except _KV.DoesNotExist:
             return None
 
     def store(self, key, value):
+        db = _DBManager.get_database()
         try:
             if value is None:
-                q = KV.delete().where(KV.key == key)
+                q = _KV.delete().where(_KV.key == key)
                 q.execute()
                 return
             with db.atomic():
-                KV.insert(key=key, value=value).execute()
-        except IntegrityError:
+                _KV.insert(key=key, value=value).execute()
+        except _peewee.IntegrityError:
             # Integrity error means the key already exists. Try updating the key.
             old_value = self.lookup(key)
             if old_value != value:
                 get_yf_logger().debug(f"Value for key {key} changed from {old_value} to {value}.")
                 with db.atomic():
-                    q = KV.update(value=value).where(KV.key == key)
+                    q = _KV.update(value=value).where(_KV.key == key)
                     q.execute()
 
-    def close(self):
-        db.close()
-
-
-class _TzCacheDummy:
-    """Dummy cache to use if tz cache is disabled"""
-
-    def lookup(self, tkr):
-        return None
-
-    def store(self, tkr, tz):
-        pass
-
-    @property
-    def tz_db(self):
-        return None
-
 
 def get_tz_cache():
     """
@@ -979,20 +1037,7 @@ def get_tz_cache():
     """
     # as this can be called from multiple threads, protect it.
     with _cache_init_lock:
-        global _tz_cache
-        if _tz_cache is None:
-            try:
-                _tz_cache = _TzCache()
-            except _TzCacheException as err:
-                get_yf_logger().info(f"Failed to create TzCache, reason: {err}. "
-                                     "TzCache will not be used. "
-                                     "Tip: You can direct cache to use a different location with 'set_tz_cache_location(mylocation)'")
-                _tz_cache = _TzCacheDummy()
-
-        return _tz_cache
-
-
-_cache_init_lock = Lock()
+        return _TzCacheManager.get_tz()
 
 
 def set_tz_cache_location(cache_dir: str):
@@ -1003,6 +1048,4 @@ def set_tz_cache_location(cache_dir: str):
     :param cache_dir: Path to use for caches
     :return: None
     """
-    global _cache_dir, _tz_cache
-    assert _tz_cache is None, "Time Zone cache already initialized, setting path must be done before cache is created"
-    _cache_dir = cache_dir
+    _DBManager.change_location(cache_dir)