linkmerce 0.3.2__tar.gz → 0.3.4__tar.gz

{linkmerce-0.3.2 → linkmerce-0.3.4}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.3
 Name: linkmerce
-Version: 0.3.2
+Version: 0.3.4
 Summary: E-commerce API integration management
 Requires-Dist: aiohttp>=3.12.15
 Requires-Dist: bs4>=0.0.2
 Requires-Dist: duckdb>=1.3.2
 Requires-Dist: jinja2>=3.1.6
+Requires-Dist: nest-asyncio>=1.6.0
 Requires-Dist: pytz>=2025.2
 Requires-Dist: requests>=2.32.4
 Requires-Dist: tqdm>=4.67.1

{linkmerce-0.3.2 → linkmerce-0.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "linkmerce"
-version = "0.3.2"
+version = "0.3.4"
 description = "E-commerce API integration management"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -9,6 +9,7 @@ dependencies = [
     "bs4>=0.0.2",
     "duckdb>=1.3.2",
     "jinja2>=3.1.6",
+    "nest-asyncio>=1.6.0",
     "pytz>=2025.2",
     "requests>=2.32.4",
     "tqdm>=4.67.1",

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/api/searchad/manage.py

@@ -15,7 +15,7 @@ def get_module(name: str) -> str:
     return (".searchad.manage" + name) if name.startswith('.') else name
 
 
-def exposure_diagnosis(
+def diagnose_exposure(
         customer_id: str,
         cookies: str,
         keyword: str | Iterable[str],
@@ -34,7 +34,7 @@ def exposure_diagnosis(
     return run_with_duckdb(get_module(".exposure"), "ExposureDiagnosis", "ExposureDiagnosis", connection, "sync", table, return_type, args, **options)
 
 
-def exposure_rank(
+def rank_exposure(
         customer_id: str,
         cookies: str,
         keyword: str | Iterable[str],

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/common/load.py

@@ -21,6 +21,12 @@ def concat_sql(*statement: str, drop_empty: bool = True, sep=' ', terminate: boo
     return query + ';' if terminate and not query.endswith(';') else query
 
 
+def where(where_clause: str | None = None, default: str | None = None) -> str:
+    if not (where_clause or default):
+        return str()
+    return f"WHERE {where_clause or default}"
+
+
 def csv_to_json(obj: list[tuple], header: int | list[str] = 0) -> list[dict]:
     if isinstance(header, int):
         header, obj = obj[header], obj[header+1:]
@@ -112,16 +118,12 @@ class Connection(metaclass=ABCMeta):
 
     ############################ Expression ###########################
 
-    def expr(self, value: Any, type: str, alias: str = str(), safe: bool = False) -> str:
-        type = type.upper()
-        if type == "DATE":
-            return self.expr_date(value, alias, safe)
-        else:
-            func = "TRY_CAST" if safe else "CAST"
-            alias = f" AS {alias}" if alias else str()
-            return f"{func}({value} AS {type})" + alias
+    def expr_cast(self, value: Any | None, type: str, alias: str = str(), safe: bool = False) -> str:
+        cast = "TRY_CAST" if safe else "CAST"
+        alias = f" AS {alias}" if alias else str()
+        return f"{cast}({self.expr_value(value)} AS {type.upper()})" + alias
 
-    def expr_create(self, option: Literal["replace", "ignore"] | None = None, temp: bool = False) -> str:
+    def expr_create(self, option: Literal["replace","ignore"] | None = None, temp: bool = False) -> str:
         temp = "TEMP" if temp else str()
         if option == "replace":
             return f"CREATE OR REPLACE {temp} TABLE"
@@ -130,20 +132,14 @@ class Connection(metaclass=ABCMeta):
         else:
             return f"CREATE {temp} TABLE"
 
-    def expr_date(self, value: dt.date | str | None = None, alias: str = str(), safe: bool = False) -> str:
-        alias = f" AS {alias}" if alias else str()
-        if safe:
-            return (f"DATE '{value}'" if value is not None else "NULL") + alias
-        else:
-            return f"DATE '{value}'" + alias
-
-    def expr_interval(value: str | int | None = None) -> str:
-        if isinstance(value, str):
-            return value
-        elif isinstance(value, int):
-            return "{} INTERVAL {} DAY".format('-' if value < 0 else '+', abs(value))
+    def expr_value(self, value: Any | None) -> str:
+        import datetime as dt
+        if value is None:
+            return "NULL"
+        elif isinstance(value, (float,int)):
+            return str(value)
         else:
-            return str()
+            return f"'{value}'"
 
     def expr_now(
             self,
@@ -173,14 +169,22 @@ class Connection(metaclass=ABCMeta):
             return f"STRFTIME({expr}, '{format}')"
         return expr if type.upper() == "DATE" else "NULL"
 
+    def expr_interval(days: str | int | None = None) -> str:
+        if isinstance(days, str):
+            return days
+        elif isinstance(days, int):
+            return "{} INTERVAL {} DAY".format('-' if days < 0 else '+', abs(days))
+        else:
+            return str()
+
 
 ###################################################################
 ############################## DuckDB #############################
 ###################################################################
 
 class DuckDBConnection(Connection):
-    def __init__(self, **kwargs):
-        self.set_connection(**kwargs)
+    def __init__(self, tzinfo: str | None = None, **kwargs):
+        self.set_connection(tzinfo, **kwargs)
 
     @property
     def conn(self) -> DuckDBPyConnection:
@@ -189,9 +193,11 @@ class DuckDBConnection(Connection):
     def get_connection(self) -> DuckDBPyConnection:
         return self.__conn
 
-    def set_connection(self, **kwargs):
+    def set_connection(self, tzinfo: str | None = None, **kwargs):
         import duckdb
         self.__conn = duckdb.connect(**kwargs)
+        if tzinfo is not None:
+            self.conn.execute(f"SET TimeZone = '{tzinfo}';")
 
     def close(self):
         try:
@@ -258,7 +264,7 @@ class DuckDBConnection(Connection):
             save_to: str | Path | None = None,
         ) -> list[tuple] | None:
         relation = self.conn.execute(query, parameters=params)
-        results = [self.get_columns(relation)] + relation.fetchall()
+        results = [tuple(self.get_columns(relation))] + relation.fetchall()
         if save_to:
             return save_to_csv(results, save_to, delimiter=',')
         else:
@@ -358,7 +364,7 @@ class DuckDBConnection(Connection):
             table: str,
             values: list[tuple] | list[dict] | bytes | str | Path,
             format: Literal["csv","json","parquet"],
-            option: Literal["replace", "ignore"] | None = None,
+            option: Literal["replace","ignore"] | None = None,
             temp: bool = False,
             params: dict | None = None,
         ) -> DuckDBPyConnection:
@@ -371,7 +377,7 @@ class DuckDBConnection(Connection):
             self,
             table: str,
             values: list[tuple] | str | Path,
-            option: Literal["replace", "ignore"] | None = None,
+            option: Literal["replace","ignore"] | None = None,
             temp: bool = False,
             params: dict | None = None,
         ) -> DuckDBPyConnection:
@@ -381,7 +387,7 @@ class DuckDBConnection(Connection):
             self,
             table: str,
             values: list[dict] | str | Path,
-            option: Literal["replace", "ignore"] | None = None,
+            option: Literal["replace","ignore"] | None = None,
             temp: bool = False,
             params: dict | None = None,
         ) -> DuckDBPyConnection:
@@ -391,7 +397,7 @@ class DuckDBConnection(Connection):
             self,
             table: str,
             values: bytes | str | Path,
-            option: Literal["replace", "ignore"] | None = None,
+            option: Literal["replace","ignore"] | None = None,
             temp: bool = False,
             params: dict | None = None,
         ) -> DuckDBPyConnection:
@@ -399,14 +405,14 @@ class DuckDBConnection(Connection):
 
     def copy_table(
             self,
-            table: str,
-            copy_to: str,
+            source_table: str,
+            target_table: str,
             limit: int | None = 0,
-            option: Literal["replace", "ignore"] | None = None,
+            option: Literal["replace","ignore"] | None = None,
             temp: bool = False,
         ) -> DuckDBPyConnection:
         limit_ = f"LIMIT {limit}" if isinstance(limit, int) else None
-        query = concat_sql(f"{self.expr_create(option, temp)} {copy_to} AS SELECT * FROM {table}", limit_)
+        query = concat_sql(f"{self.expr_create(option, temp)} {target_table} AS SELECT * FROM {source_table}", limit_)
         return self.conn.execute(query)
 
     ############################## Insert #############################
@@ -478,23 +484,26 @@ class DuckDBConnection(Connection):
                 elif agg in {"first","last","list"}:
                     return f"{agg.upper()}({col}) FILTER (WHERE {col} IS NOT NULL)"
                 else:
-                    return agg
+                    return f"{agg}({col})"
             return ", ".join([f"{render(col, agg)} AS {col}" for col, agg in func.items()])
         else:
             return func
 
     ############################## Utils ##############################
 
-    def has_table(self, table: str) -> bool:
+    def table_exists(self, table: str) -> bool:
         query = f"SELECT 1 FROM information_schema.tables WHERE table_name = '{table}' LIMIT 1;"
         return bool(self.conn.execute(query).fetchone())
 
-    def exists_table(self, table: str) -> bool:
-        if self.has_table(table):
+    def table_has_rows(self, table: str) -> bool:
+        if self.table_exists(table):
             query = f"SELECT 1 FROM {table} LIMIT 1;"
             return bool(self.conn.execute(query).fetchone())
         return False
 
+    def count_table(self, table: str) -> int:
+        return self.conn.execute(f"SELECT COUNT(*) FROM {table}").fetchall()[0][0]
+
     def get_columns(self, obj: str | DuckDBPyConnection) -> list[str]:
         if isinstance(obj, str):
             obj = self.conn.execute(f"DESCRIBE {obj}")
@@ -505,31 +514,12 @@ class DuckDBConnection(Connection):
     def has_column(self, obj: str | DuckDBPyConnection, column: str) -> bool:
         return column in self.get_columns(obj)
 
-    def unique(self, table: str, expr: str, ascending: bool | None = None, condition: str | None = None) -> list:
+    def unique(self, table: str, expr: str, ascending: bool | None = None, where_clause: str | None = None) -> list:
         select = f"SELECT DISTINCT {expr} AS expr FROM {table}"
         order_by = "ORDER BY expr {}".format({True:"ASC", False:"DESC"}[ascending]) if isinstance(ascending, bool) else None
-        query = concat_sql(select, self.expr_where(condition), order_by)
+        query = concat_sql(select, where(where_clause), order_by)
         return [row[0] for row in self.conn.execute(query).fetchall()]
 
-    def expr_where(self, condition: str | None = None, column: str | None = None, default: str = "WHERE TRUE") -> str:
-        if condition:
-            if condition.split(' ', maxsplit=1)[0].upper() == "WHERE":
-                return condition
-            else:
-                return concat_sql("WHERE", column, condition, terminate=False)
-        else:
-            return default
-
-    def expr_value(self, value: Any) -> str:
-        import datetime as dt
-        if isinstance(value, (float,int)):
-            return str(value)
-        elif isinstance(value, dt.date):
-            dtype = "TIMESTAMP" if isinstance(value, dt.datetime) else "DATE"
-            return f"{dtype} '{value}'"
-        else:
-            return f"'{value}'"
-
 
 ###################################################################
 ############################# Iterator ############################
@@ -564,7 +554,7 @@ class DuckDBIterator(Task):
             self,
             by: str | list[str],
             ascending: bool | None = True,
-            condition: str | None = None,
+            where_clause: str | None = None,
             if_errors: Literal["ignore","raise"] = "raise",
         ) -> DuckDBIterator:
         from linkmerce.utils.tqdm import _expand_kwargs
@@ -573,11 +563,11 @@ class DuckDBIterator(Task):
             if if_errors == "ignore":
                 from duckdb import BinderException
                 try:
-                    map_partitions[expr] = self.conn.unique(self.table, expr, ascending, condition)
+                    map_partitions[expr] = self.conn.unique(self.table, expr, ascending, where_clause)
                 except BinderException:
                     continue
             else:
-                map_partitions[expr] = self.conn.unique(self.table, expr, ascending, condition)
+                map_partitions[expr] = self.conn.unique(self.table, expr, ascending, where_clause)
         return self.setattr("partitions", _expand_kwargs(**map_partitions))
 
     def __iter__(self) -> DuckDBIterator:
@@ -588,8 +578,8 @@ class DuckDBIterator(Task):
         if self.index >= len(self):
             raise StopIteration
         map_partition = self.partitions[self.index]
-        condition = " AND ".join([f"{expr} = {self.conn.expr_value(value)}" for expr, value in map_partition.items()])
-        query = f"SELECT * FROM {self.table} WHERE {condition or 'TRUE'};"
+        where_clause = " AND ".join([f"{expr} = {self.conn.expr_value(value)}" for expr, value in map_partition.items()])
+        query = concat_sql(f"SELECT * FROM {self.table}", where(where_clause))
         results = self.conn.fetch_all(self.format, query)
         self.index += 1
         return results

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/core/naver/openapi/search/models.sql

@@ -215,10 +215,10 @@ CREATE OR REPLACE TABLE {{ table }} (
     keyword VARCHAR
   , nvMid BIGINT
   , mallPid BIGINT
-  , productType TINYINT -- {0: "가격비교 상품", 1: "가격비교 비매칭 일반상품", 2: "가격비교 매칭 일반상품"}
+  , productType TINYINT -- {0: "가격비교 상품", 1: "가격비교 비매칭 일반상품", 2: "가격비교 매칭 일반상품", 3: "광고상품"}
   , displayRank SMALLINT
   , createdAt TIMESTAMP NOT NULL
-  , PRIMARY KEY (keyword, nvMid)
+  , PRIMARY KEY (keyword, displayRank)
 );
 
 -- ShoppingRank: select_rank
@@ -239,9 +239,12 @@ INSERT INTO {{ table }} {{ values }} ON CONFLICT DO NOTHING;
 CREATE OR REPLACE TABLE {{ table }} (
     nvMid BIGINT PRIMARY KEY
   , mallPid BIGINT
+  , productType TINYINT -- {0: "가격비교 상품", 1: "일반상품", 3: "광고상품"}
   , productName VARCHAR
+  , wholeCategoryName VARCHAR
   , mallName VARCHAR
   , brandName VARCHAR
+  , salesPrice INTEGER
   , updatedAt TIMESTAMP NOT NULL
 );
 
@@ -249,9 +252,12 @@ CREATE OR REPLACE TABLE {{ table }} (
 SELECT
     TRY_CAST(productId AS BIGINT) AS nvMid
   , TRY_CAST(REGEXP_EXTRACT(link, '/products/(\d+)$', 1) AS BIGINT) AS mallPid
+  , IF(link LIKE '%/catalog/%', 0, 1) AS productType
   , REGEXP_REPLACE(title, '<[^>]+>', '', 'g') AS productName
+  , CONCAT_WS('>', category1, category2, category3, category4) AS wholeCategoryName
   , NULLIF(mallName, '네이버') AS mallName
   , NULLIF(brand, '') AS brandName
+  , TRY_CAST(lprice AS INTEGER) AS salesPrice
   , CAST(DATE_TRUNC('second', CURRENT_TIMESTAMP) AS TIMESTAMP) AS updatedAt
 FROM {{ array }}
 WHERE TRY_CAST(productId AS BIGINT) IS NOT NULL;
@@ -261,6 +267,7 @@ INSERT INTO {{ table }} {{ values }}
 ON CONFLICT DO UPDATE SET
     mallPid = COALESCE(excluded.mallPid, mallPid)
   , productName = COALESCE(excluded.productName, productName)
+  , wholeCategoryName = COALESCE(excluded.wholeCategoryName, wholeCategoryName)
   , mallName = COALESCE(excluded.mallName, mallName)
   , brandName = COALESCE(excluded.brandName, brandName)
   , updatedAt = excluded.updatedAt;

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/core/searchad/manage/exposure/models.sql

@@ -29,7 +29,7 @@ SELECT
   , NULLIF(fmpBrand, '') AS mallName
   , NULLIF(fmpMaker, '') AS makerName
   , imageUrl
-  , TRY_CAST(COALESCE(lowPrice, mobileLowPrice, NULL) AS INTEGER) AS salesPrice
+  , TRY_CAST(COALESCE(lowPrice, mobileLowPrice) AS INTEGER) AS salesPrice
 FROM {{ array }}
 WHERE ($is_own IS NULL) OR (isOwn = $is_own);
 
@@ -40,10 +40,10 @@ INSERT INTO {{ table }} {{ values }} ON CONFLICT DO NOTHING;
 -- ExposureRank: create_rank
 CREATE OR REPLACE TABLE {{ table }} (
     keyword VARCHAR
-  , productId BIGINT
+  , nvMid BIGINT
   , displayRank SMALLINT
   , createdAt TIMESTAMP NOT NULL
-  , PRIMARY KEY (keyword, productId)
+  , PRIMARY KEY (keyword, displayRank)
 );
 
 -- ExposureRank: select_rank
@@ -56,24 +56,26 @@ FROM (
           TRY_CAST(REGEXP_EXTRACT(imageUrl, '^https://[^/]+/main_\d+/(\d+)', 1) AS BIGINT)
         WHEN PREFIX(imageUrl, 'https://searchad-') THEN
           TRY_CAST(TRY_CAST(FROM_BASE64(REGEXP_EXTRACT(imageUrl, '^https://[^/]+/[^/]+/([^.]+)', 1)) AS VARCHAR) AS BIGINT)
-        ELSE NULL END) AS productId
+        ELSE NULL END) AS nvMid
     , rank AS displayRank
     , CAST(DATE_TRUNC('second', CURRENT_TIMESTAMP) AS TIMESTAMP) AS createdAt
   FROM {{ array }}
   WHERE ($is_own IS NULL) OR (isOwn = $is_own)
 ) AS exposure
-WHERE exposure.productId IS NOT NULL;
+WHERE exposure.nvMid IS NOT NULL;
 
 -- ExposureRank: insert_rank
 INSERT INTO {{ table }} {{ values }} ON CONFLICT DO NOTHING;
 
 -- ExposureRank: create_product
 CREATE OR REPLACE TABLE {{ table }} (
-    productId BIGINT PRIMARY KEY
-  , isNvMid BOOLEAN
+    nvMid BIGINT PRIMARY KEY
+  , mallPid BIGINT
+  , productType TINYINT -- {0: "가격비교 상품", 1: "일반상품", 3: "광고상품"}
   , productName VARCHAR
   , wholeCategoryName VARCHAR
   , mallName VARCHAR
+  , brandName VARCHAR
   , salesPrice INTEGER
   , updatedAt TIMESTAMP NOT NULL
 );
@@ -87,12 +89,14 @@ FROM (
           TRY_CAST(REGEXP_EXTRACT(imageUrl, '^https://[^/]+/main_\d+/(\d+)', 1) AS BIGINT)
         WHEN PREFIX(imageUrl, 'https://searchad-') THEN
           TRY_CAST(TRY_CAST(FROM_BASE64(REGEXP_EXTRACT(imageUrl, '^https://[^/]+/[^/]+/([^.]+)', 1)) AS VARCHAR) AS BIGINT)
-        ELSE NULL END) AS productId
-    , PREFIX(imageUrl, 'https://shopping-') AS isNvMid
+        ELSE NULL END) AS nvMid
+    , NULL AS mallPid
+    , IF(PREFIX(imageUrl, 'https://shopping-'), 0, 3) AS productType
     , productTitle AS productName
     , categoryNames AS wholeCategoryName
     , NULLIF(fmpBrand, '') AS mallName
-    , TRY_CAST(COALESCE(lowPrice, mobileLowPrice, NULL) AS INTEGER) AS salesPrice
+    , NULL AS brandName
+    , TRY_CAST(COALESCE(lowPrice, mobileLowPrice) AS INTEGER) AS salesPrice
     , CAST(DATE_TRUNC('second', CURRENT_TIMESTAMP) AS TIMESTAMP) AS updatedAt
   FROM {{ array }}
   WHERE ($is_own IS NULL) OR (isOwn = $is_own)

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/core/searchad/manage/exposure/transform.py

@@ -52,8 +52,12 @@ class ExposureRank(ExposureDiagnosis):
             params: dict = dict(),
             **kwargs
         ) -> tuple[DuckDBPyRelation,DuckDBPyRelation]:
+        def reparse_object(obj: list[dict]) -> list[dict]:
+            obj[0] = dict(obj[0], lowPrice=obj[0].get("lowPrice", None), mobileLowPrice=obj[0].get("mobileLowPrice", None))
+            return obj
         def split_params(keyword: str, is_own: bool | None = None, **kwargs) -> tuple[dict,dict]:
             return dict(keyword=keyword, is_own=is_own), dict(is_own=is_own)
+        obj = reparse_object(obj)
         rank_params, product_params = split_params(**params)
         rank = super().insert_into_table(obj, key="insert_rank", table=rank_table, values=":select_rank:", params=rank_params)
         product = super().insert_into_table(obj, key="upsert_product", table=product_table, values=":select_product:", params=product_params)

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/core/smartstore/brand/sales/transform.py

@@ -26,9 +26,9 @@ class _SalesTransformer(DuckDBTransformer):
         ):
         if isinstance(obj, dict):
             if "error" not in obj:
+                sales = obj["data"][f"{self.sales_type}Sales"]
                 params = dict(mall_seq=mall_seq, end_date=end_date)
                 if self.start_date:
-                    sales = obj["data"][f"{self.sales_type}Sales"]
                     params.update(start_date=start_date)
                 return self.insert_into_table(sales, params=params) if sales else None
             else:

linkmerce-0.3.4/src/linkmerce/extensions/bigquery.py

@@ -0,0 +1,468 @@
+from __future__ import annotations
+
+from linkmerce.common.load import Connection, concat_sql, where
+
+from typing import Sequence, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, IO, Literal, Type, TypeVar
+    from types import TracebackType
+    JsonString = TypeVar("JsonString", str)
+    Path = TypeVar("Path", str)
+    TableId = TypeVar("TableId", str)
+
+    from google.cloud.bigquery import Client
+    from google.cloud.bigquery import SchemaField, Table
+    from google.cloud.bigquery.job import LoadJob, LoadJobConfig, QueryJob
+    from google.cloud.bigquery.table import Row, RowIterator
+
+    from linkmerce.common.load import DuckDBConnection
+    DuckDBTable = TypeVar("DuckDBTable", str)
+    BigQueryTable = TypeVar("BigQueryTable", str)
+
+
+DEFAULT_ACCOUNT = "env/service_account.json"
+
+TEMP_TABLE = "temp_table"
+
+
+class ServiceAccount(dict):
+    def __init__(self, info: JsonString | Path | dict[str,str]):
+        super().__init__(self.read_account(info))
+
+    def read_account(self, info: JsonString | Path | dict[str,str]) -> dict:
+        if isinstance(info, dict):
+            return info
+        elif isinstance(info, str):
+            import json
+            if info.startswith('{') and info.endswith('}'):
+                return json.loads(info)
+            else:
+                with open(info, 'r', encoding="utf-8") as file:
+                    return json.loads(file.read())
+        else:
+            raise ValueError("Unrecognized service account.")
+
+
+class PartitionOptions(dict):
+    def __init__(
+            self,
+            by: str | list[str] | None = None,
+            ascending: bool | None = True,
+            where_clause: str | None = None,
+            if_errors: Literal["ignore","raise"] = "raise",
+            **kwargs
+        ):
+        super().__init__(by=by, ascending=ascending, where_clause=where_clause, if_errors=if_errors)
+
+
+###################################################################
+######################### BigQuery Client #########################
+###################################################################
+
+class BigQueryClient(Connection):
+    def __init__(self, account: ServiceAccount):
+        self.set_connection(account)
+
+    @property
+    def conn(self) -> Client:
+        return self.get_connection()
+
+    def get_connection(self) -> Client:
+        return self.__conn
+
+    def set_connection(self, account: ServiceAccount):
+        from google.cloud.bigquery import Client
+        account = account if isinstance(account, ServiceAccount) else ServiceAccount(account)
+        self.__conn = Client.from_service_account_info(account, project=account["project_id"])
+        self.project_id = account["project_id"]
+
+    def close(self):
+        try:
+            self.conn.close()
+        except:
+            pass
+
+    def execute(self, query: str, **kwargs) -> QueryJob:
+        return self.conn.query(query, **kwargs)
+
+    def execute_job(self, query: str, **kwargs) -> RowIterator:
+        return self.conn.query_and_wait(query, **kwargs)
+
+    def __enter__(self) -> BigQueryClient:
+        return self
+
+    def __exit__(self, type: Type[BaseException], value: BaseException, traceback: TracebackType):
+        self.close()
+
+    ############################## Fetch ##############################
+
+    def fetch_all(self, format: Literal["csv","json"], query: str) -> list[dict]:
+        try:
+            return getattr(self, f"fetch_all_to_{format}")(query)
+        except AttributeError:
+            raise ValueError("Invalid value for data format. Supported formats are: csv, json.")
+
+    def fetch_all_to_csv(self, query: str) -> list[tuple]:
+        def row_keys(row: Row) -> tuple:
+            return tuple(row.values())
+        def row_values(row: Row) -> tuple:
+            return tuple(row.values())
+        rows = self.execute_job(query)
+        return ([row_keys(rows[0])] if rows else list()) + [row_values(row) for row in rows]
+
+    def fetch_all_to_json(self, query: str) -> list[dict]:
+        def row_to_dict(row: Row) -> dict[str,Any]:
+            return dict(row.items())
+        return [row_to_dict(row) for row in self.execute_job(query)]
+
+    ############################ CRUD Table ###########################
+
+    def create_table(
+            self,
+            table: str,
+            schema: TableId | Sequence[dict | SchemaField],
+            exists_ok: bool = True,
+            **kwargs
+        ) -> Table:
+        table_ref = self.ref_table(table, schema)
+        return self.conn.create_table(table_ref, exists_ok=exists_ok, **kwargs)
+
+    def copy_table(
+            self,
+            source_table: str,
+            target_table: str,
+            where_clause: str | None = None,
+            limit: int | None = 0,
+            option: Literal["replace","ignore"] | None = None,
+        ) -> RowIterator:
+        select = f"SELECT * FROM `{self.project_id}.{source_table}`"
+        limit_ = f"LIMIT {limit}" if isinstance(limit, int) else None
+        query = concat_sql(f"{self.expr_create(option)} `{self.project_id}.{target_table}` AS", select, where(where_clause), limit_)
+        return self.execute_job(query)
+
+    def delete_table(self, table: str, where: str = "TRUE") -> RowIterator:
+        return self.execute_job(f"DELETE FROM `{self.project_id}.{table}` WHERE {where};")
+
+    def select_table_to_json(self, table: str) -> list[dict]:
+        return self.fetch_all_to_json(f"SELECT * FROM `{self.project_id}.{table}`;")
+
+    def table_exists(self, table: str) -> bool:
+        from google.api_core.exceptions import NotFound
+        try:
+            self.conn.get_table(f"{self.project_id}.{table}")
+            return True
+        except NotFound:
+            return False
+
+    def table_has_rows(self, table: str, where_clause: str | None = None) -> bool:
+        if self.table_exists(table):
+            query = concat_sql(f"SELECT COUNT(*) FROM `{self.project_id}.{table}`", where(where_clause))
+            rows = list(self.execute_job(query))
+            return bool(list(rows[0].values())[0])
+        return False
+
+    def get_table(self, table: str) -> Table:
+        return self.conn.get_table(f"{self.project_id}.{table}")
+
+    def ref_table(self, table: str, schema: TableId | Sequence[dict | SchemaField]) -> Table:
+        from google.cloud.bigquery import Table
+        if isinstance(schema, str):
+            schema = self.get_schema(schema)
+        if not isinstance(schema, Sequence):
+            raise ValueError("Invalid schema: expected sequence of schema fields.")
+        return Table(table, schema)
+
+    def get_schema(self, table: str) -> list[SchemaField]:
+        return self.conn.get_table(f"{self.project_id}.{table}").schema
+
+    ############################# Load Job ############################
+
+    def load_table_from_json(
+            self,
+            table: str,
+            values: list[dict] | str | Path,
+            serialize: bool = True,
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+            write: Literal["append","empty","truncate","truncate_data"] = "append",
+            if_not_found: Literal["create","errors","ignore"] = "errors",
+        ) -> LoadJob:
+        schema = self._auto_detect_schema(table, schema)
+        if self._find_table(table, schema, if_not_found):
+            if serialize:
+                import json
+                values = json.loads(json.dumps(values, ensure_ascii=False, default=str))
+            job_config = self.build_load_job_config(schema, write_disposition=write.upper())
+            return self.conn.load_table_from_json(values, f"{self.project_id}.{table}", job_config=job_config).result()
+
+    def load_table_from_file(
+            self,
+            table: str,
+            file_obj: IO[bytes],
+            format: Literal["avgo","csv","json","orc","parquet"],
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+            write: Literal["append","empty","truncate","truncate_data"] = "append",
+            if_not_found: Literal["create","errors","ignore"] = "errors",
+        ) -> LoadJob:
+        schema = self._auto_detect_schema(table, schema)
+        if self._find_table(table, schema, if_not_found):
+            job_config = self.build_load_job_config(schema, source_format=format.upper(), write_disposition=write.upper())
+            return self.conn.load_table_from_file(file_obj, f"{self.project_id}.{table}", job_config=job_config).result()
+
+    def build_load_job_config(
+            self,
+            schema: Sequence[dict | SchemaField] | None = None,
+            source_format: Literal["AVRO", "CSV", "JSON", "ORC", "PARQUET"] | None = None,
+            write_disposition: Literal["APPEND", "EMPTY", "TRUNCATE", "TRUNCATE_DATA"] | None = None,
+            **kwargs
+        ) -> LoadJobConfig:
+        from google.cloud.bigquery.job import LoadJobConfig
+        if schema is not None:
+            kwargs["schema"] = self.build_schema(schema)
+        if source_format is not None:
+            kwargs["source_format"] = "NEWLINE_DELIMITED_JSON" if source_format == "JSON" else source_format
+        if write_disposition is not None:
+            kwargs["write_disposition"] = f"WRITE_{write_disposition}"
+        return LoadJobConfig(**kwargs)
+
+    def build_schema(self, fields: Sequence[dict | SchemaField]) -> list[SchemaField]:
+        from google.cloud.bigquery import SchemaField
+        def build(type: str | None = None, **kwargs) -> SchemaField:
+            if type is not None:
+                kwargs["field_type"] = type
+            return SchemaField(**kwargs)
+        return [field if isinstance(field, SchemaField) else build(**field) for field in fields]
+
+    def _auto_detect_schema(
+            self,
+            table: str,
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+        ) -> list[dict | SchemaField]:
+        if isinstance(schema, str):
+            return self.get_schema(table if schema == "auto" else schema)
+        else:
+            return schema
+
+    def _find_table(
+            self,
+            table: str,
+            schema: list[dict | SchemaField],
+            if_not_found: Literal["create","errors","ignore"] = "errors",
+        ) -> bool:
+        if if_not_found == "create":
+            if not self.table_exists(table):
+                self.create_table(table, schema)
+        elif if_not_found == "ignore":
+            return self.table_exists(table)
+        return True
+
+    ############################## Merge ##############################
+
+    def merge_into_table(
+            self,
+            source_table: str,
+            target_table: str,
+            on: str | Sequence[str],
+            matched: str | dict[str,Literal["source_first","target_first","greatest","least","replace","ignore"]],
+            not_matched: str | Sequence[str],
+            where_clause: str | None = None,
+        ) -> LoadJob:
+        where = [f"T.{where_clause}"] if where_clause else list()
+        on = " AND ".join([f"T.{col} = S.{col}" for col in ([on] if isinstance(on, str) else on)]+where)
+        query = f"MERGE INTO `{self.project_id}.{target_table}` AS T USING `{self.project_id}.{source_table}` AS S ON {on}"
+        query = concat_sql(query, self._merge_update(matched), self._merge_insert(not_matched))
+        self.execute_job(query)
+
+    def _merge_update(self, matched: str | dict[str,Literal["count","sum","avg","min","max","first","last","list"]]) -> str:
+        prefix = "WHEN MATCHED THEN UPDATE SET "
+        if isinstance(matched, dict):
+            def render(col: str, agg: str) -> str:
+                if agg in {"source_first","target_first"}:
+                    kwargs = dict(zip(["left","right"], ('S','T') if agg == "source_first" else ('T','S')))
+                    return "COALESCE({left}.{col}, {right}.{col})".format(col=col, **kwargs)
+                elif agg in {"greatest","least"}:
+                    return f"{agg.upper()}(S.{col}, T.{col})"
+                elif agg in {"replace","ignore"}:
+                    return f"S.{agg}" if agg == "replace" else f"T.{agg}"
+                else:
+                    return f"{agg}({col})"
+            return prefix + ", ".join([f"T.{col} = {render(col, agg)}" for col, agg in matched.items()])
+        else:
+            return prefix + str(matched)
+
+    def _merge_insert(self, not_matched: Sequence[str]) -> str:
+        prefix = "WHEN NOT MATCHED THEN "
+        if (not isinstance(not_matched, str)) and isinstance(not_matched, Sequence):
+            return prefix + "INSERT ({}) VALUES ({})".format(", ".join(not_matched), "S."+", S.".join(not_matched))
+        else:
+            return prefix + not_matched
+
+    ########################## Load and Merge #########################
+
+    def merge_into_table_from_file(
+            self,
+            stage_table: str,
+            target_table: str,
+            source_file: IO[bytes],
+            source_format: Literal["avgo","csv","json","orc","parquet"],
+            on: str | Sequence[str],
+            matched: str | dict[str,Literal["source_first","target_first","greatest","least","replace","ignore"]],
+            not_matched: str | Sequence[str],
+            where_clause: str | None = None,
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+            write: Literal["append","empty","truncate","truncate_data"] = "truncate",
+            if_not_found: Literal["create","errors","ignore"] = "errors",
+            table_lock_wait_interval: float | int | None = 1.,
+            table_lock_wait_timeout: float | int | None = 60.,
+            drop_stage_after_merge: bool = True,
+        ) -> LoadJob:
+        try:
+            self._wait_until_table_not_found(stage_table, table_lock_wait_interval, table_lock_wait_timeout)
+            self.load_table_from_file(source_file, stage_table, source_format, schema, write, if_not_found)
+            return self.merge_into_table(stage_table, target_table, on, matched, not_matched, where_clause)
+        finally:
+            if drop_stage_after_merge and self.table_exists(stage_table):
+                self.execute_job(f"DROP TABLE `{self.project_id}.{stage_table}`")
+
+    def _wait_until_table_not_found(self, table: str, interval: float | int | None = 1., timeout: float | int | None = 60.):
+        has_interval, has_timeout = isinstance(interval, (float,int)), isinstance(timeout, (float,int))
+        if has_interval:
+            import time
+            total = 0
+            while self.table_exists(table):
+                if has_timeout and (total > timeout):
+                    raise TimeoutError("Timed out waiting until the table does not exist.")
+                total += interval
+                time.sleep(interval)
+
+    ############################ Expression ###########################
+
+    def expr_cast(self, value: Any | None, type: str, alias: str = str(), safe: bool = False) -> str:
+        cast = "SAFE_CAST" if safe else "CAST"
+        alias = f" AS {alias}" if alias else str()
+        return f"{cast}({self.expr_value(value)} AS {type.upper()})" + alias
+
+    def expr_interval(expr: str, days: int | None = None, time: bool = True) -> str:
+        if isinstance(days, int):
+            func = ("DATE{}_SUB" if days < 0 else "DATE{}_ADD").format("TIME" if time else str())
+            return f"{func}({expr}, INTERVAL {abs(days)} DAY)"
+        else:
+            return expr
+
+    def expr_now(
+            self,
+            type: Literal["DATETIME","STRING"] = "DATETIME",
+            format: str | None = "%Y-%m-%d %H:%M:%S",
+            interval: str | int | None = None,
+            tzinfo: str | None = None,
+        ) -> str:
+        expr = "CURRENT_DATETIME({})".format(f"'{tzinfo}'" if tzinfo else str())
+        expr = self.expr_interval(expr, interval, time=True)
+        if format:
+            expr = f"FORMAT_DATE('{format}', {expr})"
+            if type.upper() == "DATETIME":
+                return f"CAST({expr} AS DATETIME)"
+        return expr if type.upper() == "DATETIME" else "NULL"
+
+    def expr_today(
+            self,
+            type: Literal["DATE","STRING"] = "DATE",
+            format: str | None = "%Y-%m-%d",
+            interval: str | int | None = None,
+            tzinfo: str | None = None,
+        ) -> str:
+        expr = "CURRENT_DATE({})".format(f"'{tzinfo}'" if tzinfo else str())
+        expr = self.expr_interval(expr, interval, time=False)
+        if (type.upper() == "STRING") and format:
+            return f"FORMAT_DATE('{format}', {expr})"
+        return expr if type.upper() == "DATE" else "NULL"
+
+    ############################## DuckDB #############################
+
+    def load_table_from_duckdb(
+            self,
+            connection: DuckDBConnection,
+            source_table: DuckDBTable,
+            target_table: BigQueryTable,
+            partition_by: PartitionOptions = dict(),
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+            write: Literal["append","empty","truncate","truncate_data"] = "append",
+            if_not_found: Literal["create","errors","ignore"] = "errors",
+            progress: bool = True,
+        ) -> bool:
+        from linkmerce.common.load import DuckDBIterator
+        from linkmerce.utils.tqdm import import_tqdm
+        from io import BytesIO
+
+        if not connection.table_exists(source_table):
+            return True
+        schema = self._auto_detect_schema(target_table, schema)
+
+        iterator = DuckDBIterator(connection, format="parquet").from_table(source_table)
+        if partition_by:
+            iterator = iterator.partition_by(**partition_by)
+
+        tqdm = import_tqdm()
+        for bytes_ in tqdm(iterator, desc=f"Uploading data to '{target_table}'", disable=(not progress)):
+            self.load_table_from_file(target_table, BytesIO(bytes_), "parquet", schema, write, if_not_found)
+        return True
+
+    def overwrite_table_from_duckdb(
+            self,
+            connection: DuckDBConnection,
+            source_table: DuckDBTable,
+            target_table: BigQueryTable,
+            where_clause: str = "TRUE",
+            partition_by: PartitionOptions = dict(),
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+            if_not_found: Literal["create","errors","ignore"] = "errors",
+            progress: bool = True,
+            backup_table: DuckDBTable | None = TEMP_TABLE,
+        ) -> bool:
+        if not connection.table_exists(source_table):
+            return True
+        elif not self.table_has_rows(target_table, where_clause):
+            return self.load_table_from_duckdb(connection, source_table, target_table, partition_by, schema, "append", if_not_found, progress)
+
+        success = False
+        from_clause = concat_sql(f"FROM `{self.project_id}.{target_table}`", where(where_clause))
+        existing_values = self.fetch_all_to_json(f"SELECT * {from_clause};")
+        self.conn.query(f"DELETE {from_clause};")
+
+        try:
+            success = self.load_table_from_duckdb(connection, source_table, target_table, partition_by, schema, "append", if_not_found, progress)
+            return success
+        finally:
+            if (not success) and (backup_table is not None) and existing_values:
+                connection.copy_table(source_table, backup_table, option="replace", temp=True)
+                connection.insert_into_table_from_json(backup_table, existing_values)
+
+    def merge_into_table_from_duckdb(
+            self,
+            connection: DuckDBConnection,
+            source_table: DuckDBTable,
+            staging_table: BigQueryTable,
+            target_table: BigQueryTable,
+            on: str | Sequence[str],
+            matched: str | dict[str,Literal["source_first","target_first","greatest","least","replace","ignore"]],
+            not_matched: str | Sequence[str],
+            schema: Literal["auto"] | TableId | Sequence[dict | SchemaField] = "auto",
+            where_clause: str | None = None,
+            table_lock_wait_interval: float | int | None = 1.,
+            table_lock_wait_timeout: float | int | None = 60.,
+            drop_stage_after_merge: bool = True,
+        ) -> bool:
+        if not connection.table_exists(source_table):
+            return True
+        elif not self.table_has_rows(target_table, where_clause):
+            return self.load_table_from_duckdb(connection, source_table, target_table, schema=schema, write="append")
+
+        try:
+            self._wait_until_table_not_found(staging_table, table_lock_wait_interval, table_lock_wait_timeout)
+            self.copy_table(target_table, staging_table, option="replace")
+            self.load_table_from_duckdb(connection, source_table, staging_table, schema=schema, write="append")
+            self.merge_into_table(staging_table, target_table, on, matched, not_matched, where_clause)
+            return True
+        finally:
+            if drop_stage_after_merge and self.table_exists(staging_table):
+                self.execute_job(f"DROP TABLE `{self.project_id}.{staging_table}`")

{linkmerce-0.3.2 → linkmerce-0.3.4}/src/linkmerce/utils/tqdm.py

@@ -4,13 +4,31 @@ from typing import Sequence, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Coroutine, Hashable, Iterable, TypeVar
+    from types import ModuleType
     from numbers import Real
     _KT = TypeVar("_KT", Hashable)
     _VT = TypeVar("_VT", Any)
 
 
+def import_tqdm() -> ModuleType:
+    try:
+        from tqdm import tqdm
+        return tqdm
+    except:
+        return lambda x, **kwargs: x
+
+
+def import_tqdm_asyncio() -> ModuleType:
+    try:
+        from tqdm.asyncio import tqdm_asyncio
+        return tqdm_asyncio
+    except:
+        import asyncio
+        return asyncio
+
+
 ###################################################################
-############################## Select #############################
+############################## Gather #############################
 ###################################################################
 
 def gather(
@@ -20,12 +38,13 @@ def gather(
         delay: Real | Sequence[Real,Real] = 0.,
         tqdm_options: dict = dict(),
     ) -> list:
+    import time
     try:
         from tqdm import tqdm
     except:
         tqdm = lambda x: x
         tqdm_options = dict()
-    import time
+
     def run_with_delay(args: tuple[_VT,...] | dict[_KT,_VT]) -> Any:
         try:
             if isinstance(args, dict):
@@ -51,6 +70,7 @@ async def gather_async(
     except:
         tqdm_asyncio = asyncio
         tqdm_options = dict()
+
     async def run_with_delay(args: tuple[_VT,...] | dict[_KT,_VT]) -> Any:
         try:
             if isinstance(args, dict):
@@ -81,7 +101,7 @@ def _get_seconds(value: Real | Sequence[Real,Real]) -> Real:
 
 
 ###################################################################
-############################## Select #############################
+############################## Expand #############################
 ###################################################################
 
 def expand(