db.py

import os
import psycopg2

DATABASE_URL = os.environ.get('DATABASE_URL', False)


def init():
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('CREATE TABLE IF NOT EXISTS public.numbers(id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, number text NOT NULL)')
  cursor.execute('CREATE TABLE IF NOT EXISTS public.born(id int GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, message text NOT NULL)')
  conn.commit()


def number_exists(number):
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('SELECT exists(SELECT number FROM public.numbers WHERE number = %s)', (number,))
  return cursor.fetchone()[0]


def get_all_numbers():
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('SELECT number FROM public.numbers')
  return cursor.fetchall()


def store_number(number):
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('INSERT INTO public.numbers(number) VALUES (%s)', (number,))
  conn.commit()


def remove_number(number):
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('DELETE FROM public.numbers WHERE number = %s', (number,))
  conn.commit()


def store_born_message(message):
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('INSERT INTO public.born(message) VALUES (%s)', (message,))
  conn.commit()


def get_born_message():
  conn = psycopg2.connect(DATABASE_URL)
  cursor = conn.cursor()
  cursor.execute('SELECT message FROM public.born')
  return cursor.fetchone()