test: include sqlalchemy 1.4 test in python client testing (#14491)

integration_tests/client-library/python/client.py

@@ -1,9 +1,10 @@
 import psycopg2
 
-class client:
-    def __init__(self, host, port,database, user, password):
+
+class Client:
+    def __init__(self, host, port, database, user, password):
         self.host = host
-        self.port=port
+        self.port = port
         self.database = database
         self.user = user
         self.password = password

integration_tests/client-library/python/crud.py

@@ -1,14 +1,12 @@
 import psycopg2
-from client import client
+from client import Client
+
+
+class SampleTableCrud:
+    # Represents the table `sample_table_py`.
 
-class crud:
     def __init__(self, host, port, database, user, password):
-        self.host = host
-        self.database = database
-        self.user = user
-        self.password = password
-        self.connection = None
-        self.port=port
+        self.client = Client(host, port, database, user, password)
 
     def create_table(self):
         create_table_query = """
@@ -19,10 +17,9 @@ def create_table(self):
             );
         """
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(create_table_query)
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("Table created successfully.")
         except psycopg2.Error as e:
             print("Table creation failed: ", str(e))
@@ -33,10 +30,9 @@ def insert_data(self, name, age, salary):
             VALUES (%s, %s,%s);
         """
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(insert_data_query, (name, age, salary))
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("Data inserted successfully.")
         except psycopg2.Error as e:
             print("Data insertion failed: ", str(e))
@@ -48,10 +44,9 @@ def update_data(self, name, salary):
             WHERE name=%s;
         """
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(update_data_query, (salary, name))
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("Data updated successfully.")
         except psycopg2.Error as e:
             print("Data updation failed: ", str(e))
@@ -61,10 +56,9 @@ def delete_data(self, name):
             DELETE FROM sample_table_py WHERE name='%s';
         """
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(insert_data_query, (name,))
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("Data deletion successfully.")
         except psycopg2.Error as e:
             print("Data deletion failed: ", str(e))
@@ -74,17 +68,9 @@ def table_drop(self):
             DROP TABLE sample_table_py;
         """
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(reset_query)
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("Table Dropped successfully")
         except psycopg2.Error as e:
             print("Table Drop Failed: ", str(e))
-
-crud_ins=crud(host="risingwave-standalone",
-        port="4566",
-        database="dev",
-        user="root",
-        password="")
-crud_ins.create_table()

integration_tests/client-library/python/materializeview.py

@@ -1,46 +1,41 @@
 import psycopg2
-from client import client
-from crud import crud
+from client import Client
+from crud import SampleTableCrud
 
+
+# Represents the materialized view `average_salary_view_py`.
 class MaterializeView:
     def __init__(self, host, port, database, user, password):
-        self.host = host
-        self.database = database
-        self.user = user
-        self.password = password
-        self.connection = None
-        self.port=port
+        self.client = Client(host, port, database, user, password)
+        self.crud = SampleTableCrud(host, port, database, user, password)
 
     def create_mv(self):
-        crud_ins = crud(self.host, self.port, self.database, self.user, self.password)
-        crud_ins.create_table()
-        crud_ins.insert_data("John",25,10000)
-        crud_ins.insert_data("Shaun",25,11000)
-        crud_ins.insert_data("Caul",25,14000)
-        crud_ins.insert_data("Mantis",28,18000)
-        crud_ins.insert_data("Tony",28,19000)
-        mv_query="""
+        self.crud.create_table()
+        self.crud.insert_data("John", 25, 10000)
+        self.crud.insert_data("Shaun", 25, 11000)
+        self.crud.insert_data("Caul", 25, 14000)
+        self.crud.insert_data("Mantis", 28, 18000)
+        self.crud.insert_data("Tony", 28, 19000)
+        mv_query = """
                     CREATE MATERIALIZED VIEW average_salary_view_py AS
                     SELECT age, AVG(salary) AS average_salary
                     FROM sample_table_py
                     GROUP BY age;
                 """
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(mv_query)
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("MV created successfully.")
         except psycopg2.Error as e:
             print("MV creation failed: ", str(e))
 
     def drop_mv(self):
         mv_drop_query = "DROP materialized view average_salary_view_py;"
         try:
-            databaseconnection = client(self.host, self.port,self.database, self.user, self.password)
-            cursor=databaseconnection.connect()
+            cursor = self.client.connect()
             cursor.execute(mv_drop_query)
-            databaseconnection.connection.commit()
+            self.client.connection.commit()
             print("MV dropped successfully.")
         except psycopg2.Error as e:
             print("MV drop failed: ", str(e))

integration_tests/client-library/python/requirements.txt

@@ -1,2 +1,4 @@
 psycopg2-binary
-pytest
+pytest
+sqlalchemy-risingwave
+SQLAlchemy==1.4.51

integration_tests/client-library/python/test_database.py

@@ -1,12 +1,12 @@
 import pytest
-from client import client
-from crud import crud
+from client import Client
+from crud import SampleTableCrud
 from materializeview import MaterializeView
 
 
 @pytest.fixture
 def db_connection():
-    db = client(
+    db = Client(
         host="risingwave-standalone",
         port="4566",
         database="dev",
@@ -19,7 +19,7 @@ def db_connection():
 
 @pytest.fixture
 def crud_instance():
-    return crud(
+    return SampleTableCrud(
         host="risingwave-standalone",
         port="4566",
         database="dev",
@@ -48,19 +48,22 @@ def test_disconnect(db_connection):
     db_connection.disconnect()
     assert db_connection.connection is None
 
+
 def test_table_creation(crud_instance, db_connection):
     cursor = db_connection.connect()
     cursor.execute("SET TRANSACTION READ WRITE;")
     crud_instance.create_table()
 
     cursor.execute("FLUSH;")
-    cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_name = 'sample_table_py';")
+    cursor.execute(
+        "SELECT table_name FROM information_schema.tables WHERE table_name = 'sample_table_py';")
     result = cursor.fetchone()[0]
     cursor.close()
     assert result == 'sample_table_py'
 
+
 def test_data_insertion(crud_instance, db_connection):
-    crud_instance.insert_data("John Doe", 25,10000)
+    crud_instance.insert_data("John Doe", 25, 10000)
 
     cursor = db_connection.connect()
     cursor.execute("FLUSH;")
@@ -71,6 +74,7 @@ def test_data_insertion(crud_instance, db_connection):
 
     assert result == 1
 
+
 def test_data_updation(crud_instance, db_connection):
     crud_instance.update_data("John Doe", 12000)
 
@@ -82,56 +86,65 @@ def test_data_updation(crud_instance, db_connection):
     cursor.close()
     assert result == 12000
 
+
 def test_data_deletion(crud_instance, db_connection):
     crud_instance.delete_data("John Doe")
 
     cursor = db_connection.connect()
     cursor.execute("FLUSH;")
-    cursor.execute("SELECT EXISTS (SELECT 1 FROM sample_table_py WHERE name = 'John Doe');")
+    cursor.execute(
+        "SELECT EXISTS (SELECT 1 FROM sample_table_py WHERE name = 'John Doe');")
     result = cursor.fetchone()
     result = result[0]
     cursor.close()
 
     assert result == True
 
+
 def test_table_drop(crud_instance, db_connection):
     crud_instance.table_drop()
 
     cursor = db_connection.connect()
     cursor.execute("FLUSH;")
-    cursor.execute("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'sample_table_py');")
+    cursor.execute(
+        "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'sample_table_py');")
     result = cursor.fetchone()
     result = result[0]
     cursor.close()
 
     assert result is False
 
-def test_mv_creation(mv_instance,db_connection):
+
+def test_mv_creation(mv_instance, db_connection):
     mv_instance.create_mv()
     cursor = db_connection.connect()
     cursor.execute("FLUSH;")
-    cursor.execute("SELECT EXISTS (SELECT 1 FROM pg_matviews WHERE matviewname = 'average_salary_view_py');")
+    cursor.execute(
+        "SELECT EXISTS (SELECT 1 FROM pg_matviews WHERE matviewname = 'average_salary_view_py');")
     result = cursor.fetchone()[0]
     cursor.close()
     assert result is True
 
-def test_mv_updation(db_connection,crud_instance):
+
+def test_mv_updation(db_connection, crud_instance):
     crud_instance.insert_data("Stark", 25, 13000)
     cursor = db_connection.connect()
     cursor.execute("FLUSH;")
-    cursor.execute("SELECT average_salary FROM average_salary_view_py WHERE age=25;")
+    cursor.execute(
+        "SELECT average_salary FROM average_salary_view_py WHERE age=25;")
     result = cursor.fetchone()[0]
     cursor.close()
     # assert result == 11250
     assert result == 12000
 
 
-def test_mv_drop(crud_instance,mv_instance,db_connection):
+def test_mv_drop(crud_instance, mv_instance, db_connection):
     mv_instance.drop_mv()
     crud_instance.table_drop()
     cursor = db_connection.connect()
     cursor.execute("FLUSH;")
-    cursor.execute("SELECT EXISTS (SELECT 1 FROM pg_matviews WHERE matviewname = 'average_salary_view_py');")
+    cursor.execute(
+        "SELECT EXISTS (SELECT 1 FROM pg_matviews WHERE matviewname = 'average_salary_view_py');")
     result = cursor.fetchone()
     result = result[0]
     cursor.close()

integration_tests/client-library/python/test_sqlalchemy.py

@@ -0,0 +1,62 @@
+from sqlalchemy import Column, BigInteger, Integer, String, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+import pytest
+# Create a base class for declarative class definitions
+Base = declarative_base()
+
+
+class User(Base):
+    # Define a simple User class as an example
+    __tablename__ = 'users'
+
+    id = Column('id', BigInteger, primary_key=True)
+    name = Column('name', String)
+    age = Column('age', Integer)
+
+
+# Pytest fixture to create and destroy the database session
+@pytest.fixture
+def db_session():
+    DB_URI = 'risingwave+psycopg2://root@risingwave-standalone:4566/dev'
+    # Create an SQLAlchemy engine to manage connections to the database
+    engine = create_engine(DB_URI)
+
+    # The automatically created table is incorrect. The BigInteger will be translated into BIGSERIAL somehow, which is not supported.
+    create_table = """
+        CREATE TABLE IF NOT EXISTS users (
+            id BIGINT PRIMARY KEY,
+            name VARCHAR,
+            age INTEGER
+        )
+    """
+    with engine.connect() as conn:
+        conn.execute(create_table)
+        conn.execute('SET RW_IMPLICIT_FLUSH=true')
+
+    Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+    session = Session()
+    yield session
+    session.close()
+
+    Base.metadata.drop_all(engine)
+
+
+# Pytest test functions to perform CRUD operations
+def test_create_user(db_session):
+    new_user = User(id=1, name='John Doe', age=30)
+    db_session.add(new_user)
+    db_session.commit()
+    assert new_user.id is not None
+
+    all_users = db_session.query(User).all()
+    assert len(all_users) > 0
+
+
+def test_delete_user(db_session):
+    user_to_delete = db_session.query(User).filter_by(name='John Doe').first()
+    if user_to_delete:
+        db_session.delete(user_to_delete)
+        db_session.commit()
+        deleted_user = db_session.query(User).get(user_to_delete.id)
+        assert deleted_user is None