[Feature] Add support for SQLAlchemy 2.0 (#52392)

Signed-off-by: Maciek Bryński <[email protected]>
Signed-off-by: Maciej Bryński <[email protected]>

contrib/starrocks-python-client/README.md

@@ -27,17 +27,21 @@ starrocks://<User>:<Password>@<Host>:<Port>/<Catalog>.<Database>
 ```
 
 ## Example
-It is recommended to use python 3.x to connect to the StarRocks database, eg:
+Python connector supports only Python 3 and SQLAlchemy 2:
 ```
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, Integer, insert
 from sqlalchemy.schema import Table, MetaData, Column
 from sqlalchemy.sql.expression import select, text
 
 engine = create_engine('starrocks://root:xxx@localhost:9030/hive_catalog.hive_db')
-connection = engine.connect()
 
-rows = connection.execute(text("SELECT * FROM hive_table")).fetchall()
+### Querying data
+with engine.connect() as connection:
+    rows = connection.execute(text("SELECT * FROM hive_table")).fetchall()
+    print(rows)
 
+
+### DDL Operation
 meta = MetaData()
 tbl = Table(
     'table1',
@@ -47,11 +51,17 @@ tbl = Table(
     starrocks_comment='table comment',
     starrocks_properties=(
         ("storage_medium", "SSD"),
-        ("storage_cooldown_time", "2015-06-04 00:00:00"),
+        ("storage_cooldown_time", "2025-06-04 00:00:00"),
+        ("replication_num", "1")
     ))
 
-meta.createall()
-with connection.begin() as con:
-    tbl.insert().values(id=1)
-rows = connection.execute(tbl.select()).fetchall()
+meta.create_all(engine)
+
+### Insert data
+stmt = insert(tbl).values(id=1)
+stmt.compile()
+with engine.connect() as connection:
+    connection.execute(stmt)
+    rows = connection.execute(tbl.select()).fetchall()
+    print(rows)
 ```

contrib/starrocks-python-client/setup.py

@@ -54,8 +54,8 @@
         "Topic :: Database :: Front-Ends",
     ],
     install_requires=[
-        "sqlalchemy>=1.4, <2",
-        "sqlalchemy-utils>=0.38.3, <0.39",
+        "sqlalchemy>=2.0",
+        "sqlalchemy-utils>=0.41.2",
         "pymysql>=1.1.0",
     ],
     setup_requires=["pytest-runner"],

contrib/starrocks-python-client/starrocks/__init__.py

@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "1.0.6"
+__version__ = "1.1.0"

contrib/starrocks-python-client/starrocks/dialect.py

@@ -387,7 +387,7 @@ def _get_server_version_info(self, connection):
         cursor.execute("SELECT CURRENT_VERSION()")
         val = cursor.fetchone()[0]
         cursor.close()
-        if util.py3k and isinstance(val, bytes):
+        if isinstance(val, bytes):
             val = val.decode()
 
         return self._parse_server_version(val)
@@ -525,3 +525,11 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
 
             indexes.append(index_d)
         return indexes
+
+    def has_table(self, connection, table_name, schema=None, **kw):
+        try:
+            return super().has_table(connection, table_name, schema, **kw)
+        except exc.DBAPIError as e:
+            if self._extract_error_code(e.orig) in (5501, 5502):
+                return False
+            raise