feat: Adds ability to connect to datashare databases for clusters and serverless workgroups running the PREVIEW_2023 track

Author: Brooke-white

redshift_connector/core.py

@@ -853,9 +853,44 @@ def _is_multi_databases_catalog_enable_in_server(self: "Connection") -> bool:
             # if we don't receive this param from the server, we do not support
             return False
 
+    @property
+    def _is_cross_datasharing_enable_in_server(self: "Connection") -> bool:
+        """
+        Returns True if cross datasharing is enabled in the server. Returns False if disabled or not received.
+        :return:
+        :rtype:
+        """
+        cross_datasharing_enable_in_server: bool = False
+
+        for parameter in self.parameter_statuses:
+            if parameter[0] == b"external_database":
+                if parameter[1] == b"on":
+                    cross_datasharing_enable_in_server = True
+                elif parameter[1] == b"off":
+                    cross_datasharing_enable_in_server = False
+                else:
+                    raise InterfaceError(
+                        "Protocol error. Session setup failed. Invalid value of external_database parameter. Only on/off are valid values"
+                    )
+                break
+        return cross_datasharing_enable_in_server
+
     @property
     def is_single_database_metadata(self):
-        return self._database_metadata_current_db_only or not self._is_multi_databases_catalog_enable_in_server
+        """
+        Returns True if single database metadata enabled using ``database_metadata_current_db_only`` connection
+        parameter or if server has neither multi-database catalog nor cross datashare enabled.
+        Returns False if ``database_metadata_current_db_only`` connection parameter is disabled or if server
+        has neither multi-database catalog nor cross datashare enabled.
+        :return:
+        :rtype:
+        """
+        # for cross datasharing we always return False
+        if self._is_cross_datasharing_enable_in_server:
+            return False
+
+        else:
+            return self._database_metadata_current_db_only or not self._is_multi_databases_catalog_enable_in_server
 
     def handle_ERROR_RESPONSE(self: "Connection", data, ps):
         """

redshift_connector/cursor.py

@@ -542,7 +542,7 @@ def __is_valid_table(self: "Cursor", table: str) -> bool:
         if len(split_table_name) > 2:
             return False
 
-        q: str = "select 1 from information_schema.tables where table_name = ?"
+        q: str = "select 1 from pg_catalog.svv_all_tables where table_name = ?"
 
         temp = self.paramstyle
         self.paramstyle = DbApiParamstyle.QMARK.value

test/conftest.py

@@ -130,6 +130,30 @@ def serverless_cname_db_kwargs() -> typing.Dict[str, typing.Union[str, bool]]:
 
     return db_connect
 
+@pytest.fixture(scope="class")
+def ds_consumer_db_kwargs() -> typing.Dict[str, str]:
+    db_connect = {
+        "database": conf.get("redshift-ds-consumer", "database", fallback="mock_database"),
+        "host": conf.get("redshift-ds-consumer", "host", fallback="cname.mytest.com"),
+        "user": conf.get("redshift-ds-consumer", "user", fallback="mock_user"),
+        "password": conf.get("redshift-ds-consumer", "password", fallback="mock_password"),
+        "extra": conf.get("redshift-ds-consumer", "extra", fallback="mock_extra"),
+    }
+
+    return db_connect
+
+
+@pytest.fixture(scope="class")
+def ds_producer_db_kwargs() -> typing.Dict[str, str]:
+    db_connect = {
+        "database": conf.get("redshift-ds-producer", "database", fallback="mock_database"),
+        "host": conf.get("redshift-ds-producer", "host", fallback="cname.mytest.com"),
+        "user": conf.get("redshift-ds-producer", "user", fallback="mock_user"),
+        "password": conf.get("redshift-ds-producer", "password", fallback="mock_password"),
+    }
+
+    return db_connect
+
 
 @pytest.fixture(scope="class")
 def okta_idp() -> typing.Dict[str, typing.Union[str, bool, int]]:

test/integration/metadata/test_list_catalog.py

@@ -6,7 +6,17 @@
 
 import redshift_connector
 
-# from warnings import filterwarnings
+# Pre-reqs
+# These tests require the following:
+# - A single Redshift instance OR two Redshift instances configured for datasharing as consumer and producer.
+# - Config.ini entry for [ci-cluster] file populated with the host and port of the single Redshift instance or the
+#   producer instance if you want to test datasharing metadata methods
+#
+# How to use this test file
+# The included tests can be run to ensure the expected Redshift system tables are queried. Running this test via CLI
+# and using tee can enable to you compare and contrast result sets returned by metadata APIs on multiple clusters with
+# different patch versions. Please note the result set output is only useful if the same data is accessible by the
+# database user in all test situations.
 
 
 conf: configparser.ConfigParser = configparser.ConfigParser()
@@ -145,23 +155,27 @@ def test_get_schemas(mocker, _input, db_kwargs) -> None:
     database_metadata_current_db_only_val, _args = _input
     db_kwargs["database_metadata_current_db_only"] = database_metadata_current_db_only_val
     with redshift_connector.connect(**db_kwargs) as conn:
-        assert conn.is_single_database_metadata is database_metadata_current_db_only_val
-
         with conn.cursor() as cursor:
             spy = mocker.spy(cursor, "execute")
             result: typing.Tuple = cursor.get_schemas(**_args)
-            # ensure query was executed with arguments passed to get_schemas
+            print(result)
+            # ensure execute was called
             assert spy.called
-            assert spy.call_count == 1
-
-            if _args["schema_pattern"] is not None:  # should be in parameterized portion
-                assert _args["schema_pattern"] in spy.call_args[0][1]
+            assert spy.call_count == 1  # call in get_schemas()
 
+            # ensure execute was called with the catalog value in the prepared statement
             if _args["catalog"] is not None:
                 assert _args["catalog"] in spy.call_args[0][0]
 
-            assert len(result) > 0, print(spy.call_args, "\n", result)
-            assert len(result[0]) == 2
+            # ensure execute was called with below bind parameters
+            if _args["schema_pattern"] is not None:
+                assert _args["schema_pattern"] in spy.call_args[0][1]
+
+            # assert query text executed contains the target table name
+            if conn.is_single_database_metadata:
+                assert "FROM pg_catalog.pg_namespace" in spy.call_args[0][0]
+            else:
+                assert "FROM PG_CATALOG.SVV_ALL_SCHEMAS" in spy.call_args[0][0]
 
 
 def get_tables_test_data() -> typing.List[typing.Optional[typing.Tuple[bool, typing.Dict[str, typing.Any]]]]:
@@ -272,35 +286,54 @@ def test_get_tables(mocker, _input, db_kwargs) -> None:
     database_metadata_current_db_only_val, _args = _input
     db_kwargs["database_metadata_current_db_only"] = database_metadata_current_db_only_val
     with redshift_connector.connect(**db_kwargs) as conn:
-        assert conn.is_single_database_metadata is database_metadata_current_db_only_val
-
         with conn.cursor() as cursor:
             spy = mocker.spy(cursor, "execute")
             result: typing.Tuple = cursor.get_tables(**_args)
             print(result)
-            # ensure query was executed with arguments passed to get_schemas
+            # ensure execute was called
             assert spy.called
 
-            if _args["schema_pattern"] is not None and database_metadata_current_db_only_val:
+            if _args["schema_pattern"] is not None and conn.is_single_database_metadata:
                 assert spy.call_count == 2  # call in __schema_pattern_match(), get_tables()
             else:
-                assert spy.call_count == 1
-
-            if _args["schema_pattern"] is not None:  # should be in parameterized portion
-                print(spy.call_args)
-                print(spy.call_args[0][1])
-                print(_args, database_metadata_current_db_only_val, TestListCatalog.test_schema)
-                assert _args["schema_pattern"] in spy.call_args[0][1], print(spy.call_args)
+                assert spy.call_count == 1  # call in get_tables()
 
+            # ensure execute was called with the catalog value in the prepared statement
             if _args["catalog"] is not None:
                 assert _args["catalog"] in spy.call_args[0][0]
 
+            # ensure execute was called with below bind parameters
             for arg in (_args["schema_pattern"], _args["table_name_pattern"]):
                 if arg is not None:
                     assert arg in spy.call_args[0][1]
 
-            assert len(result) > 0, print(spy.call_args, "\n", result)
-            assert len(result[0]) == 10
+            # we cannot easily know what schema pattern matches in Python driver, so
+            # we check table is one of a few options based on whether is_single_database_metadata
+            # is true or false
+
+            possible_not_ds_tables = (
+                "FROM svv_tables",  # universal
+                "FROM pg_catalog.pg_namespace n, pg_catalog.pg_class",  # local
+                "FROM svv_external_tables",  # external
+            )
+            possible_ds_tables = (
+                "FROM PG_CATALOG.SVV_ALL_TABLES",  # universal
+                "FROM pg_catalog.pg_namespace n, pg_catalog.pg_class",  # local
+                "FROM svv_external_tables",  # external
+            )
+
+            if conn.is_single_database_metadata:
+                for table in possible_not_ds_tables:
+                    if table in spy.call_args[0][0]:
+                        assert 1 == 1
+                        return
+                assert 1 == 0, spy.call_args[0][0]
+            else:
+                for table in possible_ds_tables:
+                    if table in spy.call_args[0][0]:
+                        assert 1 == 1
+                        return
+                assert 1 == 0, spy.call_args[0][0]
 
 
 def get_columns_test_data() -> typing.List[typing.Tuple[bool, typing.Dict[str, typing.Optional[str]]]]:
@@ -482,19 +515,19 @@ def test_get_columns(mocker, _input, db_kwargs) -> None:
     database_metadata_current_db_only_val, _args = _input
     db_kwargs["database_metadata_current_db_only"] = database_metadata_current_db_only_val
     with redshift_connector.connect(**db_kwargs) as conn:
-        assert conn.is_single_database_metadata is database_metadata_current_db_only_val
-
         with conn.cursor() as cursor:
             spy = mocker.spy(cursor, "execute")
             result: typing.Tuple = cursor.get_columns(**_args)
-            # ensure query was executed with arguments passed to get_schemas
+            print(result)
+            # ensure execute was called
             assert spy.called
 
-            if _args["schema_pattern"] is not None and database_metadata_current_db_only_val:
+            if _args["schema_pattern"] is not None and conn.is_single_database_metadata:
                 assert spy.call_count == 2  # call in __schema_pattern_match(), get_columns()
             else:
-                assert spy.call_count == 1
+                assert spy.call_count == 1  # call in get_columns()
 
+            # ensure execute was called with below bind parameters
             for arg in (
                 _args["catalog"],
                 _args["schema_pattern"],
@@ -504,5 +537,54 @@ def test_get_columns(mocker, _input, db_kwargs) -> None:
                 if arg is not None:
                     assert arg in spy.call_args[0][0]
 
-            assert len(result) > 0, print(spy.call_args, "\n", result)
-            assert len(result[0]) == 24
+            # we cannot easily know what schema pattern matches in Python driver, so
+            # we check table is one of a few options based on whether is_single_database_metadata
+            # is true or false
+
+            possible_not_ds_tables = (
+                "FROM svv_columns",  # universal
+                "FROM pg_catalog.pg_namespace",  # local
+                "FROM svv_external_columns",  # external
+            )
+            possible_ds_tables = (
+                "FROM PG_CATALOG.svv_all_columns",  # universal
+                "FROM pg_catalog.pg_namespace",  # local
+                "FROM svv_external_columns",  # external
+            )
+
+            if conn.is_single_database_metadata:
+                for table in possible_not_ds_tables:
+                    if table in spy.call_args[0][0]:
+                        assert 1 == 1
+                        return
+                assert 1 == 0, spy.call_args[0][0]
+            else:
+                for table in possible_ds_tables:
+                    if table in spy.call_args[0][0]:
+                        assert 1 == 1
+                        return
+                assert 1 == 0, spy.call_args[0][0]
+
+
+def get_catalogs_test_data() -> typing.List[bool]:
+    return [True, False]
+
+
+@pytest.mark.skip(reason="manual")
+@pytest.mark.parametrize("_input", get_catalogs_test_data())
+def test_get_catalogs(mocker, _input, db_kwargs) -> None:
+    database_metadata_current_db_only_val = _input
+    db_kwargs["database_metadata_current_db_only"] = database_metadata_current_db_only_val
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            spy = mocker.spy(cursor, "execute")
+            result: typing.Tuple = cursor.get_catalogs()
+            print(result)
+            # ensure execute was called
+            assert spy.called
+            assert spy.call_count == 1  # call in get_catalogs()
+
+            if conn.is_single_database_metadata:
+                assert "current_database()" in spy.call_args[0][0]
+            else:
+                assert "PG_CATALOG.SVV_REDSHIFT_DATABASES" in spy.call_args[0][0]

test/unit/test_connection.py

@@ -296,26 +296,76 @@ def test_is_multidatabases_catalog_enable_in_server(_input) -> None:
     assert mock_connection._is_multi_databases_catalog_enable_in_server == exp_val
 
 
-test_is_single_database_metadata_data: typing.List[typing.Tuple[typing.Optional[str], bool, bool]] = [
-    ("on", True, True),
-    ("on", False, False),
-    ("off", True, True),
-    ("off", False, True),
-    (None, True, True),
-    (None, False, True),
+test_is_cross_datasharing_enable_in_server_data: typing.List[typing.Tuple[typing.Optional[str], bool]] = [
+    ("on", True),
+    ("off", False),
+    (None, False),
+]
+
+
+@pytest.mark.parametrize("_input", test_is_cross_datasharing_enable_in_server_data)
+def test_is_cross_datasharing_enable_in_server(_input):
+    param_status, exp_val = _input
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_connection.parameter_statuses: deque = deque()
+
+    if param_status is not None:
+        mock_connection.parameter_statuses.append((b"external_database", param_status.encode()))
+
+    assert mock_connection._is_cross_datasharing_enable_in_server == exp_val
+
+
+def test_is_cross_datasharing_enable_in_server_raises_on_invalid():
+    param_status = "garbage"
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_connection.parameter_statuses: deque = deque()
+
+    if param_status is not None:
+        mock_connection.parameter_statuses.append((b"external_database", param_status.encode()))
+
+    with pytest.raises(
+        InterfaceError,
+        match="Protocol error. Session setup failed. Invalid value of external_database parameter. Only on/off are valid values",
+    ):
+        mock_connection._is_cross_datasharing_enable_in_server
+
+
+test_is_single_database_metadata_data: typing.List[
+    typing.Tuple[typing.Optional[str], typing.Optional[str], bool, bool]
+] = [
+    ("on", "on", True, False),
+    ("on", "on", False, False),
+    ("on", None, True, True),
+    ("on", None, False, False),
+    (None, "on", True, False),
+    (None, "on", False, False),
+    ("off", "on", True, False),
+    ("off", "on", False, False),
+    ("on", "off", True, True),
+    ("on", "off", False, False),
+    ("off", "off", False, True),
+    ("off", "off", True, True),
+    ("off", None, False, True),
+    ("off", None, True, True),
+    (None, "off", False, True),
+    (None, "off", True, True),
+    (None, None, True, True),
+    (None, None, False, True),
 ]
 
 
 @pytest.mark.parametrize("_input", test_is_single_database_metadata_data)
-def test_is_single_database_metadata(_input) -> None:
-    param_status, database_metadata_current_db_only_val, exp_val = _input
+def test_is_single_database_metadata(_input):
+    datashare_enabled, dsw_enabled, database_metadata_current_db_only_val, exp_val = _input
 
     mock_connection: Connection = Connection.__new__(Connection)
     mock_connection.parameter_statuses = deque()
     mock_connection._database_metadata_current_db_only = database_metadata_current_db_only_val
 
-    if param_status is not None:
-        mock_connection.parameter_statuses.append((b"datashare_enabled", param_status.encode()))
+    if datashare_enabled is not None:
+        mock_connection.parameter_statuses.append((b"datashare_enabled", datashare_enabled.encode()))
+    if dsw_enabled is not None:
+        mock_connection.parameter_statuses.append((b"external_database", dsw_enabled.encode()))
 
     assert mock_connection.is_single_database_metadata == exp_val