add doctests

airflow/hooks/sqlite_hook.py

@@ -23,29 +23,14 @@ class SqliteHook(BaseHook):
         conn = sqlite3.connect(conn.host)
         return conn
 
-    def get_records(self, sql):
-        '''
-        Executes the sql and returns a set of records.
-        '''
-        conn = self.get_conn()
-        cur = conn.cursor()
-        cur.execute(sql)
-        rows = cur.fetchall()
-        cur.close()
-        conn.close()
-        return rows
-
-    def get_pandas_df(self, sql):
-        '''
-        Executes the sql and returns a pandas dataframe
-        '''
-        import pandas.io.sql as psql
-        conn = self.get_conn()
-        df = psql.read_sql(sql, con=conn)
-        conn.close()
-        return df
-
     def run(self, sql):
+        """
+        Runs a command
+
+        >>> h = SqliteHook()
+        >>> sql = "CREATE TABLE IF NOT EXISTS test_table (i INTEGER);"
+        >>> h.run(sql)
+        """
         conn = self.get_conn()
         cur = conn.cursor()
         cur.execute(sql)
@@ -57,6 +42,9 @@ class SqliteHook(BaseHook):
         """
         A generic way to insert a set of tuples into a table,
         the whole set of inserts is treated as one transaction
+
+        >>> h = SqliteHook()
+        >>> h.insert_rows('test_table', [[1]])
         """
         if target_fields:
             target_fields = ", ".join(target_fields)
@@ -88,3 +76,35 @@ class SqliteHook(BaseHook):
         conn.close()
         logging.info(
             "Done loading. Loaded a total of {i} rows".format(**locals()))
+
+    def get_records(self, sql):
+        """
+        Executes the sql and returns a set of records.
+
+        >>> h = SqliteHook()
+        >>> sql = "SELECT * FROM test_table WHERE i=1 LIMIT 1;"
+        >>> h.get_records(sql)
+        [(1,)]
+        """
+        conn = self.get_conn()
+        cur = conn.cursor()
+        cur.execute(sql)
+        rows = cur.fetchall()
+        cur.close()
+        conn.close()
+        return rows
+
+    def get_pandas_df(self, sql):
+        """
+        Executes the sql and returns a pandas dataframe
+        >>> h = SqliteHook()
+        >>> sql = "SELECT * FROM test_table WHERE i=1 LIMIT 1;"
+        >>> h.get_pandas_df(sql)
+           i
+        0  1
+        """
+        import pandas.io.sql as psql
+        conn = self.get_conn()
+        df = psql.read_sql(sql, con=conn)
+        conn.close()
+        return df