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         if resumable_file_path:
  61             location = response.headers['Location']
  62             with open(resumable_file_path, 'rb') as f:
  63                 return session.put(
  64                     url=location,
  65                     data=StreamingIterator(file_size, f),
  66                     headers={"Content-Type": "application/octet-stream"},
  67                 )