diff --git a/csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs b/csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs index 3203502e57..e5b2b3044f 100644 --- a/csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs +++ b/csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs @@ -405,27 +405,37 @@ protected virtual async Task GetCrossReferenceAsForeignTableAsync(C return await GetQueryResult(resp.DirectResults, cancellationToken); } + /// + /// Gets the cross reference (foreign key) information for the specified tables. + /// Note: Unlike other metadata queries, this method does not escape underscores in names + /// since the backend treats these as exact match queries rather than pattern matches. + /// protected virtual async Task GetCrossReferenceAsync(CancellationToken cancellationToken = default) { TGetCrossReferenceResp resp = await Connection.GetCrossReferenceAsync( - EscapeUnderscoreInName(CatalogName), - EscapeUnderscoreInName(SchemaName), - EscapeUnderscoreInName(TableName), - EscapeUnderscoreInName(ForeignCatalogName), - EscapeUnderscoreInName(ForeignSchemaName), - EscapeUnderscoreInName(ForeignTableName), + CatalogName, + SchemaName, + TableName, + ForeignCatalogName, + ForeignSchemaName, + ForeignTableName, cancellationToken); OperationHandle = resp.OperationHandle; return await GetQueryResult(resp.DirectResults, cancellationToken); } + /// + /// Gets the primary key information for the specified table. + /// Note: Unlike other metadata queries, this method does not escape underscores in names + /// since the backend treats these as exact match queries rather than pattern matches. + /// protected virtual async Task GetPrimaryKeysAsync(CancellationToken cancellationToken = default) { TGetPrimaryKeysResp resp = await Connection.GetPrimaryKeysAsync( - EscapeUnderscoreInName(CatalogName), - EscapeUnderscoreInName(SchemaName), - EscapeUnderscoreInName(TableName), + CatalogName, + SchemaName, + TableName, cancellationToken); OperationHandle = resp.OperationHandle; diff --git a/csharp/src/Drivers/Databricks/DatabricksStatement.cs b/csharp/src/Drivers/Databricks/DatabricksStatement.cs index 1dc0f7a501..d832a7edc9 100644 --- a/csharp/src/Drivers/Databricks/DatabricksStatement.cs +++ b/csharp/src/Drivers/Databricks/DatabricksStatement.cs @@ -421,10 +421,17 @@ internal bool ShouldReturnEmptyPkFkResult() if (!enablePKFK) return true; - // Handle special catalog cases - if (string.IsNullOrEmpty(CatalogName) || + var catalogInvalid = string.IsNullOrEmpty(CatalogName) || string.Equals(CatalogName, "SPARK", StringComparison.OrdinalIgnoreCase) || - string.Equals(CatalogName, "hive_metastore", StringComparison.OrdinalIgnoreCase)) + string.Equals(CatalogName, "hive_metastore", StringComparison.OrdinalIgnoreCase); + + var foreignCatalogInvalid = string.IsNullOrEmpty(ForeignCatalogName) || + string.Equals(ForeignCatalogName, "SPARK", StringComparison.OrdinalIgnoreCase) || + string.Equals(ForeignCatalogName, "hive_metastore", StringComparison.OrdinalIgnoreCase); + + // Handle special catalog cases + // Only when both catalog and foreignCatalog is Invalid, we return empty results + if (catalogInvalid && foreignCatalogInvalid) { return true; } @@ -473,6 +480,7 @@ protected override async Task GetCrossReferenceAsync(CancellationTo return await base.GetCrossReferenceAsync(cancellationToken); } + protected override async Task GetCrossReferenceAsForeignTableAsync(CancellationToken cancellationToken = default) { if (ShouldReturnEmptyPkFkResult()) diff --git a/csharp/test/Apache.Arrow.Adbc.Tests/TestBase.cs b/csharp/test/Apache.Arrow.Adbc.Tests/TestBase.cs index 919a4cc286..6e56f0f50e 100644 --- a/csharp/test/Apache.Arrow.Adbc.Tests/TestBase.cs +++ b/csharp/test/Apache.Arrow.Adbc.Tests/TestBase.cs @@ -658,7 +658,7 @@ private static string GetNameWithoutFirstChatacter(string name) return name.Substring(1); } - protected void CreateNewTableName(out string tableName, out string fullTableName) + protected virtual void CreateNewTableName(out string tableName, out string fullTableName) { string catalogName = TestConfiguration.Metadata.Catalog; string schemaName = TestConfiguration.Metadata.Schema; diff --git a/csharp/test/Drivers/Apache/Common/StatementTests.cs b/csharp/test/Drivers/Apache/Common/StatementTests.cs index 2234d80802..6438fa5737 100644 --- a/csharp/test/Drivers/Apache/Common/StatementTests.cs +++ b/csharp/test/Drivers/Apache/Common/StatementTests.cs @@ -497,7 +497,8 @@ private static async Task ValidateGetCrossReference(string? catalogName, string? Assert.Equal(StringType.Default, queryResult.Stream.Schema.FieldsList[5].DataType); // FK_SCHEMA_NAME Assert.Equal(StringType.Default, queryResult.Stream.Schema.FieldsList[6].DataType); // FK_TABLE_NAME Assert.Equal(StringType.Default, queryResult.Stream.Schema.FieldsList[7].DataType); // FK_COLUMN_NAME - Assert.Equal(Int32Type.Default, queryResult.Stream.Schema.FieldsList[8].DataType); // FK_INDEX + // Databricks return Int16(SmallInt) + Assert.True(queryResult.Stream.Schema.FieldsList[8].DataType is Int32Type or Int16Type, "FK_INDEX should be either Int32 or Int16"); // FK_INDEX Assert.Equal(expectedBatchLength, batch.Length); actualBatchLength += batch.Length; for (int i = 0; i < batch.Length; i++) diff --git a/csharp/test/Drivers/Databricks/StatementTests.cs b/csharp/test/Drivers/Databricks/StatementTests.cs index d6ac50e1d5..8d8aa19370 100644 --- a/csharp/test/Drivers/Databricks/StatementTests.cs +++ b/csharp/test/Drivers/Databricks/StatementTests.cs @@ -100,6 +100,15 @@ public LongRunningStatementTimeoutTestData() } } + protected override void CreateNewTableName(out string tableName, out string fullTableName) + { + string catalogName = TestConfiguration.Metadata.Catalog; + string schemaName = TestConfiguration.Metadata.Schema; + tableName = Guid.NewGuid().ToString("N") + "`!@#$%^&*()_+-="; + string catalogFormatted = string.IsNullOrEmpty(catalogName) ? string.Empty : DelimitIdentifier(catalogName) + "."; + fullTableName = $"{catalogFormatted}{DelimitIdentifier(schemaName)}.{DelimitIdentifier(tableName)}"; + } + [SkippableFact] public async Task CanGetPrimaryKeysDatabricks() { @@ -109,6 +118,8 @@ public async Task CanGetPrimaryKeysDatabricks() [SkippableFact] public async Task CanGetCrossReferenceFromParentTableDatabricks() { + // TODO: Get cross reference from Parent is not currently supported in Databricks + Skip.If(true, "GetCrossReference is not supported in Databricks"); await base.CanGetCrossReferenceFromParentTable(TestConfiguration.Metadata.Catalog, TestConfiguration.Metadata.Schema); } @@ -435,6 +446,17 @@ protected override void PrepareCreateTableWithPrimaryKeys(out string sqlUpdate, primaryKeys = ["index", "name"]; } + + protected override void PrepareCreateTableWithForeignKeys(string fullTableNameParent, out string sqlUpdate, out string tableNameChild, out string fullTableNameChild, out IReadOnlyList foreignKeys) + { + CreateNewTableName(out tableNameChild, out fullTableNameChild); + sqlUpdate = $"CREATE TABLE IF NOT EXISTS {fullTableNameChild} \n" + + " (INDEX INT, USERINDEX INT, USERNAME STRING, ADDRESS STRING, \n" + + " PRIMARY KEY (INDEX), \n" + + $" FOREIGN KEY (USERINDEX, USERNAME) REFERENCES {fullTableNameParent} (INDEX, NAME))"; + foreignKeys = ["userindex", "username"]; + } + // NOTE: this is a thirty minute test. As of writing, databricks commands have 20 minutes of idle time (and checked every 5 mintues) [SkippableTheory] [InlineData(false, "CloudFetch disabled")] @@ -833,5 +855,6 @@ public async Task OlderDBRVersion_ShouldSetSchemaViaUseStatement() Assert.True(rowCount > 0, "Should have results even without catalog specified"); Assert.True(foundSchemas.Count == 1, "Should have exactly one schema"); } + } }