server/tags: export also tag categories

5 files changed, 103 insertions, 32 deletions

diff --git a/server/szurubooru/db/tag.py b/server/szurubooru/db/tag.py
index bf98745..55e1857 100644
--- a/server/szurubooru/db/tag.py
+++ b/server/szurubooru/db/tag.py
@@ -47,18 +47,20 @@ class Tag(Base):
     creation_time = Column('creation_time', DateTime, nullable=False)
     last_edit_time = Column('last_edit_time', DateTime)
 
-    category = relationship('TagCategory')
-    names = relationship('TagName', cascade='all, delete-orphan')
+    category = relationship('TagCategory', lazy='joined')
+    names = relationship('TagName', cascade='all, delete-orphan', lazy='joined')
     suggestions = relationship(
         'Tag',
         secondary='tag_suggestion',
         primaryjoin=tag_id == TagSuggestion.parent_id,
-        secondaryjoin=tag_id == TagSuggestion.child_id)
+        secondaryjoin=tag_id == TagSuggestion.child_id,
+        lazy='joined')
     implications = relationship(
         'Tag',
         secondary='tag_implication',
         primaryjoin=tag_id == TagImplication.parent_id,
-        secondaryjoin=tag_id == TagImplication.child_id)
+        secondaryjoin=tag_id == TagImplication.child_id,
+        lazy='joined')
 
     post_count = column_property(
         select([func.count('Post.post_id')]) \
diff --git a/server/szurubooru/search/search_executor.py b/server/szurubooru/search/search_executor.py
index 04d1aa8..d441cd2 100644
--- a/server/szurubooru/search/search_executor.py
+++ b/server/szurubooru/search/search_executor.py
@@ -21,13 +21,17 @@ class SearchExecutor(object):
         entities = filter_query \
             .offset((page - 1) * page_size).limit(page_size).all()
         count_query = filter_query.statement \
-            .with_only_columns([sqlalchemy.func.count()]).order_by(None)
-        count = filter_query.session.execute(count_query).scalar()
+            .with_only_columns([sqlalchemy.func.count()]) \
+            .order_by(None)
+        count = filter_query \
+            .session.execute(count_query) \
+            .scalar()
         return (count, entities)
 
     def _prepare(self, query_text):
         ''' Parse input and return SQLAlchemy query. '''
-        query = self._search_config.create_query()
+        query = self._search_config.create_query() \
+            .options(sqlalchemy.orm.lazyload('*'))
         for token in re.split(r'\s+', (query_text or '').lower()):
             if not token:
                 continue
diff --git a/server/szurubooru/tests/api/test_tag_export.py b/server/szurubooru/tests/api/test_tag_export.py
index 0840acf..3e0eb1c 100644
--- a/server/szurubooru/tests/api/test_tag_export.py
+++ b/server/szurubooru/tests/api/test_tag_export.py
@@ -4,17 +4,27 @@ import json
 from szurubooru import config, db
 from szurubooru.util import tags
 
-def test_export(tmpdir, session, config_injector, tag_factory):
+def test_export(
+        tmpdir,
+        query_counter,
+        session,
+        config_injector,
+        tag_factory,
+        tag_category_factory):
     config_injector({
         'data_dir': str(tmpdir)
     })
-    sug1 = tag_factory(names=['sug1'])
-    sug2 = tag_factory(names=['sug2'])
-    imp1 = tag_factory(names=['imp1'])
-    imp2 = tag_factory(names=['imp2'])
-    tag = tag_factory(names=['alias1', 'alias2'])
+    cat1 = tag_category_factory(name='cat1', color='black')
+    cat2 = tag_category_factory(name='cat2', color='white')
+    session.add_all([cat1, cat2])
+    session.flush()
+    sug1 = tag_factory(names=['sug1'], category=cat1)
+    sug2 = tag_factory(names=['sug2'], category=cat1)
+    imp1 = tag_factory(names=['imp1'], category=cat1)
+    imp2 = tag_factory(names=['imp2'], category=cat1)
+    tag = tag_factory(names=['alias1', 'alias2'], category=cat2)
     tag.post_count = 1
-    session.add_all([tag, sug1, sug2, imp1, imp2])
+    session.add_all([tag, sug1, sug2, imp1, imp2, cat1, cat2])
     session.flush()
     session.add_all([
         db.TagSuggestion(tag.tag_id, sug1.tag_id),
@@ -24,19 +34,29 @@ def test_export(tmpdir, session, config_injector, tag_factory):
     ])
     session.flush()
 
-    tags.export_to_json()
+    with query_counter:
+        tags.export_to_json()
+        assert len(query_counter.statements) == 2
+
     export_path = os.path.join(config.config['data_dir'], 'tags.json')
     assert os.path.exists(export_path)
     with open(export_path, 'r') as handle:
-        assert json.loads(handle.read()) == [
-            {
-                'names': ['alias1', 'alias2'],
-                'usages': 1,
-                'suggestions': ['sug1', 'sug2'],
-                'implications': ['imp1', 'imp2'],
-            },
-            {'names': ['sug1'], 'usages': 0},
-            {'names': ['sug2'], 'usages': 0},
-            {'names': ['imp1'], 'usages': 0},
-            {'names': ['imp2'], 'usages': 0},
-        ]
+        assert json.loads(handle.read()) == {
+            'tags': [
+                {
+                    'names': ['alias1', 'alias2'],
+                    'usages': 1,
+                    'category': 'cat2',
+                    'suggestions': ['sug1', 'sug2'],
+                    'implications': ['imp1', 'imp2'],
+                },
+                {'names': ['sug1'], 'usages': 0, 'category': 'cat1'},
+                {'names': ['sug2'], 'usages': 0, 'category': 'cat1'},
+                {'names': ['imp1'], 'usages': 0, 'category': 'cat1'},
+                {'names': ['imp2'], 'usages': 0, 'category': 'cat1'},
+            ],
+            'categories': [
+                {'name': 'cat1', 'color': 'black'},
+                {'name': 'cat2', 'color': 'white'},
+            ]
+        }
diff --git a/server/szurubooru/tests/conftest.py b/server/szurubooru/tests/conftest.py
index 2981e10..6b70fa7 100644
--- a/server/szurubooru/tests/conftest.py
+++ b/server/szurubooru/tests/conftest.py
@@ -7,6 +7,32 @@ import sqlalchemy
 from szurubooru import api, config, db
 from szurubooru.util import misc
 
+class QueryCounter(object):
+    def __init__(self):
+        self._statements = []
+
+    def __enter__(self):
+        self._statements = []
+
+    def __exit__(self, *args, **kwargs):
+        self._statements = []
+
+    def create_before_cursor_execute(self):
+        def before_cursor_execute(
+                _conn, _cursor, statement, _parameters, _context, _executemany):
+            self._statements.append(statement)
+        return before_cursor_execute
+
+    @property
+    def statements(self):
+        return self._statements
+
+_query_counter = QueryCounter()
+
+@pytest.fixture
+def query_counter():
+    return _query_counter
+
 def get_unique_name():
     return str(uuid.uuid4())
 
@@ -21,11 +47,15 @@ def fake_datetime():
     return injector
 
 @pytest.yield_fixture
-def session(autoload=True):
+def session(query_counter, autoload=True):
     import logging
     logging.basicConfig()
     logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
     engine = sqlalchemy.create_engine('sqlite:///:memory:')
+    sqlalchemy.event.listen(
+        engine,
+        'before_cursor_execute',
+        query_counter.create_before_cursor_execute())
     session_maker = sqlalchemy.orm.sessionmaker(bind=engine)
     session = sqlalchemy.orm.scoped_session(session_maker)
     db.Base.query = session.query_property()
diff --git a/server/szurubooru/util/tags.py b/server/szurubooru/util/tags.py
index 24b8cad..69cc95e 100644
--- a/server/szurubooru/util/tags.py
+++ b/server/szurubooru/util/tags.py
@@ -28,11 +28,21 @@ def _check_name_intersection(names1, names2):
     return len(set(_lower_list(names1)).intersection(_lower_list(names2))) > 0
 
 def export_to_json():
-    output = []
-    for tag in db.session().query(db.Tag).all():
+    output = {
+        'tags': [],
+        'categories': [],
+    }
+    all_tags = db.session() \
+        .query(db.Tag) \
+        .options(
+            sqlalchemy.orm.joinedload('suggestions'),
+            sqlalchemy.orm.joinedload('implications')) \
+        .all()
+    for tag in all_tags:
         item = {
             'names': [tag_name.name for tag_name in tag.names],
-            'usages': tag.post_count
+            'usages': tag.post_count,
+            'category': tag.category.name,
         }
         if len(tag.suggestions):
             item['suggestions'] = \
@@ -40,7 +50,12 @@ def export_to_json():
         if len(tag.implications):
             item['implications'] = \
                 [rel.names[0].name for rel in tag.implications]
-        output.append(item)
+        output['tags'].append(item)
+    for category in tag_categories.get_all_categories():
+        output['categories'].append({
+            'name': category.name,
+            'color': category.color,
+        })
     export_path = os.path.join(config.config['data_dir'], 'tags.json')
     with open(export_path, 'w') as handle:
         handle.write(json.dumps(output, separators=(',', ':')))