castor-extractor 0.12.1__tar.gz → 0.12.2__tar.gz

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.12.2 - 2023-12-27
+
+* Metabase: Align postgres connection to postgres connector
+
 ## 0.12.1 - 2023-12-27
 
 * Tableau: replace luid with id during dashboards extraction

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: castor-extractor
-Version: 0.12.1
+Version: 0.12.2
 Summary: Extract your metadata assets.
 Home-page: https://www.castordoc.com/
 License: EULA

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/utils/__init__.py

@@ -1,3 +1,10 @@
+from .client import (
+    AbstractSourceClient,
+    ExtractionQuery,
+    PostgresClient,
+    SqlalchemyClient,
+    uri_encode,
+)
 from .collection import group_by
 from .constants import OUTPUT_DIR
 from .deprecate import deprecate_python
@@ -26,7 +33,6 @@ from .time import (
     past_date,
 )
 from .type import Callback, Getter, JsonType, SerializedAsset
-from .uri import uri_encode
 from .validation import validate_baseurl
 from .write import (
     get_output_filename,

castor_extractor-0.12.2/castor_extractor/utils/client/__init__.py

@@ -0,0 +1,4 @@
+from .abstract import AbstractSourceClient, SqlalchemyClient
+from .postgres import PostgresClient
+from .query import ExtractionQuery
+from .uri import uri_encode

castor_extractor-0.12.1/castor_extractor/warehouse/abstract/client.py → castor_extractor-0.12.2/castor_extractor/utils/client/abstract.py

@@ -7,7 +7,7 @@ from sqlalchemy.engine import Connection, ResultProxy, create_engine
 from .query import ExtractionQuery
 
 
-class AbstractWarehouseClient(ABC):
+class AbstractSourceClient(ABC):
     """interface for the client to connect to the source"""
 
     def __init__(self, **credentials: dict):
@@ -26,7 +26,7 @@ class AbstractWarehouseClient(ABC):
         pass
 
 
-class SqlalchemyClient(AbstractWarehouseClient, ABC):
+class SqlalchemyClient(AbstractSourceClient, ABC):
     def __init__(self, credentials: dict):
         super().__init__(**credentials)
         self._uri = self._build_uri(credentials)

castor_extractor-0.12.1/castor_extractor/warehouse/postgres/client.py → castor_extractor-0.12.2/castor_extractor/utils/client/postgres.py

@@ -1,5 +1,5 @@
-from ...utils import uri_encode
-from ..abstract import SqlalchemyClient
+from .abstract import SqlalchemyClient
+from .uri import uri_encode
 
 POSTGRES_URI = "postgresql://{user}:{password}@{host}:{port}/{database}"
 

castor_extractor-0.12.2/castor_extractor/utils/client/query.py

@@ -0,0 +1,6 @@
+class ExtractionQuery:
+    """contains both statement and parameters for the query"""
+
+    def __init__(self, statement: str, params: dict):
+        self.statement = statement
+        self.params = params

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/visualization/metabase/client/db/client.py

@@ -7,7 +7,12 @@ from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import OperationalError
 
-from .....utils import SerializedAsset, from_env
+from .....utils import (
+    ExtractionQuery,
+    PostgresClient,
+    SerializedAsset,
+    from_env,
+)
 from ...assets import EXPORTED_FIELDS, MetabaseAsset
 from ...errors import EncryptionSecretKeyRequired, MetabaseLoginError
 from ..decryption import decrypt
@@ -18,8 +23,6 @@ logger = logging.getLogger(__name__)
 
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
-PG_URL = "postgresql://{user}:{password}@{host}:{port}/{database}"
-SSL_SUFFIX = "?sslmode=require"
 SQL_FILE_PATH = "queries/{name}.sql"
 
 ENCRYPTION_SECRET_KEY = "CASTOR_METABASE_ENCRYPTION_SECRET_KEY"  # noqa: S105
@@ -33,56 +36,53 @@ def _optional_arg(args: dict, client_key: str, env_key: str) -> Optional[str]:
     return from_env(env_key, allow_missing=True)
 
 
-class DbClient:
+class DbClient(PostgresClient):
     """
     Connect to Metabase Database and fetch main assets.
     """
 
+    @staticmethod
+    def name() -> str:
+        """return the name of the client"""
+        return "Metabase/DB"
+
+    def _engine_options(self, credentials: dict) -> dict:
+        sslmode = "require" if self.require_ssl else "prefer"
+        return {"connect_args": {"sslmode": sslmode}}
+
     def __init__(
         self,
         **kwargs,
     ):
-        self._credentials = CredentialsDb(
-            host=get_value(CredentialsDbKey.HOST, kwargs),
-            port=get_value(CredentialsDbKey.PORT, kwargs),
-            database=get_value(CredentialsDbKey.DATABASE, kwargs),
-            schema=get_value(CredentialsDbKey.SCHEMA, kwargs),
-            user=get_value(CredentialsDbKey.USER, kwargs),
-            password=get_value(CredentialsDbKey.PASSWORD, kwargs),
-        )
+        self._credentials = CredentialsDb.from_env(kwargs)
         self.encryption_secret_key = _optional_arg(
             kwargs, "encryption_secret_key", ENCRYPTION_SECRET_KEY
         )
         self.require_ssl = bool(
             _optional_arg(kwargs, "require_ssl", REQUIRE_SSL_KEY)
         )
-        self._engine = self._login()
-
-    def _url(self) -> str:
-        url = PG_URL.format(**self._credentials.to_dict(hide=False))
-        if self.require_ssl:
-            url += SSL_SUFFIX
-        return url
-
-    def _login(self) -> Engine:
-        url = self._url()
         try:
-            engine = create_engine(url)
-            engine.connect()
-            return engine
+            super().__init__(self._credentials.to_dict())
         except OperationalError as err:
             raise MetabaseLoginError(
                 credentials_info=self._credentials.to_dict(hide=True),
                 error_details=err.args,
             )
 
-    def _load_query(self, name: str) -> str:
+    def _load_query(self, name: str) -> ExtractionQuery:
         """load SQL text from file"""
         filename = SQL_FILE_PATH.format(name=name)
         path = os.path.join(CURRENT_DIR, filename)
         with open(path, "r") as f:
             content = f.read()
-            return content.format(schema=self._credentials.schema)
+            statement = content.format(schema=self._credentials.schema)
+            return ExtractionQuery(statement=statement, params={})
+
+    def base_url(self) -> str:
+        """Fetches the `base_url` of the Metabase instance"""
+        query = self._load_query(name="base_url")
+        rows = list(self.execute(query))
+        return rows[0]["value"]
 
     def _database_specifics(
         self,
@@ -106,27 +106,10 @@ class DbClient:
 
         return databases
 
-    @staticmethod
-    def name() -> str:
-        """return the name of the client"""
-        return "Metabase/DB"
-
-    def base_url(self) -> str:
-        """Fetches the `base_url` of the Metabase instance"""
-        query = self._load_query(name="base_url")
-        result = self._engine.execute(query)
-        return result.fetchone()["value"]
-
     def fetch(self, asset: MetabaseAsset) -> SerializedAsset:
         """fetches the given asset"""
         query = self._load_query(asset.value.lower())
-        result = self._engine.execute(query)
-
-        # convert row proxies into dictionaries
-        assets = [
-            {column: value for column, value in rowproxy.items()}
-            for rowproxy in result
-        ]
+        assets = list(self.execute(query))
 
         if asset == MetabaseAsset.DATABASE:
             assets = self._database_specifics(assets)

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/visualization/metabase/client/db/credentials.py

@@ -52,6 +52,18 @@ class CredentialsDb:
         self.user = user
         self.password = password
 
+    @classmethod
+    def from_env(cls, env: dict) -> "CredentialsDb":
+        """returns a new CredentialsDb with values from ENV"""
+        return CredentialsDb(
+            host=get_value(CredentialsDbKey.HOST, env),
+            port=get_value(CredentialsDbKey.PORT, env),
+            database=get_value(CredentialsDbKey.DATABASE, env),
+            schema=get_value(CredentialsDbKey.SCHEMA, env),
+            user=get_value(CredentialsDbKey.USER, env),
+            password=get_value(CredentialsDbKey.PASSWORD, env),
+        )
+
     def to_dict(self, hide: bool = False) -> Dict[str, str]:
         safe = (
             CredentialsDbKey.HOST,

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/abstract/__init__.py

@@ -6,7 +6,6 @@ from .asset import (
     WarehouseAsset,
     WarehouseAssetGroup,
 )
-from .client import AbstractWarehouseClient, SqlalchemyClient
 from .extract import ExtractionProcessor, common_args
 from .query import (
     QUERIES_DIR,

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/abstract/extract.py

@@ -2,9 +2,15 @@ import logging
 from itertools import chain
 from typing import Callable, Iterator, List, Optional, Tuple
 
-from ...utils import OUTPUT_DIR, AbstractStorage, SafeMode, from_env, safe_mode
+from ...utils import (
+    OUTPUT_DIR,
+    AbstractSourceClient,
+    AbstractStorage,
+    SafeMode,
+    from_env,
+    safe_mode,
+)
 from .asset import WarehouseAsset
-from .client import AbstractWarehouseClient
 from .query import AbstractQueryBuilder, ExtractionQuery
 
 logger = logging.getLogger(__name__)
@@ -22,7 +28,7 @@ class ExtractionProcessor:
 
     def __init__(
         self,
-        client: AbstractWarehouseClient,
+        client: AbstractSourceClient,
         query_builder: AbstractQueryBuilder,
         storage: AbstractStorage,
         safe_mode: Optional[SafeMode] = None,

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/bigquery/client.py

@@ -7,8 +7,7 @@ from google.cloud.bigquery import Client as GoogleCloudClient  # type: ignore
 from google.cloud.bigquery.dataset import Dataset  # type: ignore
 from google.oauth2.service_account import Credentials  # type: ignore
 
-from ...utils import retry
-from ..abstract import SqlalchemyClient
+from ...utils import SqlalchemyClient, retry
 
 logger = logging.getLogger(__name__)
 

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/postgres/__init__.py

@@ -1,3 +1,2 @@
-from .client import PostgresClient
 from .extract import POSTGRES_ASSETS, extract_all
 from .query import PostgresQueryBuilder

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/postgres/extract.py

@@ -1,7 +1,7 @@
 import logging
 from collections import OrderedDict
 
-from ...utils import LocalStorage, from_env, write_summary
+from ...utils import LocalStorage, PostgresClient, from_env, write_summary
 from ..abstract import (
     CATALOG_ASSETS,
     ExtractionProcessor,
@@ -10,7 +10,6 @@ from ..abstract import (
     WarehouseAssetGroup,
     common_args,
 )
-from .client import PostgresClient
 from .query import PostgresQueryBuilder
 
 logger = logging.getLogger(__name__)

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/redshift/client.py

@@ -4,8 +4,7 @@ from typing import Iterator
 from psycopg2 import extensions  # type: ignore
 from sqlalchemy.engine import Connection, ResultProxy
 
-from ...utils import decode_when_bytes, uri_encode
-from ..abstract import SqlalchemyClient
+from ...utils import SqlalchemyClient, decode_when_bytes, uri_encode
 
 REDSHIFT_URI = "redshift+psycopg2://{user}:{password}@{host}:{port}/{database}"
 

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/redshift/client_test.py

@@ -15,7 +15,7 @@ def fake_execute(statement: str, params: dict):
 
 @patch.object(RedshiftClient, "_build_uri")
 @patch.object(RedshiftClient, "_engine_options")
-@patch(f"{CASTOR_EXTRACTOR_PATH}.warehouse.abstract.client.create_engine")
+@patch(f"{CASTOR_EXTRACTOR_PATH}.utils.client.abstract.create_engine")
 @patch.object(RedshiftClient, "close")
 @patch.object(RedshiftClient, "connect")
 def test_redshift_client_execute(

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/snowflake/client.py

@@ -7,8 +7,7 @@ from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from sqlalchemy.engine import Connection
 
-from ...utils import to_string_array, uri_encode
-from ..abstract import SqlalchemyClient
+from ...utils import SqlalchemyClient, to_string_array, uri_encode
 
 SNOWFLAKE_URI = "snowflake://{user}:{password}@{account}/?application=castor"
 PK_URI = "snowflake://{user}:@{account}/?application=castor"

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/snowflake/client_test.py

@@ -2,7 +2,7 @@
 
 from unittest.mock import patch
 
-from ..abstract import SqlalchemyClient
+from ...utils import SqlalchemyClient
 from .client import SnowflakeClient, _scalar
 
 

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/castor_extractor/warehouse/sqlserver/client.py

@@ -3,8 +3,7 @@ from typing import Iterator
 import pymssql  # type: ignore
 from sqlalchemy import text
 
-from ...utils import uri_encode
-from ..abstract import ExtractionQuery, SqlalchemyClient
+from ...utils import ExtractionQuery, SqlalchemyClient, uri_encode
 
 SERVER_URI = "{user}:{password}@{host}:{port}/{database}"
 MSSQL_URI = f"mssql+pymssql://{SERVER_URI}"

{castor_extractor-0.12.1 → castor_extractor-0.12.2}/pyproject.toml

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