schema_sentinel.metadata_manager.engine

  1import logging as log
  2import os
  3from abc import ABC, abstractmethod
  4from configparser import ConfigParser
  5
  6import pandas as pd
  7import snowflake.connector
  8from cryptography.hazmat.backends import default_backend
  9from cryptography.hazmat.primitives import serialization
 10from snowflake.sqlalchemy import URL
 11from sqlalchemy import MetaData, create_engine, text
 12from sqlalchemy.engine import Engine
 13
 14from .enums import ConnectMode
 15
 16
 17class DBEngineStrategy(ABC):
 18    engine = None
 19    conn = None
 20    metadata: MetaData = None
 21    connect_mode = ConnectMode.KEY_PAIR.value
 22    env: str = "dev"
 23    url: str = None
 24
 25    def __init__(self, config: dict, env: str = "dev"):
 26        self.account = config.get("account")
 27        self.user = config.get("user")
 28        self.warehouse = config.get("warehouse")
 29        self.database = config.get("database")
 30        self.private_key = config.get("private_key") if "private_key" in config else b""
 31        self.role = config.get("role")
 32        self.engine = None
 33        self.private_key_passphrase = config.get("private_key_passphrase") if "private_key_passphrase" in config else ""
 34        self.default_schema = config.get("schema")
 35        self.connect_mode = config.get("connect_mode")
 36        self.env = env
 37        self.url = config.get("url")
 38
 39    @abstractmethod
 40    def get_engine(self):
 41        pass
 42
 43    @abstractmethod
 44    def get_conn(self):
 45        pass
 46
 47    @abstractmethod
 48    def close(self):
 49        pass
 50
 51    def __del__(self):
 52        self.close()
 53
 54    @abstractmethod
 55    def execute(self, statement, parameters=None, columns=None, schema=None):
 56        pass
 57
 58    def decode_private_key(self, password):
 59        p_key = serialization.load_pem_private_key(
 60            self.private_key, password=password.encode(), backend=default_backend()
 61        )
 62
 63        pkb = p_key.private_bytes(
 64            encoding=serialization.Encoding.DER,
 65            format=serialization.PrivateFormat.PKCS8,
 66            encryption_algorithm=serialization.NoEncryption(),
 67        )
 68
 69        return pkb
 70
 71
 72class SqLiteAqlAlchemyEngine(DBEngineStrategy):
 73    metadata: MetaData = None
 74    engine: Engine = None
 75    url: str = None
 76
 77    def __int__(self, config: dict, env: str, url: str):
 78        config["url"] = url
 79        super().__init__(env=None, config=config)
 80
 81    def get_engine(self):
 82        if self.engine is None:
 83            self.engine = create_engine(url=self.url)
 84        return self.engine
 85
 86    def close(self):
 87        if self.engine:
 88            self.engine.dispose()
 89
 90    def execute(self, statement, parameters=None, columns=None, schema=None):
 91        log.debug(f"statement=[{statement}]")
 92        conn = self.get_conn()
 93
 94        if not columns:
 95            if parameters is not None:
 96                conn.execute(text(statement), parameters)
 97            else:
 98                conn.execute(text(statement))
 99            return True, None
100
101        if parameters is not None:
102            results = pd.read_sql(sql=text(statement), con=conn, params=parameters)
103        else:
104            results = pd.read_sql(sql=text(statement), con=conn)
105
106        return True, results
107
108    def get_conn(self):
109        return self.get_engine()
110
111
112class SfAlchemyEngine(DBEngineStrategy):
113    snowflake.connector.paramstyle = "numeric"
114
115    def get_engine(self):
116        if self.engine is None:
117            if self.connect_mode == ConnectMode.SSO.value:
118                self.engine = create_engine(
119                    URL(
120                        account=self.account,
121                        user=self.user,
122                        warehouse=self.warehouse,
123                        database=self.database,
124                        role=self.role,
125                        cache_column_metadata=True,
126                        schema=self.default_schema,
127                        authenticator="externalbrowser",
128                    )
129                )
130            if self.connect_mode == ConnectMode.PWD.value:
131                self.engine = create_engine(
132                    URL(
133                        account=self.account,
134                        user=self.user,
135                        warehouse=self.warehouse,
136                        database=self.database,
137                        role=self.role,
138                        cache_column_metadata=True,
139                        schema=self.default_schema,
140                        password=os.getenv("SNOWFLAKE_PASSWORD"),
141                    )
142                )
143            if self.connect_mode == ConnectMode.KEY_PAIR.value:
144                pkb = self.decode_private_key(self.private_key_passphrase)
145                self.engine = create_engine(
146                    URL(
147                        account=self.account,
148                        user=self.user,
149                        warehouse=self.warehouse,
150                        database=self.database,
151                        role=self.role,
152                        cache_column_metadata=True,
153                        schema=self.default_schema,
154                    ),
155                    connect_args={
156                        "private_key": pkb,
157                    },
158                )
159
160        return self.engine
161
162    def get_conn(self):
163        if self.conn is None:
164            engine = self.get_engine()
165            self.conn = engine.connect().execution_options(autocommit=False)
166        return self.conn
167
168    def close(self):
169        if self.conn is not None:
170            self.conn.close()
171        if self.engine is not None:
172            self.engine.dispose()
173
174    def execute(self, statement, parameters=None, columns=None, schema=None):
175        log.debug(f"statement=[{statement}]")
176        conn = self.get_conn()
177        if schema is not None:
178            # Properly quote schema identifier - escape quotes by doubling them (SQL standard)
179            safe_schema = schema.replace('"', '""')
180            conn.execute(f'USE SCHEMA "{safe_schema}";')
181
182        if columns is None:
183            conn.execute("BEGIN")
184            if parameters is not None:
185                conn.execute(text(statement), parameters)
186            else:
187                conn.execute(text(statement))
188            conn.execute("COMMIT")
189            return True, None
190
191        if parameters is not None:
192            results = pd.read_sql(sql=text(statement), con=conn, params=parameters)
193        else:
194            results = pd.read_sql(sql=text(statement), con=conn)
195
196        return True, results
197
198
199class SfConnectorEngine(DBEngineStrategy):
200    def get_engine(self, schema=None):
201        pkb = self.decode_private_key(self.private_key_passphrase)
202
203        if self.engine is None:
204            snowflake.connector.paramstyle = "numeric"
205            self.engine = snowflake.connector.connect(
206                user=self.user,
207                private_key=pkb,
208                account=self.account,
209                warehouse=self.warehouse,
210                database=self.database,
211                schema=self.default_schema,
212                role=self.role,
213                autocommit=False,
214            )
215        return self.engine
216
217    def get_conn(self):
218        if self.conn is None:
219            if self.connect_mode == ConnectMode.SSO.value:
220                snowflake.connector.paramstyle = "numeric"
221                self.conn = snowflake.connector.connect(
222                    account=self.account,
223                    user=self.user,
224                    role=self.role,
225                    database=self.database,
226                    schema=self.schema,
227                    warehouse=self.warehouse,
228                    authenticator="externalbrowser",
229                )
230
231            if self.connect_mode == ConnectMode.PWD.value:
232                self.conn = snowflake.connector.connect(
233                    account=self.account,
234                    user=self.user,
235                    role=self.role,
236                    database=self.database,
237                    schema=self.schema,
238                    warehouse=self.warehouse,
239                    password=os.getenv("SNOWFLAKE_PASSWORD"),
240                )
241
242            if self.connect_mode == ConnectMode.KEY_PAIR.value:
243                pkb = self.decode_private_key(self.private_key_passphrase)
244                self.conn = snowflake.connector.connect(
245                    account=self.account,
246                    user=self.user,
247                    role=self.role,
248                    database=self.database,
249                    schema=self.default_schema,
250                    warehouse=self.warehouse,
251                    private_key=pkb,
252                )
253        return self.conn
254
255    def close(self):
256        if self.engine is not None:
257            self.engine.close()
258
259    def execute(self, statement, parameters=None, columns=None, schema=None):
260        log.debug(f"statement=[{statement}]")
261        conn = self.get_engine()
262
263        if schema is not None:
264            # Properly quote schema identifier - escape quotes by doubling them (SQL standard)
265            safe_schema = schema.replace('"', '""')
266            conn.cursor().execute(f'USE SCHEMA "{safe_schema}";')
267
268        if columns is None:
269            conn.cursor().execute(statement, parameters)
270            conn.commit()
271            return True, None
272
273        if parameters is not None:
274            results = pd.read_sql(text(statement), conn, *parameters)
275        else:
276            results = pd.read_sql(text(statement), conn)
277
278        return True, results
279
280
281def get_config_dict(
282    config: ConfigParser,
283    *,
284    private_key: bytes,
285    password: str,
286    connect_mode: int = ConnectMode.KEY_PAIR.value,
287    user: str = None,
288    cache_column_metadata=False,
289):
290    db_section = "DB_CONNECTION"
291    config_map = {
292        "account": config.get(db_section, "account"),
293        "warehouse": config.get(db_section, "warehouse"),
294        "user": user,
295        "database": config.get(db_section, "database"),
296        "role": config.get(db_section, "role"),
297        "schema": config.get(db_section, "schema"),
298        "connect_mode": connect_mode,
299        "cache_column_metadata": cache_column_metadata,
300        "timezone": "America/Los_Angeles",
301    }
302    if connect_mode == ConnectMode.PWD.value:
303        config_map["password"] = password
304    elif connect_mode == ConnectMode.KEY_PAIR.value:
305        config_map["private_key"] = private_key
306        config_map["private_key_passphrase"] = password
307
308    return config_map