YT: Make sure responses ok.
[audio.git] / src / apiclient / models.py
1 import json
2 import os
3 from django.db import models
4 from django.contrib.auth.models import User
5 from requests_oauthlib import OAuth2Session
6 from requests_toolbelt.streaming_iterator import StreamingIterator
7 from .settings import YOUTUBE_CLIENT_ID, YOUTUBE_CLIENT_SECRET, YOUTUBE_TOKEN_URL
8
9
10 class OAuthConnection(models.Model):
11     user = models.ForeignKey(User, models.CASCADE)
12     access = models.BooleanField(default=False)
13     token = models.CharField(max_length=64, null=True, blank=True)
14     token_secret = models.CharField(max_length=64, null=True, blank=True)
15
16     @classmethod
17     def get(cls, user):
18         try:
19             return cls.objects.get(user=user)
20         except cls.DoesNotExist:
21             o = cls(user=user)
22             o.save()
23             return o
24
25
26 class YouTubeToken(models.Model):
27     token = models.TextField()
28
29     def token_updater(self, token):
30         self.token = json.dumps(token)
31         self.save()
32
33     def get_session(self):
34         return OAuth2Session(
35             client_id=YOUTUBE_CLIENT_ID,
36             auto_refresh_url=YOUTUBE_TOKEN_URL,
37             token=json.loads(self.token),
38             auto_refresh_kwargs={'client_id':YOUTUBE_CLIENT_ID,'client_secret':YOUTUBE_CLIENT_SECRET},
39             token_updater=self.token_updater
40         )
41
42     def call(self, method, url, params=None, json=None, data=None, resumable_file_path=None):
43         params = params or {}
44         if resumable_file_path:
45             params['uploadType'] = 'resumable'
46             file_size = os.stat(resumable_file_path).st_size
47
48         session = self.get_session()
49         response = session.request(
50             method=method,
51             url=url,
52             json=json,
53             data=data,
54             params=params,
55             headers={
56                 'X-Upload-Content-Length': str(file_size),
57                 'x-upload-content-type': 'application/octet-stream',
58             } if resumable_file_path else {}
59         )
60         response.raise_for_status()
61         if resumable_file_path:
62             location = response.headers['Location']
63             with open(resumable_file_path, 'rb') as f:
64                 response = session.put(
65                     url=location,
66                     data=StreamingIterator(file_size, f),
67                     headers={"Content-Type": "application/octet-stream"},
68                 )
69                 response.raise_for_status()
70                 return response