Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions exceptions/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .exceptions import JobNotFoundException, StudioException, PhotoNotFoundException

__all__ = ['JobNotFoundException', 'StudioException', 'PhotoNotFoundException']
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
class JobNotFoundException(Exception):
def __init__(self, message):
super().__init__(message)
class PhotoNotFoundException(Exception):
def __init__(self, message):
super().__init__(message)

class StudioException(Exception):
def __init__(self, status_code, message="Studio exception occurred"):
self.status_code = status_code
self.message = message
super().__init__(self.message, self.status_code)

1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@ pyvips==2.0.2
requests>=2.0.0
requests_mock>=1.5.2
decouple>=0.0.7
sentry-sdk>=2.13.0
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

setup(
name='skylab_studio',
version='0.0.14',
version='0.0.15',
author='skylabtech',
author_email='info@skylabtech.ai',
packages=find_packages(),
Expand Down
80 changes: 51 additions & 29 deletions skylab_studio/studio_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,10 @@
import base64
import hashlib
import requests
from io import BytesIO
import sentry_sdk

from .version import VERSION
from .studio_exception import StudioException
from exceptions import *

API_HEADER_KEY = 'X-SLT-API-KEY'
API_HEADER_CLIENT = 'X-SLT-API-CLIENT'
Expand Down Expand Up @@ -62,7 +62,7 @@ def __init__(self, api_key=None, **kwargs):

if 'debug' in kwargs:
self.debug = kwargs['debug']

if 'max_concurrent_downloads' in kwargs:
self.max_concurrent_downloads = kwargs['max_concurrent_downloads']

Expand All @@ -72,6 +72,19 @@ def __init__(self, api_key=None, **kwargs):
LOGGER.debug('Debug enabled')
LOGGER.propagate = True

# initialize sentry
sentry_sdk.init(
dsn="https://0b5490403ee70db8bd7869af3b10380b@o1409269.ingest.us.sentry.io/4507850876452864",
# Set traces_sample_rate to 1.0 to capture 100%
# of transactions for tracing.
traces_sample_rate=1.0,
# Set profiles_sample_rate to 1.0 to profile 100%
# of sampled transactions.
# We recommend adjusting this value in production.
profiles_sample_rate=1.0,
ignore_errors=[JobNotFoundException, PhotoNotFoundException]
)

def _build_http_auth(self):
return (self.api_key, '')

Expand Down Expand Up @@ -302,14 +315,19 @@ def _upload_photo(self, photo_path, id, model='job'):
photo_data = { f"{model}_id": id, "name": photo_name, "use_cache_upload": False }

if model == 'job':
job_type = self.get_job(id)['type']
job = self.get_job(id)
if "type" in job:
job_type = job['type']
if job_type == 'regular':
headers = { 'X-Amz-Tagging': 'job=photo&api=true' }
else:
raise JobNotFoundException(f"Unable to find job with id: {id}")

if job_type == 'regular':
headers = { 'X-Amz-Tagging': 'job=photo&api=true' }

# Ask studio to create the photo record
photo_resp = self._create_photo(photo_data)
if not photo_resp:

if not 'id' in photo_resp:
raise Exception('Unable to create the photo object, if creating profile photo, ensure enable_extract and replace_background is set to: True')

photo_id = photo_resp['id']
Expand All @@ -329,28 +347,30 @@ def _upload_photo(self, photo_path, id, model='job'):
# PUT request to presigned url with image data
headers["Content-MD5"] = b64md5

try:
upload_photo_resp = requests.put(upload_url, data, headers=headers)

if not upload_photo_resp:
print('First upload attempt failed, retrying...')
retry = 0
# retry upload
while retry < 3:
upload_photo_resp = requests.put(upload_url, data, headers=headers)

if upload_photo_resp:
break # Upload was successful, exit the loop
elif retry == 2: # Check if retry count is 2 (0-based indexing)
raise Exception('Unable to upload to the bucket after retrying.')
else:
time.sleep(1) # Wait for a moment before retrying
retry += 1
retry = 0
while retry < 3:
try:
# attempt to upload the photo to aws
upload_photo_resp = requests.put(upload_url, data, headers=headers)

except Exception as e:
print(f"An exception of type {type(e).__name__} occurred: {e}")
print('Deleting created, but unuploaded photo...')
self.delete_photo(photo_id)
# Will raise exception for any statuses 4xx-5xx
upload_photo_resp.raise_for_status()

# if raise_for_status didn't throw an exception, then we successfully uploaded, exit the loop
break

# rescue any exceptions in the loop
except Exception as e:
# if we've retried 3 times, delete the photo record and raise exception
if retry == 2:
self.delete_photo(photo_id)

raise Exception(e)
# if we haven't retried 3 times, wait for retry+1 seconds and continue the while loop
else:
print(f"Attempt #{retry + 1} to upload failed, retrying...")
retry += 1
time.sleep(retry+1)

res['upload_response'] = upload_photo_resp.status_code
return res
Expand Down Expand Up @@ -488,6 +508,9 @@ async def download_photo(self, photo_id, output_path, profile = None, options =
await semaphore.acquire()

photo = self.get_photo(photo_id)

if not 'job' in photo:
raise PhotoNotFoundException(f"Unable to find photo with id: {photo_id}")
profile_id = photo['job']['profileId']
file_name = photo['name']

Expand Down Expand Up @@ -531,4 +554,3 @@ async def download_photo(self, photo_id, output_path, profile = None, options =
finally:
if semaphore != None:
semaphore.release()

2 changes: 1 addition & 1 deletion skylab_studio/version.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@
For more information, visit https://studio.skylabtech.ai
"""

VERSION = '0.0.14'
VERSION = '0.0.15'