feat: FIR-38073 ecosystem support for geography type for python sdk (#404)

Author: stepansergeevitch

setup.cfg

@@ -26,7 +26,7 @@ project_urls =
 packages = find:
 install_requires =
     aiorwlock==1.1.0
-    anyio>=3.7.1
+    anyio>=3.7.1,<4.5.0
     appdirs>=1.4.4
     appdirs-stubs>=0.1.0
     async-generator>=1.10
@@ -55,7 +55,7 @@ dev =
     devtools==0.7.0
     mypy==1.*,<1.10.0
     pre-commit==3.5.0
-    pyfakefs>=4.5.3
+    pyfakefs>=4.5.3,<=5.6.0
     pytest==7.2.0
     pytest-cov==3.0.0
     pytest-httpx>=0.13.0

src/firebolt/client/auth/base.py

@@ -171,5 +171,9 @@ async def async_auth_flow(
             finally:
                 # token gets updated only after flow.send is called
                 # so unlock only after that
-                if self._lock.locked() and self._lock._owner_task == get_current_task():
+                # TODO: FIR-38687 Fix support for anyio 4.5.0+
+                if (
+                    self._lock.locked()
+                    and self._lock._owner_task == get_current_task()  # type: ignore
+                ):
                     self._lock.release()

src/firebolt/common/_types.py

@@ -177,6 +177,7 @@ class _InternalType(Enum):
     Boolean = "boolean"
 
     Bytea = "bytea"
+    Geography = "geography"
 
     Nothing = "Nothing"
 
@@ -198,6 +199,7 @@ def python_type(self) -> type:
             _InternalType.TimestampTz: datetime,
             _InternalType.Boolean: bool,
             _InternalType.Bytea: bytes,
+            _InternalType.Geography: str,
             # For simplicity, this could happen only during 'select null' query
             _InternalType.Nothing: str,
         }

tests/integration/dbapi/async/V2/test_errors_async.py

@@ -5,7 +5,6 @@
 from firebolt.utils.exception import (
     AccountNotFoundOrNoAccessError,
     FireboltStructuredError,
-    OperationalError,
 )
 
 
@@ -64,7 +63,7 @@ async def test_engine_name_not_exists(
     api_endpoint: str,
 ) -> None:
     """Connection properly reacts to invalid engine name error."""
-    with raises(OperationalError):
+    with raises(FireboltStructuredError) as exc_info:
         async with await connect(
             engine_name=engine_name + "_________",
             database=database_name,
@@ -74,6 +73,10 @@ async def test_engine_name_not_exists(
         ) as connection:
             await connection.cursor().execute("show tables")
 
+    assert f"Engine '{engine_name}_________' does not exist" in str(
+        exc_info.value
+    ), "Invalid engine error message."
+
 
 async def test_database_not_exists(
     engine_name: str,
@@ -84,7 +87,7 @@ async def test_database_not_exists(
 ) -> None:
     """Connection properly reacts to invalid database error."""
     new_db_name = database_name + "_"
-    with raises(OperationalError) as exc_info:
+    with raises(FireboltStructuredError) as exc_info:
         async with await connect(
             engine_name=engine_name,
             database=new_db_name,
@@ -102,23 +105,7 @@ async def test_database_not_exists(
 async def test_sql_error(connection: Connection) -> None:
     """Connection properly reacts to SQL execution error."""
     with connection.cursor() as c:
-        with raises(OperationalError) as exc_info:
-            await c.execute("select ]")
-
-        assert str(exc_info.value).startswith(
-            "Error executing query"
-        ), "Invalid SQL error message."
-
-
-async def test_structured_error(connection_system_engine_no_db: Connection) -> None:
-    """Connection properly reacts to structured error."""
-    with connection_system_engine_no_db.cursor() as c:
-        await c.execute("SET advanced_mode=1")
-        await c.execute("SET enable_json_error_output_format=true")
-
         with raises(FireboltStructuredError) as exc_info:
-            await c.execute("select 'dummy'::int")
+            await c.execute("select ]")
 
-        assert "Cannot parse string" in str(
-            exc_info.value
-        ), "Invalid structured error message"
+        assert "syntax error" in str(exc_info.value), "Invalid SQL error message."

tests/integration/dbapi/async/V2/test_queries_async.py

@@ -409,3 +409,23 @@ async def test_connection_with_mixed_case_db_and_engine(
         await cursor.execute('CREATE TABLE "test_table" (id int)')
         # This fails if we're not running on a user engine
         await cursor.execute('INSERT INTO "test_table" VALUES (1)')
+
+
+async def test_select_geography(
+    connection: Connection,
+    select_geography_query: str,
+    select_geography_description: List[Column],
+    select_geography_response: List[ColType],
+) -> None:
+    with connection.cursor() as c:
+        await c.execute(select_geography_query)
+        assert (
+            c.description == select_geography_description
+        ), "Invalid description value"
+        res = await c.fetchall()
+        assert len(res) == 1, "Invalid data length"
+        assert_deep_eq(
+            res,
+            select_geography_response,
+            "Invalid data returned by fetchall",
+        )

tests/integration/dbapi/async/V2/test_system_engine_async.py

@@ -5,7 +5,7 @@
 
 from firebolt.async_db import Connection
 from firebolt.common._types import ColType, Column
-from firebolt.utils.exception import OperationalError
+from firebolt.utils.exception import FireboltStructuredError
 from tests.integration.dbapi.utils import assert_deep_eq
 
 system_error_pattern = re.compile(
@@ -49,14 +49,14 @@ async def test_system_engine(
 
         if connection_system_engine.init_parameters.get("database"):
             await c.execute("show tables")
-            with raises(OperationalError) as e:
+            with raises(FireboltStructuredError) as e:
                 # Either one or another query fails if we're not on a user engine
                 await c.execute('create table if not exists "test_async"(id int)')
                 await c.execute('insert into "test_async" values (1)')
             assert system_error_pattern.search(str(e.value)), "Invalid error message"
         else:
             await c.execute("show databases")
-            with raises(OperationalError):
+            with raises(FireboltStructuredError):
                 await c.execute("show tables")
 
 
@@ -89,6 +89,6 @@ async def test_system_engine_use_engine(
         await cursor.execute(f'INSERT INTO "{table_name}" VALUES (1)')
         await cursor.execute('USE ENGINE "system"')
         # Werify we've switched to system by making previous query fail
-        with raises(OperationalError):
+        with raises(FireboltStructuredError):
             await cursor.execute(f'INSERT INTO "{table_name}" VALUES (1)')
         await cursor.execute(f'DROP TABLE IF EXISTS "{table_name}"')

tests/integration/dbapi/conftest.py

@@ -194,3 +194,18 @@ def create_drop_description() -> List[Column]:
         Column("num_hosts_remaining", int, None, None, None, None, None),
         Column("num_hosts_active", int, None, None, None, None, None),
     ]
+
+
+@fixture
+def select_geography_query() -> str:
+    return "SELECT 'POINT(1 1)'::geography as \"geography\""
+
+
+@fixture
+def select_geography_description() -> List[Column]:
+    return [Column("geography", str, None, None, None, None, None)]
+
+
+@fixture
+def select_geography_response() -> List[ColType]:
+    return [["0101000020E6100000FEFFFFFFFFFFEF3F000000000000F03F"]]

tests/integration/dbapi/sync/V2/test_errors.py

@@ -1,12 +1,10 @@
-from pytest import mark, raises
+from pytest import raises
 
 from firebolt.client.auth import ClientCredentials
 from firebolt.db import Connection, connect
 from firebolt.utils.exception import (
     AccountNotFoundOrNoAccessError,
-    FireboltDatabaseError,
     FireboltStructuredError,
-    OperationalError,
 )
 
 
@@ -63,7 +61,7 @@ def test_engine_name_not_exists(
     api_endpoint: str,
 ) -> None:
     """Connection properly reacts to invalid engine name error."""
-    with raises(OperationalError):
+    with raises(FireboltStructuredError) as exc_info:
         with connect(
             account_name=account_name,
             engine_name=engine_name + "_________",
@@ -73,53 +71,39 @@ def test_engine_name_not_exists(
         ) as connection:
             connection.cursor().execute("show tables")
 
+    assert f"Engine '{engine_name}_________' does not exist" in str(
+        exc_info.value
+    ), "Invalid engine error message."
+
 
-@mark.skip(reason="Behaviour is different in prod vs dev")
 def test_database_not_exists(
-    engine_url: str,
+    engine_name: str,
     database_name: str,
     auth: ClientCredentials,
-    account_name: str,
     api_endpoint: str,
+    account_name: str,
 ) -> None:
     """Connection properly reacts to invalid database error."""
     new_db_name = database_name + "_"
-    with connect(
-        account_name=account_name,
-        engine_url=engine_url,
-        database=new_db_name,
-        auth=auth,
-        api_endpoint=api_endpoint,
-    ) as connection:
-        with raises(FireboltDatabaseError) as exc_info:
+    with raises(FireboltStructuredError) as exc_info:
+        with connect(
+            engine_name=engine_name,
+            database=new_db_name,
+            auth=auth,
+            account_name=account_name,
+            api_endpoint=api_endpoint,
+        ) as connection:
             connection.cursor().execute("show tables")
 
-        assert (
-            str(exc_info.value)
-            == f"Engine {engine_name} is attached to {database_name} instead of {new_db_name}"
-        ), "Invalid database name error message"
+    assert f"Database '{new_db_name}' does not exist or not authorized" in str(
+        exc_info.value
+    ), "Invalid database error message."
 
 
 def test_sql_error(connection: Connection) -> None:
     """Connection properly reacts to sql execution error."""
     with connection.cursor() as c:
-        with raises(OperationalError) as exc_info:
-            c.execute("select ]")
-
-        assert str(exc_info.value).startswith(
-            "Error executing query"
-        ), "Invalid SQL error message"
-
-
-def test_structured_error(connection_system_engine_no_db: Connection) -> None:
-    """Connection properly reacts to structured error."""
-    with connection_system_engine_no_db.cursor() as c:
-        c.execute("SET advanced_mode=1")
-        c.execute("SET enable_json_error_output_format=true")
-
         with raises(FireboltStructuredError) as exc_info:
-            c.execute("select 'dummy'::int")
+            c.execute("select ]")
 
-        assert "Cannot parse string" in str(
-            exc_info.value
-        ), "Invalid structured error message"
+        assert "syntax error" in str(exc_info.value), "Invalid SQL error message."

tests/integration/dbapi/sync/V2/test_queries.py

@@ -492,3 +492,23 @@ def test_connection_with_mixed_case_db_and_engine(
         cursor.execute('CREATE TABLE "test_table" (id int)')
         # This fails if we're not running on a user engine
         cursor.execute('INSERT INTO "test_table" VALUES (1)')
+
+
+def test_select_geography(
+    connection: Connection,
+    select_geography_query: str,
+    select_geography_description: List[Column],
+    select_geography_response: List[ColType],
+):
+    with connection.cursor() as c:
+        c.execute(select_geography_query)
+        assert (
+            c.description == select_geography_description
+        ), "Invalid description value"
+        res = c.fetchall()
+        assert len(res) == 1, "Invalid data length"
+        assert_deep_eq(
+            res,
+            select_geography_response,
+            "Invalid data returned by fetchall",
+        )

tests/integration/dbapi/sync/V2/test_system_engine.py

@@ -5,7 +5,7 @@
 
 from firebolt.common._types import ColType, Column
 from firebolt.db import Connection
-from firebolt.utils.exception import OperationalError
+from firebolt.utils.exception import FireboltStructuredError
 from tests.integration.dbapi.utils import assert_deep_eq
 
 system_error_pattern = re.compile(
@@ -50,14 +50,14 @@ def test_system_engine(
 
         if connection_system_engine.init_parameters.get("database"):
             c.execute("show tables")
-            with raises(OperationalError) as e:
+            with raises(FireboltStructuredError) as e:
                 # Either one or another query fails if we're not on a user engine
                 c.execute('create table if not exists "test_sync"(id int)')
                 c.execute('insert into "test_sync" values (1)')
             assert system_error_pattern.search(str(e.value)), "Invalid error message"
         else:
             c.execute("show databases")
-            with raises(OperationalError):
+            with raises(FireboltStructuredError):
                 c.execute("show tables")
 
 
@@ -90,6 +90,6 @@ def test_system_engine_use_engine(
         cursor.execute(f'INSERT INTO "{table_name}" VALUES (1)')
         cursor.execute('USE ENGINE "system"')
         # Verify we've switched to system by making previous query fail
-        with raises(OperationalError):
+        with raises(FireboltStructuredError):
             cursor.execute(f'INSERT INTO "{table_name}" VALUES (1)')
         cursor.execute(f'DROP TABLE IF EXISTS "{table_name}"')