formatting

Author: benjamingregory

hooks/__init__.py

@@ -1,14 +1,13 @@
-from airflow.hooks.base_hook import BaseHook
-import json
 import stripe
 
+from airflow.hooks.base_hook import BaseHook
+
 
 class StripeHook(BaseHook):
-    def __init__(
-            self,
-            conn_id,
-            *args,
-            **kwargs):
+    def __init__(self,
+                 conn_id,
+                 *args,
+                 **kwargs):
         self.conn_id = conn_id
         self._args = args
         self._kwargs = kwargs
@@ -27,37 +26,32 @@ def get_conn(self):
         self.connection = self.get_connection(self.conn_id)
         self.extras = self.connection.extra_dejson
 
-        stripe.api_key = self.extras['api_key']
+        stripe.api_key = self.extras.get('api_key', None)
+
         self.stripe = stripe
 
         return stripe
 
-    def run_query(self, model, replication_key_value=None, **kwargs):
+    def run_query(self,
+                  stripe_object,
+                  replication_key_value=None,
+                  **kwargs):
         """
         Run a query against stripe
-        :param model:                   name of the Stripe model
+        :param stripe_object:           name of the Stripe object
         :param replication_key_value:   Stripe replicaton key value
         """
         stripe_instance = self.get_conn()
-        stripe_model = getattr(stripe_instance, model)
+        stripe_endpoint = getattr(stripe_instance, stripe_object)
 
         method_to_call = 'list'
-        if model is 'BalanceHistory':
+        if stripe_object is 'BalanceHistory':
             method_to_call = 'all'
         if replication_key_value:
-            stripe_response = getattr(stripe_model, method_to_call)(
+            stripe_response = getattr(stripe_endpoint, method_to_call)(
                 ending_before=replication_key_value, **kwargs)
         else:
-            stripe_response = getattr(stripe_model, method_to_call)(**kwargs)
+            stripe_response = getattr(stripe_endpoint, method_to_call)(**kwargs)
 
         for res in stripe_response.auto_paging_iter():
             yield res
-
-    def get_records(self, sql):
-        pass
-
-    def get_pandas_df(self, sql):
-        pass
-
-    def run(self, sql):
-        pass