BlenderKit: Basic private quota checking works now.

6 files changed, 64 insertions, 29 deletions

diff --git a/blenderkit/download.py b/blenderkit/download.py
index bffb2e8a..c8695f5e 100644
--- a/blenderkit/download.py
+++ b/blenderkit/download.py
@@ -425,7 +425,6 @@ def append_asset(asset_data, **kwargs):  # downloaders=[], location=None,
         else:
             target_object.material_slots[kwargs['material_target_slot']].material = material
 
-        props = material.blenderkit
         parent = material
 
     parent['asset_data'] = asset_data  # TODO remove this??? should write to blenderkit Props?
diff --git a/blenderkit/oauth.py b/blenderkit/oauth.py
index b0359157..d414c5a1 100644
--- a/blenderkit/oauth.py
+++ b/blenderkit/oauth.py
@@ -125,7 +125,7 @@ def write_tokens(auth_token, refresh_token):
     preferences.api_key_refresh = refresh_token
     preferences.login_attempt = False
     props = utils.get_search_props()
-
+    search.get_profile()
     props.report = 'Login success!'
     search.get_profile()
 
diff --git a/blenderkit/search.py b/blenderkit/search.py
index 9f19bf41..7faf6f18 100644
--- a/blenderkit/search.py
+++ b/blenderkit/search.py
@@ -64,8 +64,8 @@ def check_errors(rdata):
     if rdata.get('statusCode') == 401:
         if rdata.get('detail') == 'Invalid token.':
             # reset the api key, so it can be requested again.
-            user_preferences = bpy.context.preferences.addons['blenderkit'].preferences
-            user_preferences.api_key = ''
+            # user_preferences = bpy.context.preferences.addons['blenderkit'].preferences
+            # user_preferences.api_key = ''
             return False, 'Missing or wrong api_key in addon preferences'
     return True, ''
 
@@ -448,6 +448,7 @@ def generate_tooltip(mdata):
 
     return t
 
+
 def get_random_tip():
     if at == 'brush' or at == 'texture':
         t += 'click to link %s' % mdata['assetType']
@@ -459,6 +460,7 @@ def get_random_tip():
         tip = 'Tip: ' + random.choice(tips)
         t = writeblock(t, tip)
 
+
 def generate_author_textblock(adata):
     t = ''
     if adata not in (None, ''):
@@ -542,7 +544,6 @@ def fetch_author(a_id, api_key):
 def get_author(r):
     a_id = str(r['author']['id'])
     preferences = bpy.context.preferences.addons['blenderkit'].preferences
-
     authors = bpy.context.window_manager.get('bkit authors', {})
     if authors == {}:
         bpy.context.window_manager['bkit authors'] = authors
@@ -557,25 +558,34 @@ def get_author(r):
 def write_profile(adata):
     utils.p('writing profile')
     utils.p(adata.keys())
-    adata['user']['sumAssetFilesSize'] = str(round(adata['user']['sumAssetFilesSize'] / 1024 / 1024)) + ' Mb'
-    adata['user']['sumPrivateAssetFilesSize'] = str(
-        round(adata['user']['sumPrivateAssetFilesSize'] / 1024 / 1024)) + ' Mb'
-    adata['user']['remainingPrivateQuota'] = str(max(0,round(adata['user']['remainingPrivateQuota'] / 1024 / 1024))) + ' Mb'
+    print(adata)
+    user = adata['user']
+    # we have to convert to MB here, numbers too big for python int type
+    user['sumAssetFilesSize'] /= (1024 * 1024)
+    user['sumPrivateAssetFilesSize'] /= (1024 * 1024)
+    user['remainingPrivateQuota'] /= (1024 * 1024)
+
     bpy.context.window_manager['bkit profile'] = adata
 
 
+def request_profile(api_key):
+    a_url = paths.get_api_url() + 'me/'
+    headers = utils.get_headers(api_key)
+    r = requests.get(a_url, headers=headers)
+    adata = r.json()
+    if adata.get('user') is None:
+        utils.p(adata)
+        utils.p('getting profile failed')
+        return None
+    return adata
+
+
 def fetch_profile(api_key):
     utils.p('fetch profile')
     try:
-        a_url = paths.get_api_url() + 'me/'
-        headers = utils.get_headers(api_key)
-        r = requests.get(a_url, headers=headers)
-        adata = r.json()
-        if adata.get('user') is None:
-            utils.p(adata)
-            utils.p('getting profile failed')
-            return
-        tasks_queue.add_task((write_profile, (adata,)))
+        adata = request_profile(api_key)
+        if adata is not None:
+            tasks_queue.add_task((write_profile, (adata,)))
     except Exception as e:
         utils.p(e)
 
diff --git a/blenderkit/tasks_queue.py b/blenderkit/tasks_queue.py
index 2e09dcfd..9cf3720f 100644
--- a/blenderkit/tasks_queue.py
+++ b/blenderkit/tasks_queue.py
@@ -5,11 +5,13 @@ import queue
 
 from blenderkit import utils
 
+
 @persistent
 def scene_load(context):
-    if not(bpy.app.timers.is_registered(queue_worker)):
+    if not (bpy.app.timers.is_registered(queue_worker)):
         bpy.app.timers.register(queue_worker)
 
+
 def get_queue():
     # we pick just a random one of blender types, to try to get a persistent queue
     t = bpy.types.Scene
@@ -26,11 +28,12 @@ def add_task(task):
 
 def queue_worker():
     q = get_queue()
-    utils.p('queue timer')
+    # utils.p('queue timer')
     while not q.empty():
         utils.p('as a task:   ')
         print('window manager', bpy.context.window_manager)
         task = q.get()
+        utils.p(task)
         try:
             task[0](*task[1])
         except Exception as e:
@@ -43,7 +46,5 @@ def register():
     bpy.app.handlers.load_post.append(scene_load)
 
 
-
 def unregister():
     bpy.app.handlers.load_post.remove(scene_load)
-
diff --git a/blenderkit/ui_panels.py b/blenderkit/ui_panels.py
index 478a1b54..c5292f19 100644
--- a/blenderkit/ui_panels.py
+++ b/blenderkit/ui_panels.py
@@ -128,15 +128,15 @@ def draw_upload_common(layout, props, asset_type, context):
             row = layout.row()
 
             row.prop(props, 'asset_base_id', icon='FILE_TICK')
-        layout.operator("object.blenderkit_mark_for_validation", icon='EXPORT')
+        # layout.operator("object.blenderkit_mark_for_validation", icon='EXPORT')
 
     layout.prop(props, 'category')
     if asset_type == 'MODEL' and props.subcategory != '':  # by now block this for other asset types.
         layout.prop(props, 'subcategory')
 
     layout.prop(props, 'is_private')
-    layout.prop(props, 'license')
-
+    if not props.is_private:
+        layout.prop(props, 'license')
 
 
 def poll_local_panels():
@@ -416,9 +416,10 @@ class VIEW3D_PT_blenderkit_profile(Panel):
                 me = me['user']
                 layout.label(text='User: %s %s' % (me['firstName'], me['lastName']))
                 layout.label(text='Email: %s' % (me['email']))
-                layout.label(text='Public assets: %s ' % (me['sumAssetFilesSize']))
-                layout.label(text='Private assets: %s ' % (me['sumPrivateAssetFilesSize']))
-                layout.label(text='Remaining private storage: %s' % (me['remainingPrivateQuota']))
+                if me.get('sumAssetFilesSize') is not None:  # TODO remove this when production server has these too.
+                    layout.label(text='Public assets: %i Mb' % (me['sumAssetFilesSize']))
+                    layout.label(text='Private assets: %i Mb' % (me['sumPrivateAssetFilesSize']))
+                    layout.label(text='Remaining private storage: %i Mb' % (me['remainingPrivateQuota']))
             layout.operator("wm.url_open", text="See my uploads",
                             icon='URL').url = paths.BLENDERKIT_USER_ASSETS
             layout.operator("wm.blenderkit_logout", text="Logout",
diff --git a/blenderkit/upload.py b/blenderkit/upload.py
index 15c0bf83..30e17f27 100644
--- a/blenderkit/upload.py
+++ b/blenderkit/upload.py
@@ -22,11 +22,12 @@ if "bpy" in locals():
     imp.reload(asset_inspector)
     imp.reload(paths)
     imp.reload(utils)
+    imp.reload(search)
     imp.reload(bg_blender)
     imp.reload(autothumb)
     imp.reload(version_checker)
 else:
-    from blenderkit import asset_inspector, paths, utils, bg_blender, autothumb, version_checker
+    from blenderkit import asset_inspector, paths, utils, bg_blender, autothumb, version_checker, search
 
 import tempfile, os, subprocess, json, re
 
@@ -500,9 +501,32 @@ def get_upload_location(props):
         return None
     return None
 
+def check_storage_quota(props):
+    if not props.is_private:
+        return True
+
+    profile = bpy.context.window_manager.get('bkit profile')
+    if profile is None or profile.get('remainingPrivateQuota') is None:
+        preferences = bpy.context.preferences.addons['blenderkit'].preferences
+        adata = search.request_profile(preferences.api_key)
+        if adata is None:
+            props.report = 'User profile not retrieved.'
+            return False
+        search.write_profile(adata)
+        profile = adata
+    print(profile.keys())
+    if profile['user'].get('remainingPrivateQuota')>0:
+        return True
+    props.report = 'Private storage quota exceeded.'
+    return False
 
 def start_upload(self, context, asset_type, as_new, metadata_only):
     props = utils.get_upload_props()
+    storage_quota_ok = check_storage_quota(props)
+    if not storage_quota_ok:
+        self.report({'ERROR_INVALID_INPUT'}, props.report)
+        return {'CANCELLED'}
+
     location = get_upload_location(props)
     props.upload_state = 'preparing upload'
     # do this for fixing long tags in some upload cases