feat: catalogs for DuckDB (#28751)

Co-authored-by: Maxime Beauchemin <maximebeauchemin@gmail.com>
This commit is contained in:
Beto Dealmeida
2025-04-11 15:58:59 -04:00
committed by GitHub
parent 7ab8534ef6
commit bcd136cee1
2 changed files with 36 additions and 1 deletions

View File

@@ -125,7 +125,7 @@ denodo = ["denodo-sqlalchemy~=1.0.6"]
dremio = ["sqlalchemy-dremio>=1.2.1, <4"]
drill = ["sqlalchemy-drill>=1.1.4, <2"]
druid = ["pydruid>=0.6.5,<0.7"]
duckdb = ["duckdb-engine>=0.10", "duckdb>=1.1.0"]
duckdb = ["duckdb-engine>=0.12.1, <0.13"]
dynamodb = ["pydynamodb>=0.4.2"]
solr = ["sqlalchemy-solr >= 0.2.0"]
elasticsearch = ["elasticsearch-dbapi>=0.2.9, <0.3.0"]
@@ -146,6 +146,7 @@ hive = [
impala = ["impyla>0.16.2, <0.17"]
kusto = ["sqlalchemy-kusto>=3.0.0, <4"]
kylin = ["kylinpy>=2.8.1, <2.9"]
motherduck = ["duckdb==0.10.2", "duckdb-engine>=0.12.1, <0.13"]
mssql = ["pymssql>=2.2.8, <3"]
mysql = ["mysqlclient>=2.1.0, <3"]
ocient = [

View File

@@ -14,6 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import re
@@ -262,6 +263,9 @@ class MotherDuckEngineSpec(DuckDBEngineSpec):
engine_name = "MotherDuck"
engine_aliases: set[str] = {"duckdb"}
supports_catalog = True
supports_dynamic_catalog = True
sqlalchemy_uri_placeholder = (
"duckdb:///md:{database_name}?motherduck_token={SERVICE_TOKEN}"
)
@@ -296,3 +300,33 @@ class MotherDuckEngineSpec(DuckDBEngineSpec):
return str(
URL(drivername=DuckDBEngineSpec.engine, database=database, query=query)
)
@classmethod
def adjust_engine_params(
cls,
uri: URL,
connect_args: dict[str, Any],
catalog: str | None = None,
schema: str | None = None,
) -> tuple[URL, dict[str, Any]]:
if catalog:
uri = uri.set(database=f"md:{catalog}")
return uri, connect_args
@classmethod
def get_default_catalog(cls, database: Database) -> str | None:
return database.url_object.database.split(":", 1)[1]
@classmethod
def get_catalog_names(
cls,
database: Database,
inspector: Inspector,
) -> set[str]:
return {
catalog
for (catalog,) in inspector.bind.execute(
"SELECT alias FROM MD_ALL_DATABASES() WHERE is_attached;"
)
}