Fir 8531 allow passing engine name rathe

src/firebolt/db/__init__.py

@@ -13,7 +13,7 @@
     Timestamp,
     TimestampFromTicks,
 )
-from firebolt.db.connection import Connection
+from firebolt.db.connection import Connection, connect
 from firebolt.db.cursor import Cursor
 
 apilevel = "2.0"

src/firebolt/db/connection.py

@@ -2,31 +2,97 @@
 
 from inspect import cleandoc
 from types import TracebackType
-from typing import List
+from typing import List, Optional
 
 from httpx import Timeout
 from readerwriterlock.rwlock import RWLockWrite
 
 from firebolt.client import DEFAULT_API_URL, Client
-from firebolt.common.exception import ConnectionClosedError
+from firebolt.common.exception import ConnectionClosedError, InterfaceError
+from firebolt.common.settings import Settings
 from firebolt.db.cursor import Cursor
+from firebolt.service.manager import ResourceManager
 
 DEFAULT_TIMEOUT_SECONDS: int = 5
 
 
+def connect(
+    database: str = None,
+    username: str = None,
+    password: str = None,
+    engine_name: Optional[str] = None,
+    engine_url: Optional[str] = None,
+    api_endpoint: str = DEFAULT_API_URL,
+) -> Connection:
+    cleandoc(
+        """
+        Connect to Firebolt database.
+
+        Connection parameters:
+        database - name of the database to connect
+        username - user name to use for authentication
+        password - password to use for authentication
+        engine_name - name of the engine to connect to
+        engine_url - engine endpoint to use
+        note: either engine_name or engine_url should be provided, but not both
+        """
+    )
+    if engine_name and engine_url:
+        raise InterfaceError(
+            "Both engine_name and engine_url are provided. Provide only one to connect."
+        )
+    if not engine_name and not engine_url:
+        raise InterfaceError(
+            "Neither engine_name nor engine_url are provided. Provide one to connect."
+        )
+    # This parameters are optional in function signature, but are required to connect.
+    # It's recomended to make them kwargs by PEP 249
+    for param, name in (
+        (database, "database"),
+        (username, "username"),
+        (password, "password"),
+    ):
+        if not param:
+            raise InterfaceError(f"{name} is required to connect.")
+
+    if engine_name is not None:
+        rm = ResourceManager(
+            Settings(
+                user=username, password=password, server=api_endpoint, default_region=""
+            )
+        )
+        endpoint = rm.engines.get_by_name(engine_name).endpoint
+        if endpoint is None:
+            raise InterfaceError("unable to retrieve engine endpoint.")
+        else:
+            engine_url = endpoint
+
+    # Mypy checks, this should never happen
+    assert engine_url is not None
+    assert database is not None
+    assert username is not None
+    assert password is not None
+
+    engine_url = (
+        engine_url if engine_url.startswith("http") else f"https://{engine_url}"
+    )
+    return Connection(engine_url, database, username, password, api_endpoint)
+
+
 class Connection:
     cleandoc(
         """
         Firebolt database connection class. Implements PEP-249.
 
         Parameters:
+            engine_url - Firebolt database engine REST API url
+            database - Firebolt database name
             username - Firebolt account username
             password - Firebolt account password
-            engine_url - Firebolt database engine REST API url
             api_endpoint(optional) - Firebolt API endpoint. Used for authentication
 
         Methods:
-            cursor - created new Cursor object
+            cursor - create new Cursor object
             close - close the Connection and all it's cursors
 
         Firebolt currenly doesn't support transactions so commit and rollback methods
@@ -43,9 +109,6 @@ def __init__(
         password: str,
         api_endpoint: str = DEFAULT_API_URL,
     ):
-        engine_url = (
-            engine_url if engine_url.startswith("http") else f"https://{engine_url}"
-        )
         self._client = Client(
             auth=(username, password),
             base_url=engine_url,

src/firebolt/db/examples.ipynb

@@ -15,7 +15,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from firebolt.db import Connection\n",
+    "from firebolt.db import connect\n",
     "from firebolt.client import DEFAULT_API_URL"
    ]
   },
@@ -34,8 +34,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "database_name = \"\"\n",
+    "# Only one of these two parameters should be specified\n",
     "engine_url = \"\"\n",
+    "engine_name = \"\"\n",
+    "assert bool(engine_url) != bool(\n",
+    "    engine_name\n",
+    "), \"Specify only one of engine_name and engine_url\"\n",
+    "\n",
+    "database_name = \"\"\n",
     "username = \"\"\n",
     "password = \"\"\n",
     "api_endpoint = DEFAULT_API_URL"
@@ -57,8 +63,13 @@
    "outputs": [],
    "source": [
     "# create a connection based on provided credentials\n",
-    "connection = Connection(\n",
-    "    engine_url, database_name, username, password, api_endpoint=api_endpoint\n",
+    "connection = connect(\n",
+    "    engine_url=engine_url,\n",
+    "    engine_name=engine_name,\n",
+    "    database=database_name,\n",
+    "    username=username,\n",
+    "    password=password,\n",
+    "    api_endpoint=api_endpoint,\n",
     ")\n",
     "\n",
     "# create a cursor for connection\n",

tests/integration/dbapi/conftest.py

@@ -5,13 +5,14 @@
 
 from pytest import fixture
 
-from firebolt.db import ARRAY, Connection
+from firebolt.db import ARRAY, Connection, connect
 from firebolt.db._types import ColType
 from firebolt.db.cursor import Column
 
 LOGGER = getLogger(__name__)
 
 ENGINE_URL_ENV = "ENGINE_URL"
+ENGINE_NAME_ENV = "ENGINE_NAME"
 DATABASE_NAME_ENV = "DATABASE_NAME"
 USERNAME_ENV = "USERNAME"
 PASSWORD_ENV = "PASSWORD"
@@ -29,6 +30,11 @@ def engine_url() -> str:
     return must_env(ENGINE_URL_ENV)
 
 
+@fixture(scope="session")
+def engine_name() -> str:
+    return must_env(ENGINE_NAME_ENV)
+
+
 @fixture(scope="session")
 def database_name() -> str:
     return must_env(DATABASE_NAME_ENV)
@@ -53,8 +59,29 @@ def api_endpoint() -> str:
 def connection(
     engine_url: str, database_name: str, username: str, password: str, api_endpoint: str
 ) -> Connection:
-    return Connection(
-        engine_url, database_name, username, password, api_endpoint=api_endpoint
+    return connect(
+        engine_url=engine_url,
+        database=database_name,
+        username=username,
+        password=password,
+        api_endpoint=api_endpoint,
+    )
+
+
+@fixture
+def connection_engine_name(
+    engine_name: str,
+    database_name: str,
+    username: str,
+    password: str,
+    api_endpoint: str,
+) -> Connection:
+    return connect(
+        engine_name=engine_name,
+        database=database_name,
+        username=username,
+        password=password,
+        api_endpoint=api_endpoint,
     )
 
 

tests/integration/dbapi/test_errors.py

@@ -6,15 +6,19 @@
     OperationalError,
     ProgrammingError,
 )
-from firebolt.db import Connection
+from firebolt.db import Connection, connect
 
 
 def test_invalid_credentials(
     engine_url: str, database_name: str, username: str, password: str, api_endpoint: str
 ) -> None:
     """Connection properly reacts to invalid credentials error"""
-    connection = Connection(
-        engine_url, database_name, username + "_", password + "_", api_endpoint
+    connection = connect(
+        engine_url=engine_url,
+        database=database_name,
+        username=username + "_",
+        password=password + "_",
+        api_endpoint=api_endpoint,
     )
     with raises(AuthenticationError) as exc_info:
         connection.cursor().execute("show tables")
@@ -24,12 +28,35 @@ def test_invalid_credentials(
     ), "Invalid authentication error message"
 
 
-def test_engine_not_exists(
+def test_engine_url_not_exists(
     engine_url: str, database_name: str, username: str, password: str, api_endpoint: str
 ) -> None:
     """Connection properly reacts to invalid engine url error"""
-    connection = Connection(
-        engine_url + "_", database_name, username, password, api_endpoint
+    connection = connect(
+        engine_url=engine_url + "_",
+        database=database_name,
+        username=username,
+        password=password,
+        api_endpoint=api_endpoint,
+    )
+    with raises(ConnectError):
+        connection.cursor().execute("show tables")
+
+
+def test_engine_name_not_exists(
+    engine_name: str,
+    database_name: str,
+    username: str,
+    password: str,
+    api_endpoint: str,
+) -> None:
+    """Connection properly reacts to invalid engine name error"""
+    connection = connect(
+        engine_url=engine_name + "_________",
+        database=database_name,
+        username=username,
+        password=password,
+        api_endpoint=api_endpoint,
     )
     with raises(ConnectError):
         connection.cursor().execute("show tables")
@@ -40,7 +67,13 @@ def test_database_not_exists(
 ) -> None:
     """Connection properly reacts to invalid database error"""
     new_db_name = database_name + "_"
-    connection = Connection(engine_url, new_db_name, username, password, api_endpoint)
+    connection = connect(
+        engine_url=engine_url,
+        database=new_db_name,
+        username=username,
+        password=password,
+        api_endpoint=api_endpoint,
+    )
     with raises(ProgrammingError) as exc_info:
         connection.cursor().execute("show tables")
 
@@ -50,6 +83,7 @@ def test_database_not_exists(
 
 
 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 ]")

tests/integration/dbapi/test_queries.py

@@ -14,6 +14,21 @@ def assert_deep_eq(got: Any, expected: Any, msg: str) -> bool:
     ), f"{msg}: {got}(got) != {expected}(expected)"
 
 
+def test_connect_engine_name(
+    connection_engine_name: Connection,
+    all_types_query: str,
+    all_types_query_description: List[Column],
+    all_types_query_response: List[ColType],
+) -> None:
+    """Connecting with engine name is handled properly."""
+    test_select(
+        connection_engine_name,
+        all_types_query,
+        all_types_query_description,
+        all_types_query_response,
+    )
+
+
 def test_select(
     connection: Connection,
     all_types_query: str,