peers-identity/bin/rd_some_class_insert.py

#!/usr/bin/env python3

import datetime
from pathlib import Path
from sqlalchemy import Column, Integer, String
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import scoped_session, sessionmaker

Base = declarative_base()


class SomeClass(Base):
    __tablename__ = "some_table"
    id = Column(Integer, primary_key=True)
    name = Column(String(50))


def get_connection_string():
    return f"sqlite:///{Path().home()}/.peers_identity/rd_some_class.db"


def insert_record():
    connection_string = get_connection_string()
    engine = create_engine(connection_string)
    # connection = engine.connect()
    session = scoped_session(
        sessionmaker(autocommit=False, autoflush=False, bind=engine)
    )
    Base.metadata.create_all(engine)

    foo = SomeClass(name=f"foo_{datetime.datetime.now()}")
    session.add(foo)
    session.commit()

    result = session.query(SomeClass).get({"id": "1"})
    print(result.name)

    some_class = session.query(SomeClass).all()

    for foo in some_class:
        print(foo.name)

    session.close()


def drop_table():
    connection_string = get_connection_string()
    engine = create_engine(connection_string)
    # connection = engine.connect()
    session = scoped_session(
        sessionmaker(autocommit=False, autoflush=False, bind=engine)
    )
    # Base.metadata.drop_all(bind=engine, tables=[SomeClass.__table__])
    Base.metadata.drop_all(bind=engine)
    session.close()


if __name__ == "__main__":
    from argparse import ArgumentParser
    from pathlib import Path

    argument_parser = ArgumentParser(description="Test the database connection.")
    argument_parser.add_argument(
        "-d", "--drop", help="Drop the test table", required=False
    )
    args = argument_parser.parse_args()

    if args.drop == "True":
        drop_table()
    else:
        insert_record()