From eaacc693bd93afb908d7a1b1142f37713fce6acd Mon Sep 17 00:00:00 2001
From: Pere Miquel Brull <peremiquelbrull@gmail.com>
Date: Thu, 14 Dec 2023 15:46:58 +0100
Subject: [PATCH] #12027 - Add support for Python 3.11 (#14385)

* Fix datamodel codegen and bump versions

* Add 3.11 tests

* Update hive

* pandas

* pandas
---
 .github/workflows/py-tests.yml  |  2 +-
 ingestion/setup.py              | 13 ++++++++-----
 scripts/datamodel_generation.py |  3 ++-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/py-tests.yml b/.github/workflows/py-tests.yml
index b87ea842a83..255464b9707 100644
--- a/.github/workflows/py-tests.yml
+++ b/.github/workflows/py-tests.yml
@@ -36,7 +36,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        py-version: ['3.8', '3.9', '3.10']
+        py-version: ['3.8', '3.9', '3.10', '3.11']
     steps:
     - name: Free Disk Space (Ubuntu)
       uses: jlumbroso/free-disk-space@main
diff --git a/ingestion/setup.py b/ingestion/setup.py
index 0bd61e74a6c..efc5bb14109 100644
--- a/ingestion/setup.py
+++ b/ingestion/setup.py
@@ -28,7 +28,7 @@ VERSIONS = {
     "grpc-tools": "grpcio-tools>=1.47.2",
     "msal": "msal~=1.2",
     "neo4j": "neo4j~=5.3.0",
-    "pandas": "pandas==1.3.5",
+    "pandas": "pandas<=2,<3",
     "pyarrow": "pyarrow~=10.0",
     "pydomo": "pydomo~=0.3",
     "pymysql": "pymysql>=1.0.2",
@@ -45,7 +45,7 @@ VERSIONS = {
     "looker-sdk": "looker-sdk>=22.20.0",
     "lkml": "lkml~=1.3",
     "tableau": "tableau-api-lib~=0.1",
-    "pyhive": "pyhive~=0.7",
+    "pyhive": "pyhive[hive_pure_sasl]~=0.7",
     "mongo": "pymongo~=4.3",
     "redshift": "sqlalchemy-redshift==0.8.12",
     "snowflake": "snowflake-sqlalchemy~=1.4",
@@ -58,7 +58,9 @@ COMMONS = {
         VERSIONS["boto3"],
         VERSIONS["pandas"],
         VERSIONS["pyarrow"],
-        "python-snappy~=0.6.1",
+        # python-snappy does not work well on 3.11 https://github.com/aio-libs/aiokafka/discussions/931
+        # Using this as an alternative
+        "cramjam~=2.7",
     },
     "hive": {
         "presto-types-parser>=0.0.2",
@@ -189,7 +191,8 @@ plugins: Dict[str, Set[str]] = {
     "hive": {
         *COMMONS["hive"],
         "thrift>=0.13,<1",
-        "sasl~=0.3",
+        # Replacing sasl with pure-sasl based on https://github.com/cloudera/python-sasl/issues/30 for py 3.11
+        "pure-sasl",
         "thrift-sasl~=0.4",
         "impyla~=0.18.0",
     },
@@ -197,7 +200,7 @@ plugins: Dict[str, Set[str]] = {
         "presto-types-parser>=0.0.2",
         "impyla[kerberos]~=0.18.0",
         "thrift>=0.13,<1",
-        "sasl~=0.3",
+        "pure-sasl",
         "thrift-sasl~=0.4",
     },
     "kafka": {*COMMONS["kafka"]},
diff --git a/scripts/datamodel_generation.py b/scripts/datamodel_generation.py
index f97315036ab..6b26c5545c2 100644
--- a/scripts/datamodel_generation.py
+++ b/scripts/datamodel_generation.py
@@ -46,6 +46,7 @@ main(args)
 for file_path in UNICODE_REGEX_REPLACEMENT_FILE_PATHS:
     with open(file_path, "r", encoding="UTF-8") as file_:
         content = file_.read()
-        content = content.replace("(?U)", "(?u)")
+        # Python now requires to move the global flags at the very start of the expression
+        content = content.replace("^(?U)", "(?u)^")
     with open(file_path, "w", encoding="UTF-8") as file_:
         file_.write(content)