castor-extractor 0.20.4__tar.gz → 0.20.5__tar.gz

{castor_extractor-0.20.4 → castor_extractor-0.20.5}/CHANGELOG.md

@@ -1,6 +1,10 @@
 
 # Changelog
 
+## 0.20.5 - 2024-10-09
+
+* Redshift: enable extraction from a Redshift Serverless instance
+
 ## 0.20.4 - 2024-10-09
 
 * Salesforce warehouse: `Labels` instead of `api_names` for columns

{castor_extractor-0.20.4 → castor_extractor-0.20.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.20.4
+Version: 0.20.5
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA
@@ -208,6 +208,10 @@ For any questions or bug report, contact us at [support@castordoc.com](mailto:su
 
 # Changelog
 
+## 0.20.5 - 2024-10-09
+
+* Redshift: enable extraction from a Redshift Serverless instance
+
 ## 0.20.4 - 2024-10-09
 
 * Salesforce warehouse: `Labels` instead of `api_names` for columns

{castor_extractor-0.20.4 → castor_extractor-0.20.5}/castor_extractor/commands/extract_redshift.py

@@ -23,6 +23,11 @@ def main():
         action="store_true",
         help="Skips files already extracted instead of replacing them",
     )
+    parser.add_argument(
+        "--serverless",
+        action="store_true",
+        help="Enables extraction for Redshift Serverless",
+    )
     parser.set_defaults(skip_existing=False)
 
     args = parser.parse_args()
@@ -34,5 +39,6 @@ def main():
         user=args.user,
         password=args.password,
         output_directory=args.output,
+        serverless=args.serverless,
         skip_existing=args.skip_existing,
     )

{castor_extractor-0.20.4 → castor_extractor-0.20.5}/castor_extractor/warehouse/redshift/extract.py

@@ -34,6 +34,7 @@ REDSHIFT_PASSWORD = "CASTOR_REDSHIFT_PASSWORD"  # noqa: S105
 REDSHIFT_HOST = "CASTOR_REDSHIFT_HOST"
 REDSHIFT_PORT = "CASTOR_REDSHIFT_PORT"
 REDSHIFT_DATABASE = "CASTOR_REDSHIFT_DATABASE"
+REDSHIFT_SERVERLESS = "CASTOR_REDSHIFT_SERVERLESS"
 
 
 def _credentials(params: dict) -> dict:
@@ -48,6 +49,14 @@ def _credentials(params: dict) -> dict:
     }
 
 
+def _query_builder(params: dict) -> RedshiftQueryBuilder:
+    env_parameter = from_env(REDSHIFT_SERVERLESS, allow_missing=True)
+    from_env_ = str(env_parameter).lower() == "true"
+    from_params_ = params.get("serverless", False)
+    is_serverless = from_params_ or from_env_
+    return RedshiftQueryBuilder(is_serverless=is_serverless)
+
+
 def extract_all(**kwargs) -> None:
     """
     Extract all assets from Redshift and store the results in CSV files
@@ -56,7 +65,7 @@ def extract_all(**kwargs) -> None:
 
     client = RedshiftClient(credentials=_credentials(kwargs))
 
-    query_builder = RedshiftQueryBuilder()
+    query_builder = _query_builder(kwargs)
 
     storage = LocalStorage(directory=output_directory)
 

castor_extractor-0.20.5/castor_extractor/warehouse/redshift/extract_test.py

@@ -0,0 +1,26 @@
+import pytest
+
+from .extract import (
+    REDSHIFT_SERVERLESS,
+    _query_builder,
+)
+
+
+@pytest.mark.parametrize(
+    "serverless_param,env_param,expected",
+    [
+        (True, "False", True),
+        (False, "True", True),
+        (None, "TRUE", True),
+        (None, "TrUe", True),
+        (None, "FAlSE", False),
+        (None, "False", False),
+        (None, None, False),
+        (True, None, True),
+    ],
+)
+def test__query_builder(serverless_param, env_param, expected, monkeypatch):
+    params = {"serverless": serverless_param}
+    monkeypatch.setenv(REDSHIFT_SERVERLESS, env_param)
+
+    assert _query_builder(params).is_serverless == expected

castor_extractor-0.20.5/castor_extractor/warehouse/redshift/queries/query_serverless.sql

@@ -0,0 +1,69 @@
+WITH parameters AS (
+    SELECT
+        :day AS day_start,
+        :hour_min AS hour_min,
+        :hour_max AS hour_max
+),
+
+queries_deduplicated AS (
+    SELECT DISTINCT q.query_id
+    FROM SYS_QUERY_HISTORY AS q
+        CROSS JOIN parameters AS p
+    WHERE TRUE
+        AND DATE(q.start_time) = p.day_start
+        AND EXTRACT('hour' FROM q.start_time) BETWEEN p.hour_min AND p.hour_max
+),
+
+query AS (
+    SELECT
+        q.query_id,
+        qt.text,
+        qt.sequence,
+        COUNT(*) OVER(PARTITION BY q.query_id) AS sequence_count
+    FROM queries_deduplicated AS q
+        INNER JOIN SYS_QUERY_TEXT AS qt ON q.query_id = qt.query_id
+),
+
+raw_query_text AS
+(
+    SELECT
+        q.query_id,
+        LISTAGG(q.text, '') WITHIN GROUP (ORDER BY q.sequence) AS agg_text
+    FROM query AS q
+    WHERE TRUE
+    -- LISTAGG raises an error when total length >= 65535
+    -- each query text contains 4000 char max
+        AND q.sequence_count < (65535 / 4000)
+    GROUP BY q.query_id
+),
+
+query_text AS (
+    SELECT
+        query_id,
+        CASE
+            WHEN agg_text ILIKE 'INSERT INTO%%'
+                THEN REGEXP_REPLACE(agg_text, 'VALUES (.*)', 'DEFAULT VALUES')
+            ELSE agg_text
+        END AS agg_text
+    FROM raw_query_text
+)
+SELECT
+    q.query_id::VARCHAR(256) AS query_id,
+    qt.agg_text::VARCHAR(60000) AS query_text,
+    q.database_name AS database_id,
+    q.database_name AS database_name,
+    q.session_id AS process_id,
+    0 as aborted,
+    q.start_time AS start_time,
+    q.end_time AS end_time,
+    q.user_id AS user_id,
+    q.query_label,
+    u.usename AS user_name
+FROM SYS_QUERY_HISTORY AS q
+    JOIN query_text AS qt ON q.query_id = qt.query_id
+    JOIN pg_catalog.pg_user AS u ON u.usesysid = q.user_id
+    CROSS JOIN parameters AS p
+WHERE TRUE
+    AND DATE(q.start_time) = p.day_start
+    AND EXTRACT('hour' FROM q.start_time) BETWEEN p.hour_min AND p.hour_max
+    AND q.status = 'success'

castor_extractor-0.20.5/castor_extractor/warehouse/redshift/query.py

@@ -0,0 +1,35 @@
+from typing import List, Optional
+
+from ..abstract import (
+    AbstractQueryBuilder,
+    ExtractionQuery,
+    TimeFilter,
+    WarehouseAsset,
+)
+
+
+class RedshiftQueryBuilder(AbstractQueryBuilder):
+    """
+    Builds queries to extract assets from Redshift.
+    """
+
+    def __init__(
+        self,
+        is_serverless: bool = False,
+        time_filter: Optional[TimeFilter] = None,
+    ):
+        super().__init__(time_filter=time_filter)
+        self.is_serverless = is_serverless
+
+    def build_query_serverless(self) -> ExtractionQuery:
+        """To get the query history in Redshift Serverless, we cannot use STL tables."""
+        statement = self._load_from_file("query_serverless.sql")
+        params = self._time_filter.to_dict()
+        return ExtractionQuery(statement, params)
+
+    def build(self, asset: WarehouseAsset) -> List[ExtractionQuery]:
+        if asset == WarehouseAsset.QUERY and self.is_serverless:
+            query = self.build_query_serverless()
+        else:
+            query = self.build_default(asset)
+        return [query]

{castor_extractor-0.20.4 → castor_extractor-0.20.5}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=61.2"]
 
 [tool.poetry]
 name = "castor-extractor"
-version = "0.20.4"
+version = "0.20.5"
 description = "Extract your metadata assets."
 authors = ["Castor <support@castordoc.com>"]
 license = "EULA"

castor_extractor-0.20.4/castor_extractor/warehouse/redshift/query.py

@@ -1,24 +0,0 @@
-from typing import List, Optional
-
-from ..abstract import (
-    AbstractQueryBuilder,
-    ExtractionQuery,
-    TimeFilter,
-    WarehouseAsset,
-)
-
-
-class RedshiftQueryBuilder(AbstractQueryBuilder):
-    """
-    Builds queries to extract assets from Redshift.
-    """
-
-    def __init__(
-        self,
-        time_filter: Optional[TimeFilter] = None,
-    ):
-        super().__init__(time_filter=time_filter)
-
-    def build(self, asset: WarehouseAsset) -> List[ExtractionQuery]:
-        query = self.build_default(asset)
-        return [query]