ydb sqlalchemy experiment

.gitignore

@@ -4,4 +4,8 @@ ydb.egg-info/
 /tox
 /venv
 /ydb_certs
+/ydb_data
 /tmp
+.coverage
+/cov_html
+/build

examples/_sqlalchemy_example/example.py

@@ -0,0 +1,229 @@
+import datetime
+import logging
+import argparse
+import sqlalchemy as sa
+from sqlalchemy import orm, exc, sql
+from sqlalchemy import Table, Column, Integer, String, Float, TIMESTAMP
+from ydb._sqlalchemy import register_dialect
+
+from fill_tables import fill_all_tables, to_days
+from models import Base, Series, Episodes
+
+
+def describe_table(engine, name):
+    inspect = sa.inspect(engine)
+    print(f"describe table {name}:")
+    for col in inspect.get_columns(name):
+        print(f"\t{col['name']}: {col['type']}")
+
+
+def simple_select(conn):
+    stm = sa.select(Series).where(Series.series_id == 1)
+    res = conn.execute(stm)
+    print(res.one())
+
+
+def simple_insert(conn):
+    stm = Episodes.__table__.insert().values(
+        series_id=3, season_id=6, episode_id=1, title="TBD"
+    )
+    conn.execute(stm)
+
+
+def test_types(conn):
+    types_tb = Table(
+        "test_types",
+        Base.metadata,
+        Column("id", Integer, primary_key=True),
+        Column("str", String),
+        Column("num", Float),
+        Column("dt", TIMESTAMP),
+    )
+    types_tb.drop(bind=conn.engine, checkfirst=True)
+    types_tb.create(bind=conn.engine, checkfirst=True)
+
+    stm = types_tb.insert().values(
+        id=1,
+        str=b"Hello World!",
+        num=3.1415,
+        dt=datetime.datetime.now(),
+    )
+    conn.execute(stm)
+
+    # GROUP BY
+    stm = sa.select(types_tb.c.str, sa.func.max(types_tb.c.num)).group_by(
+        types_tb.c.str
+    )
+    rs = conn.execute(stm)
+    for x in rs:
+        print(x)
+
+
+def run_example_orm(engine):
+    Base.metadata.bind = engine
+    Base.metadata.drop_all()
+    Base.metadata.create_all()
+
+    session = orm.sessionmaker(bind=engine)()
+
+    rs = session.query(Episodes).all()
+    for e in rs:
+        print(f"{e.episode_id}: {e.title}")
+
+    fill_all_tables(session.connection())
+
+    try:
+        session.add_all(
+            [
+                Episodes(
+                    series_id=2,
+                    season_id=1,
+                    episode_id=1,
+                    title="Minimum Viable Product",
+                    air_date=to_days("2014-04-06"),
+                ),
+                Episodes(
+                    series_id=2,
+                    season_id=1,
+                    episode_id=2,
+                    title="The Cap Table",
+                    air_date=to_days("2014-04-13"),
+                ),
+                Episodes(
+                    series_id=2,
+                    season_id=1,
+                    episode_id=3,
+                    title="Articles of Incorporation",
+                    air_date=to_days("2014-04-20"),
+                ),
+                Episodes(
+                    series_id=2,
+                    season_id=1,
+                    episode_id=4,
+                    title="Fiduciary Duties",
+                    air_date=to_days("2014-04-27"),
+                ),
+                Episodes(
+                    series_id=2,
+                    season_id=1,
+                    episode_id=5,
+                    title="Signaling Risk",
+                    air_date=to_days("2014-05-04"),
+                ),
+            ]
+        )
+        session.commit()
+    except exc.DatabaseError:
+        print("Episodes already added!")
+        session.rollback()
+
+    rs = session.query(Episodes).all()
+    for e in rs:
+        print(f"{e.episode_id}: {e.title}")
+
+    rs = session.query(Episodes).filter(Episodes.title == "abc??").all()
+    for e in rs:
+        print(e.title)
+
+    print("Episodes count:", session.query(Episodes).count())
+
+    max_episode = session.query(sql.expression.func.max(Episodes.episode_id)).scalar()
+    print("Maximum episodes id:", max_episode)
+
+    session.add(
+        Episodes(
+            series_id=2,
+            season_id=1,
+            episode_id=max_episode + 1,
+            title="Signaling Risk",
+            air_date=to_days("2014-05-04"),
+        )
+    )
+
+    print("Episodes count:", session.query(Episodes).count())
+
+
+def run_example_core(engine):
+    with engine.connect() as conn:
+        # raw sql
+        rs = conn.execute("SELECT 1 AS value")
+        print(rs.fetchone()["value"])
+
+        fill_all_tables(conn)
+
+        for t in "series seasons episodes".split():
+            describe_table(engine, t)
+
+        tb = sa.Table("episodes", sa.MetaData(engine), autoload=True)
+        stm = (
+            sa.select([tb.c.title])
+            .where(sa.and_(tb.c.series_id == 1, tb.c.season_id == 3))
+            .where(tb.c.title.like("%"))
+            .order_by(sa.asc(tb.c.title))
+            # TODO: limit isn't working now
+            # .limit(3)
+        )
+        rs = conn.execute(stm)
+        print(rs.fetchall())
+
+        simple_select(conn)
+
+        simple_insert(conn)
+
+        # simple join
+        stm = sa.select(
+            [Episodes.__table__.join(Series, Episodes.series_id == Series.series_id)]
+        ).where(sa.and_(Series.series_id == 1, Episodes.season_id == 1))
+        rs = conn.execute(stm)
+        for row in rs:
+            print(f"{row.series_title}({row.episode_id}): {row.title}")
+
+        rs = conn.execute(sa.select(Episodes).where(Episodes.series_id == 3))
+        print(rs.fetchall())
+
+        # count
+        cnt = conn.execute(sa.func.count(Episodes.episode_id)).scalar()
+        print("Episodes cnt:", cnt)
+
+        # simple delete
+        conn.execute(sa.delete(Episodes).where(Episodes.title == "TBD"))
+        cnt = conn.execute(sa.func.count(Episodes.episode_id)).scalar()
+        print("Episodes cnt:", cnt)
+
+        test_types(conn)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="""\033[92mYandex.Database examples sqlalchemy usage.\x1b[0m\n""",
+    )
+    parser.add_argument(
+        "-d",
+        "--database",
+        help="Name of the database to use",
+        default="/local",
+    )
+    parser.add_argument(
+        "-e",
+        "--endpoint",
+        help="Endpoint url to use",
+        default="grpc://localhost:2136",
+    )
+
+    args = parser.parse_args()
+    register_dialect()
+    engine = sa.create_engine(
+        "yql:///ydb/",
+        connect_args={"database": args.database, "endpoint": args.endpoint},
+    )
+
+    logging.basicConfig(level=logging.INFO)
+    logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+
+    run_example_core(engine)
+    # run_example_orm(engine)
+
+
+if __name__ == "__main__":
+    main()

examples/_sqlalchemy_example/fill_tables.py

@@ -0,0 +1,83 @@
+import iso8601
+
+import sqlalchemy as sa
+from models import Base, Series, Seasons, Episodes
+
+
+def to_days(date):
+    timedelta = iso8601.parse_date(date) - iso8601.parse_date("1970-1-1")
+    return timedelta.days
+
+
+def fill_series(conn):
+    data = [
+        (
+            1,
+            "IT Crowd",
+            "The IT Crowd is a British sitcom produced by Channel 4, written by Graham Linehan, produced by "
+            "Ash Atalla and starring Chris O'Dowd, Richard Ayoade, Katherine Parkinson, and Matt Berry.",
+            to_days("2006-02-03"),
+        ),
+        (
+            2,
+            "Silicon Valley",
+            "Silicon Valley is an American comedy television series created by Mike Judge, John Altschuler and "
+            "Dave Krinsky. The series focuses on five young men who founded a startup company in Silicon Valley.",
+            to_days("2014-04-06"),
+        ),
+    ]
+    conn.execute(sa.insert(Series).values(data))
+
+
+def fill_seasons(conn):
+    data = [
+        (1, 1, "Season 1", to_days("2006-02-03"), to_days("2006-03-03")),
+        (1, 2, "Season 2", to_days("2007-08-24"), to_days("2007-09-28")),
+        (1, 3, "Season 3", to_days("2008-11-21"), to_days("2008-12-26")),
+        (1, 4, "Season 4", to_days("2010-06-25"), to_days("2010-07-30")),
+        (2, 1, "Season 1", to_days("2014-04-06"), to_days("2014-06-01")),
+        (2, 2, "Season 2", to_days("2015-04-12"), to_days("2015-06-14")),
+        (2, 3, "Season 3", to_days("2016-04-24"), to_days("2016-06-26")),
+        (2, 4, "Season 4", to_days("2017-04-23"), to_days("2017-06-25")),
+        (2, 5, "Season 5", to_days("2018-03-25"), to_days("2018-05-13")),
+    ]
+    conn.execute(sa.insert(Seasons).values(data))
+
+
+def fill_episodes(conn):
+    data = [
+        (1, 1, 1, "Yesterday's Jam", to_days("2006-02-03")),
+        (1, 1, 2, "Calamity Jen", to_days("2006-02-03")),
+        (1, 1, 3, "Fifty-Fifty", to_days("2006-02-10")),
+        (1, 1, 4, "The Red Door", to_days("2006-02-17")),
+        (1, 1, 5, "The Haunting of Bill Crouse", to_days("2006-02-24")),
+        (1, 1, 6, "Aunt Irma Visits", to_days("2006-03-03")),
+        (1, 2, 1, "The Work Outing", to_days("2006-08-24")),
+        (1, 2, 2, "Return of the Golden Child", to_days("2007-08-31")),
+        (1, 2, 3, "Moss and the German", to_days("2007-09-07")),
+        (1, 2, 4, "The Dinner Party", to_days("2007-09-14")),
+        (1, 2, 5, "Smoke and Mirrors", to_days("2007-09-21")),
+        (1, 2, 6, "Men Without Women", to_days("2007-09-28")),
+        (1, 3, 1, "From Hell", to_days("2008-11-21")),
+        (1, 3, 2, "Are We Not Men?", to_days("2008-11-28")),
+        (1, 3, 3, "Tramps Like Us", to_days("2008-12-05")),
+        (1, 3, 4, "The Speech", to_days("2008-12-12")),
+        (1, 3, 5, "Friendface", to_days("2008-12-19")),
+        (1, 3, 6, "Calendar Geeks", to_days("2008-12-26")),
+        (1, 4, 1, "Jen The Fredo", to_days("2010-06-25")),
+        (1, 4, 2, "The Final Countdown", to_days("2010-07-02")),
+        (1, 4, 3, "Something Happened", to_days("2010-07-09")),
+        (1, 4, 4, "Italian For Beginners", to_days("2010-07-16")),
+        (1, 4, 5, "Bad Boys", to_days("2010-07-23")),
+        (1, 4, 6, "Reynholm vs Reynholm", to_days("2010-07-30")),
+    ]
+    conn.execute(sa.insert(Episodes).values(data))
+
+
+def fill_all_tables(conn):
+    Base.metadata.drop_all(conn.engine)
+    Base.metadata.create_all(conn.engine)
+
+    fill_series(conn)
+    fill_seasons(conn)
+    fill_episodes(conn)

examples/_sqlalchemy_example/models.py

@@ -0,0 +1,34 @@
+import sqlalchemy.orm as orm
+from sqlalchemy import Column, Integer, Unicode
+
+
+Base = orm.declarative_base()
+
+
+class Series(Base):
+    __tablename__ = "series"
+
+    series_id = Column(Integer, primary_key=True)
+    title = Column(Unicode)
+    series_info = Column(Unicode)
+    release_date = Column(Integer)
+
+
+class Seasons(Base):
+    __tablename__ = "seasons"
+
+    series_id = Column(Integer, primary_key=True)
+    season_id = Column(Integer, primary_key=True)
+    title = Column(Unicode)
+    first_aired = Column(Integer)
+    last_aired = Column(Integer)
+
+
+class Episodes(Base):
+    __tablename__ = "episodes"
+
+    series_id = Column(Integer, primary_key=True)
+    season_id = Column(Integer, primary_key=True)
+    episode_id = Column(Integer, primary_key=True)
+    title = Column(Unicode)
+    air_date = Column(Integer)

test-requirements.txt

@@ -47,3 +47,5 @@ sqlalchemy==1.4.26
 pylint-protobuf
 cython
 freezegun==1.2.2
+grpcio-tools
+pytest-cov

tests/sqlalchemy/conftest.py

@@ -0,0 +1,22 @@
+import pytest
+import sqlalchemy as sa
+
+from ydb._sqlalchemy import register_dialect
+
+
+@pytest.fixture(scope="module")
+def engine(endpoint, database):
+    register_dialect()
+    engine = sa.create_engine(
+        "yql:///ydb/",
+        connect_args={"database": database, "endpoint": endpoint},
+    )
+
+    yield engine
+    engine.dispose()
+
+
+@pytest.fixture(scope="module")
+def connection(engine):
+    with engine.connect() as conn:
+        yield conn