|
|
@@ -1,15 +1,45 @@
|
|
|
from sqlalchemy import create_engine
|
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
|
+import inspect
|
|
|
+from ruamel import yaml
|
|
|
|
|
|
|
|
|
SQLABase = declarative_base()
|
|
|
|
|
|
|
|
|
+yaml = yaml.YAML()
|
|
|
+yaml.default_flow_style = False
|
|
|
+yaml.allow_unicode = True
|
|
|
+
|
|
|
class Database:
|
|
|
def __init__(self, db_url):
|
|
|
self.engine = create_engine(db_url)
|
|
|
self.session_factory = sessionmaker(bind=self.engine)
|
|
|
|
|
|
SQLABase.metadata.create_all(self.engine)
|
|
|
+ self.models = {
|
|
|
+ k: v for k, v in SQLABase._decl_class_registry.items()
|
|
|
+ if inspect.isclass(v) and issubclass(v, SQLABase)
|
|
|
+ }
|
|
|
+
|
|
|
+ def dump(self, file):
|
|
|
+ session = self.session_factory()
|
|
|
+ instances_dict = {}
|
|
|
+
|
|
|
+ for k, v in self.models.items():
|
|
|
+ instances = [i for i in session.query(v).all()]
|
|
|
+ instances_dict[k] = instances
|
|
|
+ yaml.register_class(v)
|
|
|
+
|
|
|
+ with open(file, 'w') as f:
|
|
|
+ yaml.dump(instances_dict, f)
|
|
|
+
|
|
|
+ def load(self, file):
|
|
|
+ with open(file) as f:
|
|
|
+ instances_dict = yaml.load(f)
|
|
|
|
|
|
+ session = self.session_factory()
|
|
|
+ for k, v_list in instances_dict.items():
|
|
|
+ model = self.models[k]
|
|
|
+ instances = [model(**v) for v in v_list]
|