Improve get_objects API

Author: kou

lib/activerecord_adbc_adapter/adapter.rb

@@ -46,12 +46,29 @@ def open_statement(&block)
         @connection.open_statement(&block)
       end
 
-      def get_objects(*args)
-        reader = @connection.get_objects(*args)
-        begin
-          yield(reader.read_all)
-        ensure
-          reader.unref
+      if ADBC::Connection.method_defined?(:get_objects_raw)
+        # red-adbc provides convenient wrapper
+        def get_objects(...)
+          @connection.get_objects(...)
+        end
+      else
+        def get_objects(depth: :all,
+                        catalog: nil,
+                        db_schema: nil,
+                        table_name: nil,
+                        table_types: nil,
+                        column_name: nil)
+          reader = @connection.get_objects(depth,
+                                           catalog,
+                                           db_schema,
+                                           table_name,
+                                           table_types,
+                                           column_name)
+          begin
+            reader.read_all
+          ensure
+            reader.unref
+          end
         end
       end
     end

lib/activerecord_adbc_adapter/schema_statements.rb

@@ -63,94 +63,86 @@ def adbc_view_type
     def tables
       type = adbc_table_type
       with_raw_connection do |conn|
-        conn.get_objects(:tables,
-                         adbc_catalog,
-                         adbc_db_schema,
-                         nil,
-                         [type]) do |table|
-          tables = []
-          table.raw_records.each do |_catalog_name, db_schemas|
-            db_schemas.each do |db_schema|
-              db_schema_tables = db_schema["db_schema_tables"]
-              next if db_schema_tables.nil?
-              db_schema_tables.each do |t|
-                # Some drivers may ignore table_types
-                next unless t["table_type"] == type
-                tables << t["table_name"]
-              end
+        objects = conn.get_objects(depth: :tables,
+                                   catalog: adbc_catalog,
+                                   db_schema: adbc_db_schema,
+                                   table_types: [type])
+        tables = []
+        objects.raw_records.each do |_catalog_name, db_schemas|
+          db_schemas.each do |db_schema|
+            db_schema_tables = db_schema["db_schema_tables"]
+            next if db_schema_tables.nil?
+            db_schema_tables.each do |t|
+              # Some drivers may ignore table_types
+              next unless t["table_type"] == type
+              tables << t["table_name"]
             end
           end
-          tables
         end
+        tables
       end
     end
 
     def views
       type = adbc_view_type
       with_raw_connection do |conn|
-        conn.get_objects(:tables,
-                         adbc_catalog,
-                         adbc_db_schema,
-                         nil,
-                         [type]) do |table|
-          views = []
-          table.raw_records.each do |_catalog_name, db_schemas|
-            db_schemas.each do |db_schema|
-              db_schema_tables = db_schema["db_schema_tables"]
-              next if db_schema_tables.nil?
-              db_schema_tables.each do |t|
-                # Some drivers may ignore table_types
-                next unless t["table_type"] == type
-                views << t["table_name"]
-              end
+        objects = conn.get_objects(depth: :tables,
+                                   catalog: adbc_catalog,
+                                   db_schema: adbc_db_schema,
+                                   table_types: [type])
+        views = []
+        objects.raw_records.each do |_catalog_name, db_schemas|
+          db_schemas.each do |db_schema|
+            db_schema_tables = db_schema["db_schema_tables"]
+            next if db_schema_tables.nil?
+            db_schema_tables.each do |t|
+              # Some drivers may ignore table_types
+              next unless t["table_type"] == type
+              views << t["table_name"]
             end
           end
-          views
         end
+        views
       end
     end
 
     def column_definitions(table_name)
       with_raw_connection do |conn|
-        conn.get_objects(:all,
-                         adbc_catalog,
-                         adbc_db_schema,
-                         table_name) do |table|
-          table.raw_records.each do |_catalog_name, db_schemas|
-            db_schemas.each do |db_schema|
-              db_schema_tables = db_schema["db_schema_tables"]
-              next if db_schema_tables.nil?
-              db_schema_tables.each do |t|
-                return t["table_columns"]
-              end
+        objects = conn.get_objects(catalog: adbc_catalog,
+                                   db_schema: adbc_db_schema,
+                                   table_name: table_name)
+        objects.raw_records.each do |_catalog_name, db_schemas|
+          db_schemas.each do |db_schema|
+            db_schema_tables = db_schema["db_schema_tables"]
+            next if db_schema_tables.nil?
+            db_schema_tables.each do |table|
+              return table["table_columns"]
             end
           end
-          [] # raise?
         end
+        [] # raise?
       end
     end
 
     def primary_keys(table_name)
       with_raw_connection do |conn|
-        conn.get_objects(:all,
-                         adbc_catalog,
-                         adbc_db_schema,
-                         table_name) do |table|
-          table.raw_records.each do |_catalog_name, db_schemas|
-            db_schemas.each do |db_schema|
-              db_schema_tables = db_schema["db_schema_tables"]
-              next if db_schema_tables.nil?
-              db_schema_tables.each do |t|
-                constraint = t["table_constraints"].find do |constraint|
-                  constraint["constraint_type"] == "PRIMARY KEY"
-                end
-                return [] if constraint.nil?
-                return constraint["constraint_column_names"] || []
+        objects = conn.get_objects(catalog: adbc_catalog,
+                                   db_schema: adbc_db_schema,
+                                   table_name: table_name)
+        objects.raw_records.each do |_catalog_name, db_schemas|
+          db_schemas.each do |db_schema|
+            db_schema_tables = db_schema["db_schema_tables"]
+            next if db_schema_tables.nil?
+            db_schema_tables.each do |table|
+              constraint = table["table_constraints"].find do |constraint|
+                constraint["constraint_type"] == "PRIMARY KEY"
               end
+              return [] if constraint.nil?
+              return constraint["constraint_column_names"] || []
             end
           end
-          []
         end
+        []
       end
     end