sitemap-generator/pysitemap/backends/sqlite_todo.py

import logging
import sqlite3


class SQLiteTodo(object):
    def __init__(self, db_name):
        self.connection = sqlite3.connect(db_name)
        self.__init_tables()

    def __init_tables(self):
        cursor = self.connection.cursor()
        cursor.execute("DROP TABLE IF EXISTS todo_queue;")
        cursor.execute("""
            CREATE TABLE todo_queue (
                url text(1000) primary key
            );
        """)
        self.connection.commit()
        cursor.close()

    def add(self, url):
        cursor = self.connection.cursor()
        try:
            cursor.execute("""insert into todo_queue values (?);""", (url,))
        except Exception as e:
            logging.info(e)
        finally:
            self.connection.commit()
            cursor.close()

    def remove(self, url):
        cursor = self.connection.cursor()
        try:
            cursor.execute("""delete from todo_queue where url = ?;""", (url,))
        except Exception as e:
            logging.info(e)
        finally:
            self.connection.commit()
            cursor.close()

    def __contains__(self, item):
        cursor = self.connection.cursor()
        result = False
        try:
            cursor.execute("""select 1 from todo_queue where url = ?""", (item, ))
            row = cursor.fetchone()
            if len(row):
                result = True
        except Exception as e:
            logging.info(e)
        finally:
            cursor.close()
        return result

    def __iter__(self):
        cursor = self.connection.cursor()
        result = []
        try:
            cursor.execute("""select url from todo_queue""")
            rows = cursor.fetchall()
            result = [row[0] for row in rows]
        except Exception as e:
            logging.info(e)
        finally:
            cursor.close()
        return iter(result)

    def __next__(self):
        for url in self:
            yield url

    def __len__(self):
        cursor = self.connection.cursor()
        result = []
        try:
            cursor.execute("""select count(*) as cnt from todo_queue""")
            row = cursor.fetchone()
            result = row[0]
        except Exception as e:
            logging.info(e)
        finally:
            cursor.close()
        return result